Trie.hs

module Trie where

data Trie a = Node Char (Trie a) (Trie a) (Trie a) (Maybe a)
              | Empty deriving (Show, Eq)

sanify :: Trie a -> Trie a
sanify (Node _ Empty Empty Empty Nothing) = Empty
sanify (Node _ Empty lo    Empty Nothing) = lo
sanify (Node _ Empty Empty hi    Nothing) = hi
sanify t = t  

update :: Trie a -> String -> (Maybe a -> Maybe a) -> Trie a
update _ [] _ = error "Can not insert an empty string to a Trie"
update Empty (x:[]) f = sanify $ Node x Empty Empty Empty (f Nothing) 
update Empty (x:xs) f = sanify $ Node x (update Empty xs f) Empty Empty Nothing
update (Node c eq lo hi val) xss@(x:xs) f = 
    sanify $ case x `compare` c of
        LT -> Node c eq (update lo xss f) hi val
        GT -> Node c eq lo (update hi xss f) val
        EQ -> case xs of
                [] -> Node c eq lo hi (f val)
                _  -> Node c (update eq xs f) lo hi val

get :: Trie a -> String -> Maybe a
get _ [] = Nothing
get Empty _ = Nothing
get (Node c eq lo hi val) xss@(x:xs) =
    case x `compare` c of
        LT -> get lo xss
        GT -> get hi xss
        EQ -> case xs of
                [] -> val
                _  -> get eq xs

insert :: Trie a -> String -> a -> Trie a
insert t s v = update t s (const $ Just v)

delete :: Trie a -> String -> Trie a
delete t s = update t s (const Nothing)

deleteAll :: Trie a -> [String] -> Trie a
deleteAll = foldl delete 

-- Highly inefficient, mainly for debugging
toList :: Trie a -> [(String, a)]
toList t = (\(s, v) -> (reverse s, v)) `map` acc t ""
    where acc Empty _ = []
          acc (Node c eq lo hi val) prefix =
            acc lo prefix 
            ++ case val of
                  Just x -> [(c:prefix, x)]
                  Nothing -> []
            ++ acc eq (c:prefix) 
            ++ acc hi prefix

fromList :: [(String, a)] -> Trie a
fromList = foldl (\t l -> uncurry (insert t) l) Empty