{-# LANGUAGE BangPatterns #-}

-- | 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 Data.Function
import Data.Foldable (Foldable)
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
           , a -> Trie k b -> Trie k a
(a -> b) -> Trie k a -> Trie k b
(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
<$ :: a -> Trie k b -> Trie k a
$c<$ :: forall k a b. a -> Trie k b -> Trie k a
fmap :: (a -> b) -> Trie k a -> Trie k b
$cfmap :: forall k a b. (a -> b) -> Trie k a -> Trie k b
Functor
           , Trie k a -> Bool
(a -> m) -> Trie k a -> m
(a -> b -> b) -> b -> Trie k a -> b
(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 :: Trie k a -> a
$cproduct :: forall k a. Num a => Trie k a -> a
sum :: Trie k a -> a
$csum :: forall k a. Num a => Trie k a -> a
minimum :: Trie k a -> a
$cminimum :: forall k a. Ord a => Trie k a -> a
maximum :: Trie k a -> a
$cmaximum :: forall k a. Ord a => Trie k a -> a
elem :: a -> Trie k a -> Bool
$celem :: forall k a. Eq a => a -> Trie k a -> Bool
length :: Trie k a -> Int
$clength :: forall k a. Trie k a -> Int
null :: Trie k a -> Bool
$cnull :: forall k a. Trie k a -> Bool
toList :: Trie k a -> [a]
$ctoList :: forall k a. Trie k a -> [a]
foldl1 :: (a -> a -> a) -> Trie k a -> a
$cfoldl1 :: forall k a. (a -> a -> a) -> Trie k a -> a
foldr1 :: (a -> a -> a) -> Trie k a -> a
$cfoldr1 :: forall k a. (a -> a -> a) -> Trie k a -> a
foldl' :: (b -> a -> b) -> b -> Trie k a -> b
$cfoldl' :: forall k b a. (b -> a -> b) -> b -> Trie k a -> b
foldl :: (b -> a -> b) -> b -> Trie k a -> b
$cfoldl :: forall k b a. (b -> a -> b) -> b -> Trie k a -> b
foldr' :: (a -> b -> b) -> b -> Trie k a -> b
$cfoldr' :: forall k a b. (a -> b -> b) -> b -> Trie k a -> b
foldr :: (a -> b -> b) -> b -> Trie k a -> b
$cfoldr :: forall k a b. (a -> b -> b) -> b -> Trie k a -> b
foldMap' :: (a -> m) -> Trie k a -> m
$cfoldMap' :: forall k m a. Monoid m => (a -> m) -> Trie k a -> m
foldMap :: (a -> m) -> Trie k a -> m
$cfoldMap :: forall k m a. Monoid m => (a -> m) -> Trie k a -> m
fold :: Trie k m -> m
$cfold :: forall k m. Monoid m => Trie k m -> m
Foldable
           )

-- | 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 v :: Maybe v
v t :: 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 :: Bool -> [k] -> v -> Trie k v
singletonOrEveryPrefix _           []       !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 everyPrefix :: Bool
everyPrefix (x :: k
x : xs :: [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 :: [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 :: [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 :: 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 :: (v -> v -> v) -> Trie k v -> Trie k v -> Trie k v
unionWith f :: v -> v -> v
f (Trie v :: Maybe v
v ss :: Map k (Trie k v)
ss) (Trie w :: Maybe v
w ts :: 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 :: [k] -> v -> Trie k v -> Trie k v
insert k :: [k]
k v :: v
v t :: Trie k v
t = Trie k v -> Trie k v -> Trie k v
forall k v. Ord k => Trie k v -> Trie k v -> Trie k v
union ([k] -> v -> Trie k v
forall k v. [k] -> v -> Trie k v
singleton [k]
k v
v) 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 :: (v -> v -> v) -> [k] -> v -> Trie k v -> Trie k v
insertWith f :: v -> v -> v
f k :: [k]
k v :: v
v t :: 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 :: [k] -> Trie k v -> Trie k v
delete path :: [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 :: [k] -> (Maybe v -> Maybe v) -> Trie k v -> Trie k v
adjust path :: [k]
path f :: Maybe v -> Maybe v
f t :: Trie k v
t@(Trie v :: Maybe v
v ts :: 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 :: [k]
ks | Just s :: 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
    _ -> Trie k v
t

-- | Convert to ascending list.
toList :: Ord k => Trie k v -> [([k],v)]
toList :: 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 :: Trie k v -> [([k], v)]
toAscList (Trie mv :: Maybe v
mv ts :: 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
k,  t :: 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
  , (ks :: [k]
ks, v :: 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 :: (v -> v -> Ordering) -> Trie k v -> [([k], v)]
toListOrderedBy cmp :: v -> v -> Ordering
cmp (Trie mv :: Maybe v
mv ts :: 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
k, t :: 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,
                  (ks :: [k]
ks, v :: 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' Strict.Nothing  Strict.Just{}   = Ordering
LT
    cmp' Strict.Just{}   Strict.Nothing  = Ordering
GT
    cmp' Strict.Nothing  Strict.Nothing  = Ordering
EQ
    cmp' (Strict.Just x :: v
x) (Strict.Just y :: v
y) = v -> v -> Ordering
cmp v
x v
y
    val :: Trie k v -> Maybe v
val (Trie mv :: Maybe v
mv _) = 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 :: (Trie k u -> Maybe v) -> Trie k u -> Trie k v
mapSubTries f :: Trie k u -> Maybe v
f t :: Trie k u
t@(Trie mv :: Maybe u
mv ts :: 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 :: [k] -> Trie k v -> Maybe v
lookup []       (Trie v :: Maybe v
v _)  = Maybe v -> Maybe v
forall a. Maybe a -> Maybe a
Strict.toLazy Maybe v
v
lookup (k :: k
k : ks :: [k]
ks) (Trie _ ts :: 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
  Nothing -> Maybe v
forall a. Maybe a
Nothing
  Just t :: 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 :: [k] -> Trie k v -> Bool
member ks :: [k]
ks t :: 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 :: [k] -> Trie k v -> [v]
lookupPath xs :: [k]
xs (Trie v :: Maybe v
v cs :: Map k (Trie k v)
cs) = case [k]
xs of
    []     -> Maybe v -> [v]
forall a. Maybe a -> [a]
Strict.maybeToList Maybe v
v
    x :: k
x : xs :: [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 :: [k] -> Trie k v -> Trie k v
lookupTrie []       t :: Trie k v
t           = Trie k v
t
lookupTrie (k :: k
k : ks :: [k]
ks) (Trie _ cs :: 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 :: (v -> Bool) -> Trie k v -> Trie k v
filter p :: v -> Bool
p (Trie mv :: Maybe v
mv ts :: 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 | v -> Bool
p v
v -> Maybe v
mv
        _                   -> Maybe v
forall a. Maybe a
Strict.Nothing

-- | Key lens.
valueAt :: Ord k => [k] -> Lens' (Maybe v) (Trie k v)
valueAt :: [k] -> Lens' (Maybe v) (Trie k v)
valueAt path :: [k]
path f :: Maybe v -> f (Maybe v)
f t :: 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
  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
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