{-# LANGUAGE TemplateHaskell #-}
module SimpleTrie where
import Data.Map.Strict
import Control.Lens
data Trie a b = Trie
    { _trieMVal    :: Maybe b
    , _trieChildren :: Map a (Trie a b)
    }
makeLenses ''Trie

emptyTrie :: Trie a b
emptyTrie = Trie Nothing empty

singletonTrie :: Ord a => [a] -> b -> Trie a b
singletonTrie (k:ks) x = Trie Nothing $ singleton k $ singletonTrie ks x
singletonTrie [] x = Trie (Just x) empty

insertInTrie :: Ord a => [a] -> b -> Trie a b -> Trie a b
insertInTrie [] x      = trieMVal ?~ x
insertInTrie (k:ks) x  = trieChildren %~ insertWith f k (singletonTrie ks x)
  where
    f _ = insertInTrie ks x

lookupTrie :: Ord a => [a] -> Trie a b -> Maybe b
lookupTrie (k:ks) t = _trieChildren t !? k >>= lookupTrie ks
lookupTrie []     t = _trieMVal t