module System.Console.CmdTheLine.Trie
( add, lookup, empty, isEmpty, fromList, ambiguities
, Trie, LookupFail(..)
) where
import Prelude hiding ( lookup )
import qualified Data.Map as M
import Data.List (foldl')
type CMap = M.Map Char
data Value a = Pre a
| Key a
| Amb
| Nil
deriving ( Eq )
data Trie a = Trie
{ val :: Value a
, nexts :: CMap (Trie a)
} deriving ( Eq )
data LookupFail = Ambiguous | NotFound deriving (Show)
empty :: Trie a
empty = Trie Nil M.empty
isEmpty :: Eq a => Trie a -> Bool
isEmpty = (== empty)
add :: String -> a -> Trie a -> Trie a
add k v t = go t k
where
go t s = case s of
[] -> Trie (Key v) next
c:rest ->
let t' = maybe empty id $ M.lookup c next
newNexts = M.insert c (go t' rest) next
in Trie newVal newNexts
where
next = nexts t
newVal = case val t of
Amb -> Amb
Pre _ -> Amb
v'@(Key _) -> v'
Nil -> Pre v
findNode :: Trie a -> String -> Maybe (Trie a)
findNode t = foldl' go (Just t)
where
go acc c = M.lookup c . nexts =<< acc
lookup :: String -> Trie a -> Either LookupFail a
lookup k t = case findNode t k of
Nothing -> Left NotFound
Just t' -> case val t' of
Key v -> Right v
Pre v -> Right v
Amb -> Left Ambiguous
Nil -> Left NotFound
ambiguities :: Trie a -> String -> [String]
ambiguities t pre = case findNode t pre of
Nothing -> []
Just t' -> case val t' of
Amb -> go [] pre $ M.toList (nexts t') : []
_ -> []
where
go acc pre assocs = case assocs of
[] -> error "saw lone empty list while searching for ambiguities"
[[]] -> acc
[] : rest -> go acc (init pre) rest
_ -> descend assocs
where
descend ((top : bottom) : rest) = go acc' pre' assocs'
where
( c, t'' ) = top
assocs' = M.toList (nexts t'') : bottom : rest
pre' = pre ++ return c
acc' = case val t'' of
Key _ -> pre' : acc
Nil -> error "saw Nil on descent"
_ -> acc
descend _ = undefined
fromList :: [( String, a )] -> Trie a
fromList assoc = foldl' consume empty assoc
where
consume t ( k, v ) = add k v t