module Agda.Utils.Trie
( Trie
, empty, singleton, everyPrefix, insert, insertWith, union, unionWith
, adjust, delete
, toList, toAscList
, lookup, member, lookupPath, lookupTrie
) where
import Prelude hiding (null, lookup)
import qualified Prelude
import Data.Function
import Data.Functor
import Data.Foldable (Foldable)
import qualified Data.Maybe as Lazy
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import qualified Agda.Utils.Maybe.Strict as Strict
import Agda.Utils.Null
data Trie k v = Trie !(Strict.Maybe v) !(Map k (Trie k v))
deriving (Show, Eq, Functor, Foldable)
instance Null (Trie k v) where
empty = Trie Strict.Nothing Map.empty
null (Trie v t) = null v && null t
singletonOrEveryPrefix :: Bool -> [k] -> v -> Trie k v
singletonOrEveryPrefix _ [] !v =
Trie (Strict.Just v) Map.empty
singletonOrEveryPrefix everyPrefix (x : xs) !v =
Trie (if everyPrefix then Strict.Just v else Strict.Nothing)
(Map.singleton x (singletonOrEveryPrefix everyPrefix xs v))
singleton :: [k] -> v -> Trie k v
singleton = singletonOrEveryPrefix False
everyPrefix :: [k] -> v -> Trie k v
everyPrefix = singletonOrEveryPrefix True
union :: (Ord k) => Trie k v -> Trie k v -> Trie k v
union = unionWith const
unionWith :: (Ord k) => (v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
unionWith f (Trie v ss) (Trie w ts) =
Trie (Strict.unionMaybeWith f v w) (Map.unionWith (unionWith f) ss ts)
insert :: (Ord k) => [k] -> v -> Trie k v -> Trie k v
insert k v t = union (singleton k v) t
insertWith :: (Ord k) => (v -> v -> v) -> [k] -> v -> Trie k v -> Trie k v
insertWith f k v t = unionWith f (singleton k v) t
delete :: Ord k => [k] -> Trie k v -> Trie k v
delete path = adjust path (const Strict.Nothing)
adjust ::
Ord k =>
[k] -> (Strict.Maybe v -> Strict.Maybe v) -> Trie k v -> Trie k v
adjust path f t@(Trie v ts) =
case path of
[] -> Trie (f v) ts
k : ks | Just s <- Map.lookup k ts -> Trie v $ Map.insert k (adjust ks f s) ts
_ -> t
toList :: Ord k => Trie k v -> [([k],v)]
toList = toAscList
toAscList :: Ord k => Trie k v -> [([k],v)]
toAscList (Trie mv ts) = Strict.maybeToList (([],) <$> mv) ++
[ (k:ks, v)
| (k, t) <- Map.toAscList ts
, (ks, v) <- toAscList t
]
lookup :: Ord k => [k] -> Trie k v -> Maybe v
lookup [] (Trie v _) = Strict.toLazy v
lookup (k : ks) (Trie _ ts) = case Map.lookup k ts of
Nothing -> Nothing
Just t -> lookup ks t
member :: Ord k => [k] -> Trie k v -> Bool
member ks t = Lazy.isJust (lookup ks t)
lookupPath :: Ord k => [k] -> Trie k v -> [v]
lookupPath xs (Trie v cs) = case xs of
[] -> Strict.maybeToList v
x : xs -> Strict.maybeToList v ++
maybe [] (lookupPath xs) (Map.lookup x cs)
lookupTrie :: Ord k => [k] -> Trie k v -> Trie k v
lookupTrie [] t = t
lookupTrie (k : ks) (Trie _ cs) = maybe empty (lookupTrie ks) (Map.lookup k cs)