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]

data Trie a 
  = Node ![Branch a]
  deriving (Trie a -> Trie a -> Bool
(Trie a -> Trie a -> Bool)
-> (Trie a -> Trie a -> Bool) -> Eq (Trie a)
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
[Trie a] -> ShowS
Trie a -> String
(Int -> Trie a -> ShowS)
-> (Trie a -> String) -> ([Trie a] -> ShowS) -> Show (Trie a)
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
(Branch a -> Branch a -> Bool)
-> (Branch a -> Branch a -> Bool) -> Eq (Branch a)
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
[Branch a] -> ShowS
Branch a -> String
(Int -> Branch a -> ShowS)
-> (Branch a -> String) -> ([Branch a] -> ShowS) -> Show (Branch a)
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 :: Trie a
empty = [Branch a] -> Trie a
forall a. [Branch a] -> Trie a
Node []

-------------------------------------------------------------------------------
insert :: Path -> a -> Trie a -> Trie a 
-------------------------------------------------------------------------------
insert :: Path -> a -> Trie a -> Trie a
insert []     a
v (Node [Branch a]
ts) = [Branch a] -> Trie a
forall a. [Branch a] -> Trie a
Node ((a -> Branch a
forall a. a -> Branch a
Val a
v) Branch a -> [Branch a] -> [Branch a]
forall a. a -> [a] -> [a]
: [Branch a]
ts) 
insert (Int
i:Path
is) a
v (Node [Branch a]
ts) = [Branch a] -> Trie a
forall a. [Branch a] -> Trie a
Node (Int -> Path -> a -> [Branch a] -> [Branch a]
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 :: [(Path, a)] -> Trie a
fromList = (Trie a -> (Path, a) -> Trie a) -> Trie a -> [(Path, a)] -> Trie a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' (\Trie a
t (Path
k, a
v) -> Path -> a -> Trie a -> Trie a
forall a. Path -> a -> Trie a -> Trie a
insert Path
k a
v Trie a
t) Trie a
forall a. Trie a
empty 

-- i=3 
-- 0 1 2 3 4 5 6 
  
insertKey :: Key -> Path -> a -> [Branch a] -> [Branch a]
insertKey :: 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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j              = Int -> Trie a -> Branch a
forall a. Int -> Trie a -> Branch a
Bind Int
i (Path -> a -> Trie a -> Trie a
forall a. Path -> a -> Trie a -> Trie a
insert Path
is a
v Trie a
tj) Branch a -> [Branch a] -> [Branch a]
forall a. a -> [a] -> [a]
: [Branch a]
bs' 
  | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<  Int
j              = Int -> Trie a -> Branch a
forall a. Int -> Trie a -> Branch a
Bind Int
i (Path -> a -> Trie a
forall a. Path -> a -> Trie a
pathTrie Path
is a
v)  Branch a -> [Branch a] -> [Branch a]
forall a. a -> [a] -> [a]
: [Branch a]
bs 
insertKey Int
i Path
is a
v (Branch a
b:[Branch a]
bs) = Branch a
b Branch a -> [Branch a] -> [Branch a]
forall a. a -> [a] -> [a]
: Int -> Path -> a -> [Branch a] -> [Branch 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 []     = [ Int -> Trie a -> Branch a
forall a. Int -> Trie a -> Branch a
Bind Int
i (Path -> a -> Trie a
forall a. Path -> a -> Trie a
pathTrie Path
is a
v) ] 

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

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

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


instance Show a => PPrint (Trie a) where 
  pprintTidy :: Tidy -> Trie a -> Doc
pprintTidy Tidy
_ = Trie a -> Doc
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 = [Branch Char] -> Trie Char
forall a. [Branch a] -> Trie a
Node [Int -> Trie Char -> Branch Char
forall a. Int -> Trie a -> Branch a
Bind Int
1 Trie Char
n1]
  where
    n1 :: Trie Char
n1   = [Branch Char] -> Trie Char
forall a. [Branch a] -> Trie a
Node [Int -> Trie Char -> Branch Char
forall a. Int -> Trie a -> Branch a
Bind Int
2 Trie Char
n2, Int -> Trie Char -> Branch Char
forall a. Int -> Trie a -> Branch a
Bind Int
6 Trie Char
n6]
    n2 :: Trie Char
n2   = [Branch Char] -> Trie Char
forall a. [Branch a] -> Trie a
Node [Char -> Branch Char
forall a. a -> Branch a
Val Char
'A'  , Int -> Trie Char -> Branch Char
forall a. Int -> Trie a -> Branch a
Bind Int
3 Trie Char
n3]
    n3 :: Trie Char
n3   = [Branch Char] -> Trie Char
forall a. [Branch a] -> Trie a
Node [Int -> Trie Char -> Branch Char
forall a. Int -> Trie a -> Branch a
Bind Int
4 Trie Char
n4, Int -> Trie Char -> Branch Char
forall a. Int -> Trie a -> Branch a
Bind Int
5 Trie Char
n5]
    n4 :: Trie Char
n4   = [Branch Char] -> Trie Char
forall a. [Branch a] -> Trie a
Node [Char -> Branch Char
forall a. a -> Branch a
Val Char
'B']
    n5 :: Trie Char
n5   = [Branch Char] -> Trie Char
forall a. [Branch a] -> Trie a
Node [Char -> Branch Char
forall a. a -> Branch a
Val Char
'C']
    n6 :: Trie Char
n6   = [Branch Char] -> Trie Char
forall a. [Branch a] -> Trie a
Node [Char -> Branch Char
forall a. a -> Branch a
Val Char
'D'] 


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