t.hs

{-# LANGUAGE DeriveFunctor       #-}
{-# LANGUAGE FlexibleInstances   #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE LambdaCase          #-}

import           Data.Functor.Foldable hiding (Foldable)
import qualified Data.Functor.Foldable as Functor
import           Data.Traversable
import qualified Data.Map as Map
import           Data.Maybe
import           Data.Foldable
import           Data.Coerce
import           Data.Monoid

bool :: a -> a -> Bool -> a
bool t _ True  = t
bool _ f False = f

zipo :: (Functor.Foldable g, Functor.Foldable h) 
     => (Base g (h -> c) -> Base h h -> c)
     -> g -> h -> c
zipo alg = cata zalg where zalg x = alg x . project

data Trie a where
  Trie :: Bool -> Map.Map a (Trie [a]) -> Trie [a]

data TrieF a r = 
  TrieF { endHereF :: Bool
        , childrenF :: Map.Map a r
        } deriving Functor
        
type instance Base (Trie [a]) = TrieF a
     
instance Functor.Foldable (Trie [a]) where
  project (Trie e m) = TrieF e m
  
member :: Ord a => [a] -> Trie [a] -> Bool
member = zipo alg where
  alg Nil = endHereF 
  alg (Cons x xs) = maybe False xs . Map.lookup x . childrenF

empty :: Trie [a]
empty = Trie False Map.empty
      
insert :: Ord a => [a] -> Trie [a] -> Trie [a]
insert = foldr f (\(Trie _ m) -> Trie True m) where
  f :: Ord a => a -> (Trie [a] -> Trie [a]) -> Trie [a] -> Trie [a]
  f x a (Trie e m) = Trie e (Map.alter (Just . a . fromMaybe empty) x m)

fromList :: (Foldable f, Ord a) => f [a] -> Trie [a]
fromList = foldr insert empty

instance Show a => Show (Trie [a]) where
  show = showString "fromList " . show . toList
  
instance Foldable Trie where
  foldr f i t@(Trie _ _) = foldrTrie (:) [] f i t
  foldMap g t@(Trie _ _) = foldMapTrie (Endo #. (:)) (g . flip appEndo []) t

foldrTrie :: (a -> b -> b) -> b -> (b -> c -> c) -> c -> Trie [a] -> c
foldrTrie g j h k t = cata alg t h k where
  alg (TrieF e m) f i = Map.foldrWithKey r l m where
        r key val = val (f . g key)
        l = bool (f j i) i e
        
foldMapTrie :: (Monoid m, Monoid n) => (a -> m) -> (m -> n) -> Trie [a] -> n
foldMapTrie g h t = cata alg t h where
  alg (TrieF e m) f = bool (mappend (f mempty) s) s e where
    s = Map.foldMapWithKey (r . g) m
    r key val = val (f . mappend key)

(#.) :: Coercible b c => (b -> c) -> (a -> b) -> (a -> c)
(#.) _f = coerce
{-# INLINE (#.) #-}