WrappedTrie.hs

{-# LANGUAGE GADTs #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE UndecidableInstances #-}

import Control.Comonad.Cofree
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Control.Lens hiding ((:<))
import Data.Foldable
import Data.Maybe

data Trie a where
  Trie :: { runTrie :: Cofree (Map a) Bool } -> Trie [a]

trie :: Iso (Trie [a]) (Trie [b]) (Cofree (Map a) Bool) (Cofree (Map b) Bool)
trie = iso runTrie Trie

word :: Ord a => [a] -> Lens' (Trie [a]) Bool
word xs = trie . foldr
   (\x r -> _unwrap . at x . anon (False :< mempty) (\(v :< m) -> not v && null m) . r)
   _extract
   xs
    
instance Foldable Trie where
  foldr f b (Trie xs) = frt f xs b where
    frt f (e :< xs) b = if e then f [] bs else bs where 
      bs = Map.foldrWithKey (\k -> frt (f . (:) k)) b xs

  foldMap f (Trie xs) = fmt f xs where
    fmt f ~(e :< xs) = (if e then f [] else mempty) <> Map.foldMapWithKey (\k -> fmt (f . (:) k)) xs
  
instance (Show a, as ~ [a]) => Show (Trie as) where
  showsPrec n xs = showParen (n > 10) (showString "fromList " . showsPrec 11 (toList xs))
  
insert :: Ord a => [a] -> Trie [a] -> Trie [a]
insert xs = trie %~ ins xs
  where
    ins :: Ord a => [a] -> Cofree (Map a) Bool -> Cofree (Map a) Bool
    ins = foldr (\x r -> _unwrap . at x %~ (Just . r . fromMaybe (False :< Map.empty))) (_extract .~ True)
    
instance (Eq a, as ~ [a]) => Eq (Trie as) where
  Trie xs == Trie ys = xs == ys
  
instance (Ord a, as ~ [a]) => Ord (Trie as) where
  compare (Trie xs) (Trie ys) = compare xs ys
  Trie xs <= Trie ys = xs <= ys
  
instance (Ord a, as ~ [a]) => Semigroup (Trie as) where
  xs <> ys = Trie (cmb (view trie xs) (view trie ys))
    where
      cmb ~(x :< xs) ~(y :< ys) = (x || y) :< Map.unionWith cmb xs ys
      
instance (Ord a, as ~ [a]) => Monoid (Trie as) where
  mempty = Trie (False :< Map.empty)

fromList :: Ord a => [[a]] -> Trie [a]
fromList = foldr insert mempty

prefix :: Ord a => [a] -> Lens' (Trie [a]) (Trie [a])
prefix xs = trie . 
  foldr
   (\x r -> _unwrap . at x . anon (False :< mempty) (\(v :< m) -> not v && null m) . r)
   (from trie)
   xs
   
mapTrie :: Ord b => ([a] -> [b]) -> Trie [a] -> Trie [b]
mapTrie f = fromList . map f . toList

bindTrie :: Ord b => ([a] -> Trie [b]) -> Trie [a] -> Trie [b]
bindTrie = foldMap