-- | Strict tries (based on "Data.Map.Strict" and "Agda.Utils.Maybe.Strict").

module Agda.Utils.Trie
  ( Trie(..)
  , empty, singleton, everyPrefix, insert, insertWith, union, unionWith
  , 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
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
showList :: [Trie k v] -> ShowS
$cshowList :: forall k v. (Show v, Show k) => [Trie k v] -> ShowS
show :: Trie k v -> String
$cshow :: forall k v. (Show v, Show k) => Trie k v -> String
showsPrec :: Int -> Trie k v -> ShowS
$cshowsPrec :: forall k v. (Show v, Show k) => Int -> 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
/= :: 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
Eq
           , (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
<$ :: forall a b. a -> Trie k b -> Trie k a
$c<$ :: forall k a b. a -> Trie k b -> Trie k a
fmap :: forall a b. (a -> b) -> Trie k a -> Trie k b
$cfmap :: forall k a b. (a -> b) -> Trie k a -> Trie k b
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
product :: forall a. Num a => Trie k a -> a
$cproduct :: forall k a. Num a => Trie k a -> a
sum :: forall a. Num a => Trie k a -> a
$csum :: forall k a. Num a => Trie k a -> a
minimum :: forall a. Ord a => Trie k a -> a
$cminimum :: forall k a. Ord a => Trie k a -> a
maximum :: forall a. Ord a => Trie k a -> a
$cmaximum :: forall k a. Ord a => Trie k a -> a
elem :: forall a. Eq a => a -> Trie k a -> Bool
$celem :: forall k a. Eq a => a -> Trie k a -> Bool
length :: forall a. Trie k a -> Int
$clength :: forall k a. Trie k a -> Int
null :: forall a. Trie k a -> Bool
$cnull :: forall k a. Trie k a -> Bool
toList :: forall a. Trie k a -> [a]
$ctoList :: forall k a. Trie k a -> [a]
foldl1 :: forall a. (a -> a -> a) -> Trie k a -> a
$cfoldl1 :: forall k a. (a -> a -> a) -> Trie k a -> a
foldr1 :: forall a. (a -> a -> a) -> Trie k a -> a
$cfoldr1 :: forall k a. (a -> a -> a) -> Trie k a -> a
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
$cfoldl :: forall k b a. (b -> a -> 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
$cfoldr :: forall k a b. (a -> b -> b) -> b -> Trie k a -> b
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
$cfoldMap :: forall k m a. Monoid m => (a -> m) -> Trie k a -> m
fold :: forall m. Monoid m => Trie k m -> m
$cfold :: forall k m. Monoid m => Trie k m -> m
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 () -> () -> ()
`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
&& Map k (Trie k v) -> Bool
forall a. Null a => a -> Bool
null Map k (Trie k v)
t

-- | 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

-- | Delete value at key, but leave subtree intact.
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)

-- | Adjust value at key, leave subtree intact.
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 a. Maybe a -> Maybe a
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 (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 a. Maybe a -> Maybe a
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' (Maybe v) (Trie k v)
valueAt :: forall k v. Ord k => [k] -> Lens' (Maybe v) (Trie k 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 (m :: * -> *) a b. Functor m => m a -> (a -> b) -> m 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