module Language.Fixpoint.Utils.Trie
  ( -- * Datatype
    Trie (..)
  , Branch (..)

    -- * Constructors
  , empty
  , insert
  , fromList

    -- * Visitors
  , fold
  , foldM
  )
  where

import qualified Data.List as L
import Language.Fixpoint.Types.PrettyPrint
import qualified Language.Fixpoint.Misc as Misc

type Key  = Int
type Path = [Key]

newtype Trie a
  = Node [Branch a]
  deriving (Trie a -> Trie a -> Bool
forall a. Eq a => Trie a -> Trie a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Trie a -> Trie a -> Bool
$c/= :: forall a. Eq a => Trie a -> Trie a -> Bool
== :: Trie a -> Trie a -> Bool
$c== :: forall a. Eq a => Trie a -> Trie a -> Bool
Eq, Int -> Trie a -> ShowS
forall a. Show a => Int -> Trie a -> ShowS
forall a. Show a => [Trie a] -> ShowS
forall a. Show a => Trie a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Trie a] -> ShowS
$cshowList :: forall a. Show a => [Trie a] -> ShowS
show :: Trie a -> String
$cshow :: forall a. Show a => Trie a -> String
showsPrec :: Int -> Trie a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Trie a -> ShowS
Show)

data Branch a
  = Bind !Key !(Trie a)
  | Val a
  deriving (Branch a -> Branch a -> Bool
forall a. Eq a => Branch a -> Branch a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Branch a -> Branch a -> Bool
$c/= :: forall a. Eq a => Branch a -> Branch a -> Bool
== :: Branch a -> Branch a -> Bool
$c== :: forall a. Eq a => Branch a -> Branch a -> Bool
Eq, Int -> Branch a -> ShowS
forall a. Show a => Int -> Branch a -> ShowS
forall a. Show a => [Branch a] -> ShowS
forall a. Show a => Branch a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Branch a] -> ShowS
$cshowList :: forall a. Show a => [Branch a] -> ShowS
show :: Branch a -> String
$cshow :: forall a. Show a => Branch a -> String
showsPrec :: Int -> Branch a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Branch a -> ShowS
Show)

-------------------------------------------------------------------------------
empty :: Trie a
-------------------------------------------------------------------------------
empty :: forall a. Trie a
empty = forall a. [Branch a] -> Trie a
Node []

-------------------------------------------------------------------------------
insert :: Path -> a -> Trie a -> Trie a
-------------------------------------------------------------------------------
insert :: forall a. Path -> a -> Trie a -> Trie a
insert []     a
v (Node [Branch a]
ts) = forall a. [Branch a] -> Trie a
Node (forall a. a -> Branch a
Val a
v forall a. a -> [a] -> [a]
: [Branch a]
ts)
insert (Int
i:Path
is) a
v (Node [Branch a]
ts) = forall a. [Branch a] -> Trie a
Node (forall a. Int -> Path -> a -> [Branch a] -> [Branch a]
insertKey Int
i Path
is a
v [Branch a]
ts)


-------------------------------------------------------------------------------
fromList :: [(Path, a)] -> Trie a
-------------------------------------------------------------------------------
fromList :: forall a. [(Path, a)] -> Trie a
fromList = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' (\Trie a
t (Path
k, a
v) -> forall a. Path -> a -> Trie a -> Trie a
insert Path
k a
v Trie a
t) forall a. Trie a
empty

-- i=3
-- 0 1 2 3 4 5 6

insertKey :: Key -> Path -> a -> [Branch a] -> [Branch a]
insertKey :: forall a. Int -> Path -> a -> [Branch a] -> [Branch a]
insertKey Int
i Path
is a
v bs :: [Branch a]
bs@((Bind Int
j Trie a
tj) : [Branch a]
bs')
  | Int
i forall a. Eq a => a -> a -> Bool
== Int
j              = forall a. Int -> Trie a -> Branch a
Bind Int
i (forall a. Path -> a -> Trie a -> Trie a
insert Path
is a
v Trie a
tj) forall a. a -> [a] -> [a]
: [Branch a]
bs'
  | Int
i forall a. Ord a => a -> a -> Bool
<  Int
j              = forall a. Int -> Trie a -> Branch a
Bind Int
i (forall a. Path -> a -> Trie a
pathTrie Path
is a
v)  forall a. a -> [a] -> [a]
: [Branch a]
bs
insertKey Int
i Path
is a
v (Branch a
b:[Branch a]
bs) = Branch a
b forall a. a -> [a] -> [a]
: forall a. Int -> Path -> a -> [Branch a] -> [Branch a]
insertKey Int
i Path
is a
v [Branch a]
bs
insertKey Int
i Path
is a
v []     = [ forall a. Int -> Trie a -> Branch a
Bind Int
i (forall a. Path -> a -> Trie a
pathTrie Path
is a
v) ]

pathTrie :: Path -> a -> Trie a
pathTrie :: forall a. Path -> a -> Trie a
pathTrie []     a
v = forall a. [Branch a] -> Trie a
Node [forall a. a -> Branch a
Val a
v]
pathTrie (Int
i:Path
is) a
v = forall a. [Branch a] -> Trie a
Node [forall a. Int -> Trie a -> Branch a
Bind Int
i (forall a. Path -> a -> Trie a
pathTrie Path
is a
v)]

-------------------------------------------------------------------------------
fold :: (acc -> Path -> a -> acc) -> acc -> Trie a -> acc
-------------------------------------------------------------------------------
fold :: forall acc a. (acc -> Path -> a -> acc) -> acc -> Trie a -> acc
fold = forall a. HasCallStack => a
undefined

-------------------------------------------------------------------------------
foldM :: (Monad m) => (acc -> Path -> a -> m acc) -> acc -> Trie a -> m acc
-------------------------------------------------------------------------------
foldM :: forall (m :: * -> *) acc a.
Monad m =>
(acc -> Path -> a -> m acc) -> acc -> Trie a -> m acc
foldM = forall a. HasCallStack => a
undefined


instance Show a => PPrint (Trie a) where
  pprintTidy :: Tidy -> Trie a -> Doc
pprintTidy Tidy
_ = forall a. Show a => a -> Doc
Misc.tshow


{-

     e1
        <----
 === e2
        <----
 === e3
        <----
   ? e4
..

 -}

-------------------------------------------------------------------------------
-- | Examples
-------------------------------------------------------------------------------
{-
Lets represent

    1,2,3,Z ->
    1,2,3,4 -> A
    1,2,3,5 -> B
    1,6     -> C

    insert [1, 2]    "A"
  . insert [1,2,3,4] "B"
  . insert [1,2,3,5] "C"
  . insert [1,6]     "D"
  $ empty

as the `trie`

     1 -> 2 -----------> A
           `-> 3 -> 4 -> B
     |         ` -> 5 -> C
     `-> 6 ------------> D

-}

-- >>> _example0 == _example1
-- True

_example0 :: Trie Char
_example0 :: Trie Char
_example0 = forall a. [Branch a] -> Trie a
Node [forall a. Int -> Trie a -> Branch a
Bind Int
1 Trie Char
n1]
  where
    n1 :: Trie Char
n1   = forall a. [Branch a] -> Trie a
Node [forall a. Int -> Trie a -> Branch a
Bind Int
2 Trie Char
n2, forall a. Int -> Trie a -> Branch a
Bind Int
6 Trie Char
n6]
    n2 :: Trie Char
n2   = forall a. [Branch a] -> Trie a
Node [forall a. a -> Branch a
Val Char
'A'  , forall a. Int -> Trie a -> Branch a
Bind Int
3 Trie Char
n3]
    n3 :: Trie Char
n3   = forall a. [Branch a] -> Trie a
Node [forall a. Int -> Trie a -> Branch a
Bind Int
4 Trie Char
n4, forall a. Int -> Trie a -> Branch a
Bind Int
5 Trie Char
n5]
    n4 :: Trie Char
n4   = forall a. [Branch a] -> Trie a
Node [forall a. a -> Branch a
Val Char
'B']
    n5 :: Trie Char
n5   = forall a. [Branch a] -> Trie a
Node [forall a. a -> Branch a
Val Char
'C']
    n6 :: Trie Char
n6   = forall a. [Branch a] -> Trie a
Node [forall a. a -> Branch a
Val Char
'D']


_example1 :: Trie Char
_example1 :: Trie Char
_example1 = forall a. Path -> a -> Trie a -> Trie a
insert [Int
1, Int
2]    Char
'A'
         forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Path -> a -> Trie a -> Trie a
insert [Int
1,Int
2,Int
3,Int
4] Char
'B'
         forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Path -> a -> Trie a -> Trie a
insert [Int
1,Int
2,Int
3,Int
5] Char
'C'
         forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Path -> a -> Trie a -> Trie a
insert [Int
1,Int
6]     Char
'D'
         forall a b. (a -> b) -> a -> b
$ forall a. Trie a
empty