KvTrie.hs

import Data.List(groupBy, sortBy)
import Data.Maybe(listToMaybe)
import Control.Applicative((<|>))

data Trie k v = RootTrie [Trie k v]
              | NodeTrie k [Trie k v]
              | EndTrie v
              | EmptyTrie

toTrie :: (Eq k, Ord k) => [([k], v)] -> Trie k v
toTrie = toTrieRoot . regroup . sortBy (\a b -> fst a `compare` fst b)

toTrieRoot :: Eq k => [[([k], v)]] -> Trie k v
toTrieRoot = RootTrie . map toTrieRoot'

toTrieRoot' :: Eq k => [([k], v)] -> Trie k v
toTrieRoot' l = NodeTrie name (map toTrieDeep (regroup ts))
  where name = head (fst (head l))
        ts = nextLevel l

toTrieDeep :: Eq k => [([k], v)] -> Trie k v
toTrieDeep l = if 1 == length l && null (fst (head l)) then EndTrie (snd $ head l) else maybe EmptyTrie createNode getName
  where getName = join $ fmap (listToMaybe . fst) (listToMaybe l)
        createNode name = NodeTrie name (map toTrieDeep (regroup (nextLevel l)))

nextLevel :: [([k], v)] -> [([k], v)]
nextLevel = map (\(n, v) -> (tail n, v))

regroup :: Eq k => [([k], v)] -> [[([k], v)]]
regroup = groupBy (\a b -> listToMaybe (fst a) == listToMaybe (fst b))

findTrie :: Eq k => Trie k v -> [k] -> Maybe v
findTrie t s = case t of
  NodeTrie k n  -> join $ fmap (\c -> if c == k then fn ns n else Nothing) nc
  EndTrie v     -> if null s then Just v else Nothing
  RootTrie n    -> fn s n
  EmptyTrie     -> Nothing
  where nc = listToMaybe s
        ns = tail s
        fn :: Eq k => [k] -> [Trie k v] -> Maybe v
        fn fs n = foldr (<|>) Nothing $ map (flip findTrie fs) n

-- multiple
data Trie k v = RootTrie [Trie k v]
              | NodeTrie k [Trie k v]
              | EndTrie [v]
              | EmptyTrie
              deriving Show

toTrieDeep :: Eq k => [([k], v)] -> Trie k v
toTrieDeep l = if 1 <= length l && null (fst (head l))
               then EndTrie (map snd l)
               else maybe EmptyTrie createNode getName
  where getName = join $ fmap (listToMaybe . fst) (listToMaybe l)
        createNode name = NodeTrie name (map toTrieDeep (regroup (nextLevel l)))

findsTrie :: Eq k => Trie k v -> [k] -> [v]
findsTrie t s = case t of
  NodeTrie k n  -> join $ fmap (\c -> if c == k then fn ns n else []) $ maybeToList nc
  EndTrie v     -> v
  RootTrie n    -> fn s n
  EmptyTrie     -> []
  where nc = listToMaybe s
        ns = tail s
        fn :: Eq k => [k] -> [Trie k v] -> [v]
        fn fs ft = join [findsTrie ftt fss | ftt <- ft, fss <- tails fs]

-- Samples
find :: Table -> String -> String
find t ext = maybe "UNKNOWN" id $ findTrie t ext

countElements :: Trie k v -> Int
countElements t = case t of
  NodeTrie k n  -> r n
  EndTrie v     -> 0
  RootTrie n    -> r n - 1
  EmptyTrie     -> 0
  where r n = foldr (+) 1 (map countElements n)