module Data.Trie where
import Data.Binary
import Data.DeriveTH
import qualified Data.Map as Map
import Control.Monad
data Trie = Trie Bool (Map.Map Char Trie) deriving (Show)
empty :: Trie
empty = Trie False Map.empty
insert :: String -> Trie -> Trie
insert [] (Trie _ m) = Trie True m
insert (x:xs) (Trie b m) = Trie b $ Map.alter (maybe (Just $ fromString xs) (Just . insert xs)) x m
fromString :: String -> Trie
fromString = foldr (\x xs -> Trie False (Map.singleton x xs)) (Trie True Map.empty)
fromList :: [String] -> Trie
fromList = foldr insert empty
toList :: Trie -> [String]
toList (Trie b m) =
if b then "":expand
else expand
where expand = [ char:word | (char, trie) <- Map.toList m,
word <- toList trie ]
lookupPrefix :: (MonadPlus m) => String -> Trie -> m Trie
lookupPrefix [] trie = return trie
lookupPrefix (x:xs) (Trie _ m) = liftMaybe (Map.lookup x m) >>= lookupPrefix xs
liftMaybe :: MonadPlus m => Maybe a -> m a
liftMaybe Nothing = mzero
liftMaybe (Just x) = return x
forcedNext :: Trie -> String
forcedNext (Trie _ m) =
if length ls == 1 then
let (char, trie) = head ls in
char:forcedNext trie
else []
where ls = Map.toList m
possibleSuffixes :: String -> Trie -> [String]
possibleSuffixes prefix fulltrie =
lookupPrefix prefix fulltrie >>= toList
certainSuffix :: String -> Trie -> String
certainSuffix prefix fulltrie =
lookupPrefix prefix fulltrie >>= forcedNext
$(derive makeBinary ''Trie)