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

emptyTrie :: Trie a b
emptyTrie = Trie Nothing M.empty

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

insertInTrie :: Ord a => [a] -> b -> Trie a b -> Trie a b
insertInTrie [] x      = trieMVal ?~ x
insertInTrie (k:ks) x  = trieChildren %~ M.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 M.!? k >>= lookupTrie ks
lookupTrie []     t = _trieMVal t

multiLookupTrie :: Ord a => [a] -> Trie a b -> [([a],b)]
multiLookupTrie  xs t@(Trie my ch)
    | null xs || null ch = maybe [] (\y -> [([],y)]) my
    | otherwise = let (z:zs) = xs in
        case M.lookup z ch of
            Nothing -> multiLookupTrie zs t
            Just t' -> map (_1 %~ (z:)) (multiLookupTrie zs t') ++ multiLookupTrie zs t

multiLookupTrieI :: Ord a => [(a,c)] -> Trie a b -> [([c],b)]
multiLookupTrieI  xs t@(Trie my ch)
    | null xs || null ch = maybe [] (\y -> [([],y)]) my
    | otherwise = let ((z1,z2):zs) = xs in
        case M.lookup z1 ch of
            Nothing -> multiLookupTrieI zs t
            Just t' -> map (_1 %~ (z2:)) (multiLookupTrieI zs t') ++ multiLookupTrieI zs t