toktok-0.5: Toktok/Trie.hs

module Toktok.Trie where

import Data.Map (Map)
import qualified Data.Map as Map
import Data.Char (toLower, toUpper, isUpper)

data Trie = Trie Bool (Map Char Trie)

emptyTrie = Trie False Map.empty 

mkTrie = fromList

fromList :: [String] -> Trie
fromList [] = emptyTrie
fromList (w:ws) = addWord w $ fromList ws

addWord :: String -> Trie -> Trie
addWord []     (Trie _ m) 
   = Trie True m
addWord (c:cs) (Trie b m) | Map.member c m
   = Trie b $ Map.update (return . addWord cs) c m
addWord (c:cs) (Trie b m) 
   = Trie b $ Map.insert c (addWord cs emptyTrie) m

apply :: Trie -> String -> [[String]]
apply trie = apply' trie []
   where apply' (Trie True _) w [] = [[reverse w]]
         apply' (Trie False _) w [] = []
         apply' (Trie True m) w (c:cs) = (apply' trie "" (c:cs) 
                                          >>= return . (reverse w:))
                                      ++ apply'' m c cs w
         apply' (Trie False m) w (c:cs) = apply'' m c cs w
         apply'' m c cs w = case Map.lookup c m of
                             Nothing -> []
                             Just n -> apply' n (c:w) cs
                         ++ if isUpper c
                             then case Map.lookup (toLower c) m of
                                       Nothing -> []
                                       Just n -> apply' n (toLower c:w) cs
                             else []
         

test :: [[String]]
test = apply t "aabbbc"
  where t = fromList ["a", "ab", "bb", "c", "b"]