module Data.Trie ( empty, insert, fromString, fromList
, toList, lookupPrefix, forcedNext, Trie
, possibleSuffixes, certainSuffix
) where
import Control.Monad
import Data.Binary
import qualified Data.Map as Map
#if __GLASGOW_HASKELL__ < 708
import Data.DeriveTH
#else
import GHC.Generics (Generic)
#endif
data Trie = Trie Bool (Map.Map Char Trie) deriving (Show, Eq)
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 (Just . maybe (fromString xs) (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
#if __GLASGOW_HASKELL__ < 708
$(derive makeBinary ''Trie)
#else
deriving instance Generic Trie
instance Binary Trie
#endif