Trie Implementation

Trie.hs

extract :: (a -> Bool) -> [a] -> (Maybe a, [a])
extract _ [] = (Nothing, [])
extract pred (x:xs)
  | pred x = (Just x, xs)
  | otherwise = (x:) <$> extract pred xs


data Trie a = Node { datum :: Maybe a
                   , children :: [Trie a] }

instance Show a => Show (Trie a) where
  show Node{datum = d, children = kids}
    = showD d ++ (case length kids of
                    0 -> ""
                    1 -> " => " ++ show (head kids)
                    _ -> " => {" ++ intercalate " | " (show <$> kids) ++ "}")
    where showD (Just x) = show x
          showD Nothing = "Root"

getDepth :: (Ord b, Num b) => Trie a -> b
getDepth Node{children = kids}
  | null kids = 1
  | otherwise = 1 + maximum (getDepth <$> kids)

insertStr :: Eq a => Trie a -> [a] -> Trie a
insertStr root [] = root
insertStr root@Node{children = kids} xs
  | (Just node, otherKids) <- extract (\x -> datum x == Just (head xs)) kids
  = Node{datum = datum root
        ,children = insertStr node (tail xs) : otherKids}

  | otherwise
  = Node{datum = datum root
        ,children = insertStr Node{datum = Just (head xs), children = []} (tail xs) : kids}

emptyTrie :: Trie a
emptyTrie = Node{datum = Nothing, children = []}