{-# OPTIONS -Wno-incomplete-uni-patterns #-}
{-# LANGUAGE TemplateHaskell #-}

module SimpleTrie where

import Control.Lens
import qualified Data.Map.Strict as M

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

-- # OPTIONS -Wno-incomplete-uni-patterns #-}
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