module Data.Binary.Defer.Trie(
Trie, newTrie, lookupTrie, unionsWith
) where
import Data.Binary.Defer
import Data.List
import Data.Maybe
import General.Code
data Trie a = Trie (Maybe a) [(Char,Trie a)]
deriving Eq
newTrie :: [(String,a)] -> Trie a
newTrie = newTrieOrdered . sortBy (compare `on` fst)
newTrieOrdered :: [(String,a)] -> Trie a
newTrieOrdered xs = Trie
(fmap snd $ listToMaybe as)
(map f $ groupBy ((==) `on` (head . fst)) bs)
where
(as,bs) = span (null . fst) xs
f xs = (head $ fst $ head xs, newTrieOrdered [(a,b) | (_:a,b) <- xs])
lookupTrie :: String -> Trie a -> Maybe a
lookupTrie [] (Trie a b) = a
lookupTrie (x:xs) (Trie a b) = lookup x b >>= lookupTrie xs
unionsWith :: (a -> a -> a) -> [Trie a] -> Trie a
unionsWith f [] = newTrie []
unionsWith f [x] = x
unionsWith f (x:xs) = foldl (unionWith f) x xs
unionWith :: (a -> a -> a) -> Trie a -> Trie a -> Trie a
unionWith merge (Trie x xs) (Trie y ys) = Trie xy $ f xs ys
where
xy = case (x,y) of
(Just x, Nothing) -> Just x
(Nothing, Just y) -> Just y
(Just x, Just y) -> Just $ merge x y
(Nothing,Nothing) -> Nothing
f ((x1,x2):xs) ((y1,y2):ys) = case compare x1 y1 of
EQ -> (x1,unionWith merge x2 y2) : f xs ys
LT -> (x1,x2) : f xs ((y1,y2):ys)
GT -> (y1,y2) : f ((x1,x2):xs) ys
f xs [] = xs
f [] ys = ys
instance BinaryDefer a => BinaryDefer (Trie a) where
put (Trie a b) = putDefer $ put2 a b
get = getDefer $ get2 Trie
instance Show a => Show (Trie a) where
show = unlines . f ""
where
f str (Trie res xs) =
[show str ++ " = " ++ show r | Just r <- [res]] ++
concat [f (str ++ [c]) t | (c,t) <- xs]