Trie.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE Strict #-}

module Data.Trie
    ( Trie(..)
    , singleton
    , lookup
    , toList
    , buildCompact
    ) where

import Control.Applicative
import Control.Monad
import Control.Monad.IO.Class
import Data.Binary
import qualified Data.DList as DList
import Data.DList (DList)
import qualified Data.Foldable as F
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import GHC.Compact
import GHC.Generics

import Prelude hiding (lookup)

-- | A tree with branches keyed by a 'Map', which is useful for storing
--   and searching list prefixes in a compact way. Each node is marked
--   with a type signifying list enedings.
--
--   Example: @"he"@ and @"hi"@ stored in a @'Trie' 'Char' '()'@ would look like:
--
--   @Nothing :< toList [('h', Nothing :< toList [('e', Just () :< toList []),('i', Just () :< toList [])])]@
data Trie k a = Maybe a :< Map k (Trie k a)
    deriving (Show, Eq, Ord, Functor, Foldable, Traversable, Generic)

instance Ord k => Semigroup (Trie k a) where
    ma1 :< f1 <> ma2 :< f2 = (ma2 <|> ma1) :< Map.unionWith (<>) f1 f2

instance Ord k => Monoid (Trie k a) where
    mempty = Nothing :< Map.empty

instance (Binary k, Binary a) => Binary (Trie k a)

-- | Walk the 'Trie' using the given key list looking for a matching node.
lookup :: Ord k => Trie k a -> [k] -> Maybe (Trie k a)
lookup (_ :< mapK) (k:ks) =
    Map.lookup k mapK >>= (\c -> lookup c ks)
lookup t [] = Just t

-- | Return the path and value of every node marked with 'Just'
toList :: Trie k a -> [([k],a)]
toList = DList.toList . go DList.empty DList.empty
  where
    go :: DList ([k], a) -> DList k -> Trie k a -> DList ([k], a)
    go justs path (b :< m) =
        let justs' = case b of
                Just x -> justs `DList.snoc` (DList.toList path, x)
                Nothing -> justs
        in justs' <> F.foldMap'
            (\(x,t) -> go justs' (path `DList.snoc` x) t)
            (Map.toAscList m)

-- | Create a simple 'Trie' from a key list and a value.
singleton :: [k] -> a -> Trie k a
singleton l x = go l
  where
    go (k:ks) = Nothing :< Map.singleton k (go ks)
    go []     = Just x :< Map.empty

-- | Build a 'Trie' within a compact region
buildCompact :: (MonadIO m, Ord k)
    => (a -> ([k], b))
    -> [a]
    -> m (Compact (Trie k b))
buildCompact _ [] = liftIO $ compact mempty
buildCompact f (x:xs) = do
    ct <- buildCompact f xs
    cs <- liftIO $ compactAdd ct (f x)
    let (t,(s,b)) = (getCompact ct, getCompact cs)
    cq <- liftIO $ compactAdd ct $ t <> singleton s b
    pure cq