SetTrie.hs

module SetTrie where

import Prelude hiding (lookup)
import Data.List (intercalate)
import Data.Set hiding (empty, insert, map, fromList)
import qualified Data.Set as S

data SetTrie k v = SetTrie { count :: Int, leaf :: Maybe v, branches :: [(k,SetTrie k v)] }

prefixThenIndent :: String -> [String] -> [String]
-- create a line
prefixThenIndent initial [] = [initial]
-- prefix the first line, then indent the rest
prefixThenIndent initial ls = zipWith (++) (initial : repeat spaces) ls
  where spaces = map (const ' ') initial

instance (Show k, Show v) => Show (SetTrie k v) where
  show = intercalate "\n" . nodeLines
    where nodeLines (SetTrie n l bs) = prefixThenIndent (show (n,l)) $ concatMap edgeLines bs
          edgeLines (k,t) = prefixThenIndent (" -{ " ++ show k ++ " }-> ") $ nodeLines t

empty :: SetTrie k v
empty = SetTrie 0 Nothing []

insert :: Ord k => Set k -> v -> SetTrie k v -> SetTrie k v
insert ks v (SetTrie c l bs)  | S.null ks = c' `seq` SetTrie c' (Just v) bs
                              | otherwise = c' `seq` SetTrie c' l bs'
  where c' = c + 1
        bs' = case break (flip member ks . fst) bs of
                (xs, (k,t):ys)  -> xs ++ (k, insert (delete k ks) v t) : ys
                (bs, [])        -> bs ++ (branches . foldr go (SetTrie 1 (Just v) []) $ toList ks)
        go k t = SetTrie 1 Nothing [(k,t)]

lookup :: Ord k => Set k -> SetTrie k v -> Maybe v
lookup ks (SetTrie c l bs) | S.null ks = l
                           | otherwise = foldr go Nothing bs 
  where go (k,t) l = if member k ks then lookup (delete k ks) t else l

fromList :: Ord k => [(Set k, v)] -> SetTrie k v
fromList = foldr (uncurry insert) empty 

example :: SetTrie Char Float
example = fromList  [ (S.fromList "hello", 1)
                    , (S.fromList "world", 2)
                    , (S.fromList "hell", 3)
                    , (S.fromList "war", 4)
                    ]