{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE NoImplicitPrelude #-}
module Salak.Trie(
Trie
, getPrimitive
, getMap
, empty
, singleton
, Salak.Trie.null
, member
, lookup
, subTrie
, subTries
, insert
, modify
, modify'
, update
, alter
, Salak.Trie.toList
, fromList
, filter
, unionWith
, unionWith'
) where
import Control.Applicative (pure, (<*>))
import Data.Bool
import qualified Data.DList as D
import Data.Eq
import Data.Foldable (Foldable (..))
import Data.Function
import Data.Functor
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import Data.List (concat, intercalate, map, (++))
import Data.Maybe
import Data.Traversable
import Data.Tuple (uncurry)
import Salak.Internal.Key
import Text.Show (Show (..))
data Trie v = Trie !(Maybe v) !(HashMap Key (Trie v)) deriving (Eq, Functor)
instance Show v => Show (Trie v) where
show t = intercalate "\n" $ map (\(k,v)-> show k ++ ":" ++ show v) $ Salak.Trie.toList t
instance Foldable Trie where
foldr f b (Trie v m) = foldr (flip (foldr f)) (go v) m
where
go (Just x) = f x b
go _ = b
instance Traversable Trie where
traverse f (Trie v m) = Trie <$> go v <*> traverse (traverse f) m
where
go (Just x) = Just <$> f x
go _ = pure Nothing
singleton :: v -> Trie v
singleton v = Trie (Just v) HM.empty
empty :: Trie v
empty = Trie Nothing HM.empty
getPrimitive :: Trie v -> Maybe v
getPrimitive (Trie v _) = v
getMap :: Trie v -> HashMap Key (Trie v)
getMap (Trie _ m) = m
null :: Trie v -> Bool
null (Trie Nothing e) = HM.null e
null _ = False
member :: Eq v => Keys -> Trie v -> Bool
member k t = isJust (lookup k t)
subTrie :: Key -> Trie v -> Trie v
subTrie key = fromMaybe empty . HM.lookup key . getMap
subTries :: Keys -> Trie v -> Trie v
subTries ks v = foldl (flip subTrie) v (toKeyList ks)
lookup :: Eq v => Keys -> Trie v -> Maybe v
lookup keys = getPrimitive . subTries keys
insert :: Eq v => Keys -> v -> Trie v -> Trie v
insert ks v = alter (\_ -> Just v) ks
modify :: Eq v => Key -> (Trie v -> Trie v) -> Trie v -> Trie v
modify k f (Trie v m) = Trie v $ HM.alter (go . f . fromMaybe empty) k m
where
go x = if x == empty then Nothing else Just x
modify' :: Eq v => Keys -> (Trie v -> Trie v) -> Trie v -> Trie v
modify' ks f = foldr modify f (toKeyList ks)
update :: Eq v => (Maybe v -> Maybe v) -> Trie v -> Trie v
update f = alter f mempty
alter :: Eq v => (Maybe v -> Maybe v) -> Keys -> Trie v -> Trie v
alter f keys = modify' keys (\(Trie a b) -> Trie (f a) b)
toList :: Trie v -> [(Keys, v)]
toList = go D.empty
where
go p (Trie (Just v) m) = (Keys p, v) : g2 p m
go p (Trie _ m) = g2 p m
g2 p m = concat $ g3 p <$> HM.toList m
g3 p (k,t) = go (D.snoc p k) t
fromList :: Eq v => [(Keys, v)] -> Trie v
fromList = foldr (uncurry insert) empty
filter :: Eq v => (v -> Bool) -> Trie v -> Trie v
filter f (Trie v m) = if ok v then Trie v go else Trie Nothing go
where
ok (Just x) = f x
ok _ = False
go = HM.mapMaybe (g2 . filter f) m
g2 x = if x == empty then Nothing else Just x
unionWith :: Eq v => (Maybe v -> Maybe v -> Maybe v) -> Trie v -> Trie v -> Trie v
unionWith f (Trie v1 m1) (Trie v2 m2) = Trie (f v1 v2) $ HM.unionWith (unionWith f) m1 m2
unionWith' :: (Maybe v -> Maybe v -> Maybe v3) -> Trie v -> Trie v -> Trie v3
unionWith' f (Trie v1 m1) (Trie v2 m2) = Trie (f v1 v2) $ foldr go HM.empty $ HM.keys $ HM.union m1 m2
where
go k =
let x1 = fromMaybe empty $ HM.lookup k m1
x2 = fromMaybe empty $ HM.lookup k m2
in HM.insert k (unionWith' f x1 x2)