{-# OPTIONS_GHC -Wunused-imports #-}

-- | Strict tries (based on "Data.Map.Strict" and "Agda.Utils.Maybe.Strict").
--
-- Note that if 'delete' or 'adjust' are used, one may end up with non-canonical
-- subtries, i.e., those that do not contain any value (@Data.List.null . toList@)
-- but may be @Eq@-different to other empty tries (@/= empty@).

module Agda.Utils.Trie
  ( Trie(..)
  , empty, singleton, everyPrefix, insert, insertWith, union, unionWith
  -- Andreas 2024-11-16:
  -- Deletion does not clean up the trie, thus violates equational laws
  -- with the given Eq that does not quotient out the different representations
  -- of empty subtries.
  -- E.g. @delete k . insert k v . delete k /= delete k@.
  -- @adjust@ has a similar problem.
  , adjust, delete
  , toList, toAscList, toListOrderedBy
  , lookup, member, lookupPath, lookupTrie
  , mapSubTries, filter
  , valueAt
  ) where

import Prelude hiding (null, lookup, filter)

import Control.DeepSeq

import Data.Function (on)
import qualified Data.Maybe as Lazy
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import qualified Data.List as List

import qualified Agda.Utils.Maybe.Strict as Strict
import Agda.Utils.Null
import Agda.Utils.Lens

-- | Finite map from @[k]@ to @v@.
--
--   With the strict 'Maybe' type, 'Trie' is also strict in 'v'.
data Trie k v = Trie !(Strict.Maybe v) !(Map k (Trie k v))
  deriving
    ( Int -> Trie k v -> ShowS
[Trie k v] -> ShowS
Trie k v -> String
(Int -> Trie k v -> ShowS)
-> (Trie k v -> String) -> ([Trie k v] -> ShowS) -> Show (Trie k v)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k v. (Show v, Show k) => Int -> Trie k v -> ShowS
forall k v. (Show v, Show k) => [Trie k v] -> ShowS
forall k v. (Show v, Show k) => Trie k v -> String
$cshowsPrec :: forall k v. (Show v, Show k) => Int -> Trie k v -> ShowS
showsPrec :: Int -> Trie k v -> ShowS
$cshow :: forall k v. (Show v, Show k) => Trie k v -> String
show :: Trie k v -> String
$cshowList :: forall k v. (Show v, Show k) => [Trie k v] -> ShowS
showList :: [Trie k v] -> ShowS
Show
    , Trie k v -> Trie k v -> Bool
(Trie k v -> Trie k v -> Bool)
-> (Trie k v -> Trie k v -> Bool) -> Eq (Trie k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v. (Eq v, Eq k) => Trie k v -> Trie k v -> Bool
$c== :: forall k v. (Eq v, Eq k) => Trie k v -> Trie k v -> Bool
== :: Trie k v -> Trie k v -> Bool
$c/= :: forall k v. (Eq v, Eq k) => Trie k v -> Trie k v -> Bool
/= :: Trie k v -> Trie k v -> Bool
Eq
        -- Andreas, 2024-11-16: The derived Eq isn't extensional.
        -- It disguishes different representation of the empty Trie.
        -- (E.g. the empty map and the map that only contains Nothing values.)
    , (forall a b. (a -> b) -> Trie k a -> Trie k b)
-> (forall a b. a -> Trie k b -> Trie k a) -> Functor (Trie k)
forall a b. a -> Trie k b -> Trie k a
forall a b. (a -> b) -> Trie k a -> Trie k b
forall k a b. a -> Trie k b -> Trie k a
forall k a b. (a -> b) -> Trie k a -> Trie k b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall k a b. (a -> b) -> Trie k a -> Trie k b
fmap :: forall a b. (a -> b) -> Trie k a -> Trie k b
$c<$ :: forall k a b. a -> Trie k b -> Trie k a
<$ :: forall a b. a -> Trie k b -> Trie k a
Functor
    , (forall m. Monoid m => Trie k m -> m)
-> (forall m a. Monoid m => (a -> m) -> Trie k a -> m)
-> (forall m a. Monoid m => (a -> m) -> Trie k a -> m)
-> (forall a b. (a -> b -> b) -> b -> Trie k a -> b)
-> (forall a b. (a -> b -> b) -> b -> Trie k a -> b)
-> (forall b a. (b -> a -> b) -> b -> Trie k a -> b)
-> (forall b a. (b -> a -> b) -> b -> Trie k a -> b)
-> (forall a. (a -> a -> a) -> Trie k a -> a)
-> (forall a. (a -> a -> a) -> Trie k a -> a)
-> (forall a. Trie k a -> [a])
-> (forall a. Trie k a -> Bool)
-> (forall a. Trie k a -> Int)
-> (forall a. Eq a => a -> Trie k a -> Bool)
-> (forall a. Ord a => Trie k a -> a)
-> (forall a. Ord a => Trie k a -> a)
-> (forall a. Num a => Trie k a -> a)
-> (forall a. Num a => Trie k a -> a)
-> Foldable (Trie k)
forall a. Eq a => a -> Trie k a -> Bool
forall a. Num a => Trie k a -> a
forall a. Ord a => Trie k a -> a
forall m. Monoid m => Trie k m -> m
forall a. Trie k a -> Bool
forall a. Trie k a -> Int
forall a. Trie k a -> [a]
forall a. (a -> a -> a) -> Trie k a -> a
forall k a. Eq a => a -> Trie k a -> Bool
forall k a. Num a => Trie k a -> a
forall k a. Ord a => Trie k a -> a
forall m a. Monoid m => (a -> m) -> Trie k a -> m
forall k m. Monoid m => Trie k m -> m
forall k a. Trie k a -> Bool
forall k a. Trie k a -> Int
forall k a. Trie k a -> [a]
forall b a. (b -> a -> b) -> b -> Trie k a -> b
forall a b. (a -> b -> b) -> b -> Trie k a -> b
forall k a. (a -> a -> a) -> Trie k a -> a
forall k m a. Monoid m => (a -> m) -> Trie k a -> m
forall k b a. (b -> a -> b) -> b -> Trie k a -> b
forall k a b. (a -> b -> b) -> b -> Trie k a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall k m. Monoid m => Trie k m -> m
fold :: forall m. Monoid m => Trie k m -> m
$cfoldMap :: forall k m a. Monoid m => (a -> m) -> Trie k a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Trie k a -> m
$cfoldMap' :: forall k m a. Monoid m => (a -> m) -> Trie k a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> Trie k a -> m
$cfoldr :: forall k a b. (a -> b -> b) -> b -> Trie k a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Trie k a -> b
$cfoldr' :: forall k a b. (a -> b -> b) -> b -> Trie k a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Trie k a -> b
$cfoldl :: forall k b a. (b -> a -> b) -> b -> Trie k a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Trie k a -> b
$cfoldl' :: forall k b a. (b -> a -> b) -> b -> Trie k a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Trie k a -> b
$cfoldr1 :: forall k a. (a -> a -> a) -> Trie k a -> a
foldr1 :: forall a. (a -> a -> a) -> Trie k a -> a
$cfoldl1 :: forall k a. (a -> a -> a) -> Trie k a -> a
foldl1 :: forall a. (a -> a -> a) -> Trie k a -> a
$ctoList :: forall k a. Trie k a -> [a]
toList :: forall a. Trie k a -> [a]
$cnull :: forall k a. Trie k a -> Bool
null :: forall a. Trie k a -> Bool
$clength :: forall k a. Trie k a -> Int
length :: forall a. Trie k a -> Int
$celem :: forall k a. Eq a => a -> Trie k a -> Bool
elem :: forall a. Eq a => a -> Trie k a -> Bool
$cmaximum :: forall k a. Ord a => Trie k a -> a
maximum :: forall a. Ord a => Trie k a -> a
$cminimum :: forall k a. Ord a => Trie k a -> a
minimum :: forall a. Ord a => Trie k a -> a
$csum :: forall k a. Num a => Trie k a -> a
sum :: forall a. Num a => Trie k a -> a
$cproduct :: forall k a. Num a => Trie k a -> a
product :: forall a. Num a => Trie k a -> a
Foldable
    )

instance (NFData k, NFData v) => NFData (Trie k v) where
  rnf :: Trie k v -> ()
rnf (Trie Maybe v
a Map k (Trie k v)
b) = Maybe v -> ()
forall a. NFData a => a -> ()
rnf Maybe v
a () -> () -> ()
forall a b. a -> b -> b
`seq` Map k (Trie k v) -> ()
forall a. NFData a => a -> ()
rnf Map k (Trie k v)
b

-- | Empty trie.
instance Null (Trie k v) where
  empty :: Trie k v
empty = Maybe v -> Map k (Trie k v) -> Trie k v
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
Trie Maybe v
forall a. Maybe a
Strict.Nothing Map k (Trie k v)
forall k a. Map k a
Map.empty
  null :: Trie k v -> Bool
null (Trie Maybe v
v Map k (Trie k v)
t) = Maybe v -> Bool
forall a. Null a => a -> Bool
null Maybe v
v Bool -> Bool -> Bool
&& (Trie k v -> Bool) -> Map k (Trie k v) -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Trie k v -> Bool
forall a. Null a => a -> Bool
null Map k (Trie k v)
t
    -- Andreas, 2024-11-16: since we allow non-canoncial tries,
    -- we have to check that every subtrie is empty,
    -- not just that there are no subtries.

-- | Helper function used to implement 'singleton' and 'everyPrefix'.
singletonOrEveryPrefix :: Bool -> [k] -> v -> Trie k v
singletonOrEveryPrefix :: forall k v. Bool -> [k] -> v -> Trie k v
singletonOrEveryPrefix Bool
_           []       !v
v =
  Maybe v -> Map k (Trie k v) -> Trie k v
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
Trie (v -> Maybe v
forall a. a -> Maybe a
Strict.Just v
v) Map k (Trie k v)
forall k a. Map k a
Map.empty
singletonOrEveryPrefix Bool
everyPrefix (k
x : [k]
xs) !v
v =
  Maybe v -> Map k (Trie k v) -> Trie k v
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
Trie (if Bool
everyPrefix then v -> Maybe v
forall a. a -> Maybe a
Strict.Just v
v else Maybe v
forall a. Maybe a
Strict.Nothing)
       (k -> Trie k v -> Map k (Trie k v)
forall k a. k -> a -> Map k a
Map.singleton k
x (Bool -> [k] -> v -> Trie k v
forall k v. Bool -> [k] -> v -> Trie k v
singletonOrEveryPrefix Bool
everyPrefix [k]
xs v
v))

-- | Singleton trie.
singleton :: [k] -> v -> Trie k v
singleton :: forall k v. [k] -> v -> Trie k v
singleton = Bool -> [k] -> v -> Trie k v
forall k v. Bool -> [k] -> v -> Trie k v
singletonOrEveryPrefix Bool
False

-- | @everyPrefix k v@ is a trie where every prefix of @k@ (including
-- @k@ itself) is mapped to @v@.
everyPrefix :: [k] -> v -> Trie k v
everyPrefix :: forall k v. [k] -> v -> Trie k v
everyPrefix = Bool -> [k] -> v -> Trie k v
forall k v. Bool -> [k] -> v -> Trie k v
singletonOrEveryPrefix Bool
True

-- | Left biased union.
--
--   @union = unionWith (\ new old -> new)@.
union :: (Ord k) => Trie k v -> Trie k v -> Trie k v
union :: forall k v. Ord k => Trie k v -> Trie k v -> Trie k v
union = (v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
forall k v.
Ord k =>
(v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
unionWith v -> v -> v
forall a b. a -> b -> a
const

-- | Pointwise union with merge function for values.
unionWith :: (Ord k) => (v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
unionWith :: forall k v.
Ord k =>
(v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
unionWith v -> v -> v
f (Trie Maybe v
v Map k (Trie k v)
ss) (Trie Maybe v
w Map k (Trie k v)
ts) =
  Maybe v -> Map k (Trie k v) -> Trie k v
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
Trie ((v -> v -> v) -> Maybe v -> Maybe v -> Maybe v
forall a. (a -> a -> a) -> Maybe a -> Maybe a -> Maybe a
Strict.unionMaybeWith v -> v -> v
f Maybe v
v Maybe v
w) ((Trie k v -> Trie k v -> Trie k v)
-> Map k (Trie k v) -> Map k (Trie k v) -> Map k (Trie k v)
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith ((v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
forall k v.
Ord k =>
(v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
unionWith v -> v -> v
f) Map k (Trie k v)
ss Map k (Trie k v)
ts)

-- | Insert.  Overwrites existing value if present.
--
--   @insert = insertWith (\ new old -> new)@
insert :: (Ord k) => [k] -> v -> Trie k v -> Trie k v
insert :: forall k v. Ord k => [k] -> v -> Trie k v -> Trie k v
insert [k]
k v
v Trie k v
t = ([k] -> v -> Trie k v
forall k v. [k] -> v -> Trie k v
singleton [k]
k v
v) Trie k v -> Trie k v -> Trie k v
forall k v. Ord k => Trie k v -> Trie k v -> Trie k v
`union` Trie k v
t

-- | Insert with function merging new value with old value.
insertWith :: (Ord k) => (v -> v -> v) -> [k] -> v -> Trie k v -> Trie k v
insertWith :: forall k v.
Ord k =>
(v -> v -> v) -> [k] -> v -> Trie k v -> Trie k v
insertWith v -> v -> v
f [k]
k v
v Trie k v
t = (v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
forall k v.
Ord k =>
(v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
unionWith v -> v -> v
f ([k] -> v -> Trie k v
forall k v. [k] -> v -> Trie k v
singleton [k]
k v
v) Trie k v
t

-- Andreas, 2024-11-16: delete does not clean up empty subtrees,
-- so it is not correct wrt. @Eq@.

-- | Delete value at key, but leave subtree intact.
--
-- Disclaimer: may return a non-canoncial trie
-- because it does not clean up subtries that become empty.
--
delete :: Ord k => [k] -> Trie k v -> Trie k v
delete :: forall k v. Ord k => [k] -> Trie k v -> Trie k v
delete [k]
path = [k] -> (Maybe v -> Maybe v) -> Trie k v -> Trie k v
forall k v.
Ord k =>
[k] -> (Maybe v -> Maybe v) -> Trie k v -> Trie k v
adjust [k]
path (Maybe v -> Maybe v -> Maybe v
forall a b. a -> b -> a
const Maybe v
forall a. Maybe a
Strict.Nothing)

-- Andreas, 2024-11-16: adjust does not clean up empty subtrees,
-- so it is not correct wrt. @Eq@.

-- | Adjust value at key, leave subtree intact.
--
-- Disclaimer: may return a non-canoncial trie
-- because it does not clean up subtries that become empty.
--
adjust ::
  Ord k =>
  [k] -> (Strict.Maybe v -> Strict.Maybe v) -> Trie k v -> Trie k v
adjust :: forall k v.
Ord k =>
[k] -> (Maybe v -> Maybe v) -> Trie k v -> Trie k v
adjust [k]
path Maybe v -> Maybe v
f t :: Trie k v
t@(Trie Maybe v
v Map k (Trie k v)
ts) =
  case [k]
path of
    -- case: found the value we want to adjust: adjust it!
    []                                 -> Maybe v -> Map k (Trie k v) -> Trie k v
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
Trie (Maybe v -> Maybe v
f Maybe v
v) Map k (Trie k v)
ts
    -- case: found the subtrie matching the first key: adjust recursively
    k
k : [k]
ks | Just Trie k v
s <- k -> Map k (Trie k v) -> Maybe (Trie k v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k (Trie k v)
ts -> Maybe v -> Map k (Trie k v) -> Trie k v
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
Trie Maybe v
v (Map k (Trie k v) -> Trie k v) -> Map k (Trie k v) -> Trie k v
forall a b. (a -> b) -> a -> b
$ k -> Trie k v -> Map k (Trie k v) -> Map k (Trie k v)
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k ([k] -> (Maybe v -> Maybe v) -> Trie k v -> Trie k v
forall k v.
Ord k =>
[k] -> (Maybe v -> Maybe v) -> Trie k v -> Trie k v
adjust [k]
ks Maybe v -> Maybe v
f Trie k v
s) Map k (Trie k v)
ts
    -- case: subtrie not found: leave trie untouched
    [k]
_ -> Trie k v
t

-- | Convert to ascending list.
toList :: Ord k => Trie k v -> [([k],v)]
toList :: forall k v. Ord k => Trie k v -> [([k], v)]
toList = Trie k v -> [([k], v)]
forall k v. Ord k => Trie k v -> [([k], v)]
toAscList

-- | Convert to ascending list.
toAscList :: Ord k => Trie k v -> [([k],v)]
toAscList :: forall k v. Ord k => Trie k v -> [([k], v)]
toAscList (Trie Maybe v
mv Map k (Trie k v)
ts) = Maybe ([k], v) -> [([k], v)]
forall a. Maybe a -> [a]
Strict.maybeToList (([],) (v -> ([k], v)) -> Maybe v -> Maybe ([k], v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe v
mv) [([k], v)] -> [([k], v)] -> [([k], v)]
forall a. [a] -> [a] -> [a]
++
  [ (k
kk -> [k] -> [k]
forall a. a -> [a] -> [a]
:[k]
ks, v
v)
  | (k
k,  Trie k v
t) <- Map k (Trie k v) -> [(k, Trie k v)]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map k (Trie k v)
ts
  , ([k]
ks, v
v) <- Trie k v -> [([k], v)]
forall k v. Ord k => Trie k v -> [([k], v)]
toAscList Trie k v
t
  ]

-- | Convert to list where nodes at the same level are ordered according to the
--   given ordering.
toListOrderedBy :: Ord k => (v -> v -> Ordering) -> Trie k v -> [([k], v)]
toListOrderedBy :: forall k v. Ord k => (v -> v -> Ordering) -> Trie k v -> [([k], v)]
toListOrderedBy v -> v -> Ordering
cmp (Trie Maybe v
mv Map k (Trie k v)
ts) =
  Maybe ([k], v) -> [([k], v)]
forall a. Maybe a -> [a]
Strict.maybeToList (([],) (v -> ([k], v)) -> Maybe v -> Maybe ([k], v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe v
mv) [([k], v)] -> [([k], v)] -> [([k], v)]
forall a. [a] -> [a] -> [a]
++
  [ (k
k k -> [k] -> [k]
forall a. a -> [a] -> [a]
: [k]
ks, v
v) | (k
k, Trie k v
t)  <- ((k, Trie k v) -> (k, Trie k v) -> Ordering)
-> [(k, Trie k v)] -> [(k, Trie k v)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
List.sortBy (Maybe v -> Maybe v -> Ordering
cmp' (Maybe v -> Maybe v -> Ordering)
-> ((k, Trie k v) -> Maybe v)
-> (k, Trie k v)
-> (k, Trie k v)
-> Ordering
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` Trie k v -> Maybe v
forall {k} {v}. Trie k v -> Maybe v
val (Trie k v -> Maybe v)
-> ((k, Trie k v) -> Trie k v) -> (k, Trie k v) -> Maybe v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k, Trie k v) -> Trie k v
forall a b. (a, b) -> b
snd) ([(k, Trie k v)] -> [(k, Trie k v)])
-> [(k, Trie k v)] -> [(k, Trie k v)]
forall a b. (a -> b) -> a -> b
$ Map k (Trie k v) -> [(k, Trie k v)]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map k (Trie k v)
ts,
                  ([k]
ks, v
v) <- (v -> v -> Ordering) -> Trie k v -> [([k], v)]
forall k v. Ord k => (v -> v -> Ordering) -> Trie k v -> [([k], v)]
toListOrderedBy v -> v -> Ordering
cmp Trie k v
t ]
  where
    cmp' :: Maybe v -> Maybe v -> Ordering
cmp' Maybe v
Strict.Nothing  Strict.Just{}   = Ordering
LT
    cmp' Strict.Just{}   Maybe v
Strict.Nothing  = Ordering
GT
    cmp' Maybe v
Strict.Nothing  Maybe v
Strict.Nothing  = Ordering
EQ
    cmp' (Strict.Just v
x) (Strict.Just v
y) = v -> v -> Ordering
cmp v
x v
y
    val :: Trie k v -> Maybe v
val (Trie Maybe v
mv Map k (Trie k v)
_) = Maybe v
mv

-- | Create new values based on the entire subtrie. Almost, but not quite
--   comonad extend.
mapSubTries :: Ord k => (Trie k u -> Maybe v) -> Trie k u -> Trie k v
mapSubTries :: forall k u v.
Ord k =>
(Trie k u -> Maybe v) -> Trie k u -> Trie k v
mapSubTries Trie k u -> Maybe v
f t :: Trie k u
t@(Trie Maybe u
mv Map k (Trie k u)
ts) = Maybe v -> Map k (Trie k v) -> Trie k v
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
Trie (Maybe v -> Maybe v
forall lazy strict. Strict lazy strict => lazy -> strict
Strict.toStrict (Trie k u -> Maybe v
f Trie k u
t)) ((Trie k u -> Trie k v) -> Map k (Trie k u) -> Map k (Trie k v)
forall a b. (a -> b) -> Map k a -> Map k b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Trie k u -> Maybe v) -> Trie k u -> Trie k v
forall k u v.
Ord k =>
(Trie k u -> Maybe v) -> Trie k u -> Trie k v
mapSubTries Trie k u -> Maybe v
f) Map k (Trie k u)
ts)

-- | Returns the value associated with the given key, if any.
lookup :: Ord k => [k] -> Trie k v -> Maybe v
lookup :: forall k v. Ord k => [k] -> Trie k v -> Maybe v
lookup []       (Trie Maybe v
v Map k (Trie k v)
_)  = Maybe v -> Maybe v
forall lazy strict. Strict lazy strict => strict -> lazy
Strict.toLazy Maybe v
v
lookup (k
k : [k]
ks) (Trie Maybe v
_ Map k (Trie k v)
ts) = case k -> Map k (Trie k v) -> Maybe (Trie k v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k (Trie k v)
ts of
  Maybe (Trie k v)
Nothing -> Maybe v
forall a. Maybe a
Nothing
  Just Trie k v
t  -> [k] -> Trie k v -> Maybe v
forall k v. Ord k => [k] -> Trie k v -> Maybe v
lookup [k]
ks Trie k v
t

-- | Is the given key present in the trie?
member :: Ord k => [k] -> Trie k v -> Bool
member :: forall k v. Ord k => [k] -> Trie k v -> Bool
member [k]
ks Trie k v
t = Maybe v -> Bool
forall a. Maybe a -> Bool
Lazy.isJust ([k] -> Trie k v -> Maybe v
forall k v. Ord k => [k] -> Trie k v -> Maybe v
lookup [k]
ks Trie k v
t)

-- | Collect all values along a given path.
lookupPath :: Ord k => [k] -> Trie k v -> [v]
lookupPath :: forall k v. Ord k => [k] -> Trie k v -> [v]
lookupPath [k]
xs (Trie Maybe v
v Map k (Trie k v)
cs) = case [k]
xs of
    []     -> Maybe v -> [v]
forall a. Maybe a -> [a]
Strict.maybeToList Maybe v
v
    k
x : [k]
xs -> Maybe v -> [v]
forall a. Maybe a -> [a]
Strict.maybeToList Maybe v
v [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
++
              [v] -> (Trie k v -> [v]) -> Maybe (Trie k v) -> [v]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] ([k] -> Trie k v -> [v]
forall k v. Ord k => [k] -> Trie k v -> [v]
lookupPath [k]
xs) (k -> Map k (Trie k v) -> Maybe (Trie k v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
x Map k (Trie k v)
cs)

-- | Get the subtrie rooted at the given key.
lookupTrie :: Ord k => [k] -> Trie k v -> Trie k v
lookupTrie :: forall k v. Ord k => [k] -> Trie k v -> Trie k v
lookupTrie []       Trie k v
t           = Trie k v
t
lookupTrie (k
k : [k]
ks) (Trie Maybe v
_ Map k (Trie k v)
cs) = Trie k v -> (Trie k v -> Trie k v) -> Maybe (Trie k v) -> Trie k v
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Trie k v
forall a. Null a => a
empty ([k] -> Trie k v -> Trie k v
forall k v. Ord k => [k] -> Trie k v -> Trie k v
lookupTrie [k]
ks) (k -> Map k (Trie k v) -> Maybe (Trie k v)
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k (Trie k v)
cs)

-- | Filter a trie.
filter :: Ord k => (v -> Bool) -> Trie k v -> Trie k v
filter :: forall k v. Ord k => (v -> Bool) -> Trie k v -> Trie k v
filter v -> Bool
p (Trie Maybe v
mv Map k (Trie k v)
ts) = Maybe v -> Map k (Trie k v) -> Trie k v
forall k v. Maybe v -> Map k (Trie k v) -> Trie k v
Trie Maybe v
mv' ((Trie k v -> Bool) -> Map k (Trie k v) -> Map k (Trie k v)
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (Bool -> Bool
not (Bool -> Bool) -> (Trie k v -> Bool) -> Trie k v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie k v -> Bool
forall a. Null a => a -> Bool
null) (Map k (Trie k v) -> Map k (Trie k v))
-> Map k (Trie k v) -> Map k (Trie k v)
forall a b. (a -> b) -> a -> b
$ (v -> Bool) -> Trie k v -> Trie k v
forall k v. Ord k => (v -> Bool) -> Trie k v -> Trie k v
filter v -> Bool
p (Trie k v -> Trie k v) -> Map k (Trie k v) -> Map k (Trie k v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map k (Trie k v)
ts)
  where
    mv' :: Maybe v
mv' =
      case Maybe v
mv of
        Strict.Just v
v | v -> Bool
p v
v -> Maybe v
mv
        Maybe v
_                   -> Maybe v
forall a. Maybe a
Strict.Nothing

-- | Key lens.
valueAt :: Ord k => [k] -> Lens' (Trie k v) (Maybe v)
valueAt :: forall k v. Ord k => [k] -> Lens' (Trie k v) (Maybe v)
valueAt [k]
path Maybe v -> f (Maybe v)
f Trie k v
t = Maybe v -> f (Maybe v)
f ([k] -> Trie k v -> Maybe v
forall k v. Ord k => [k] -> Trie k v -> Maybe v
lookup [k]
path Trie k v
t) f (Maybe v) -> (Maybe v -> Trie k v) -> f (Trie k v)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \ case
  Maybe v
Nothing -> [k] -> Trie k v -> Trie k v
forall k v. Ord k => [k] -> Trie k v -> Trie k v
delete [k]
path Trie k v
t
  Just v
v  -> [k] -> v -> Trie k v -> Trie k v
forall k v. Ord k => [k] -> v -> Trie k v -> Trie k v
insert [k]
path v
v Trie k v
t