{-# OPTIONS_GHC -Wunused-imports #-}
{-# OPTIONS_GHC -Wunused-matches #-}
{-# OPTIONS_GHC -Wunused-binds #-}

-- | Strict compressed tries
--   (based on "Data.Map.Strict" and "Agda.Utils.Maybe.Strict").
--
-- In contrast to "Agda.Utils.Trie", branches from a node
-- can be labeled with non-empty lists of keys that do not
-- overlap, meaning they start with different first elements.
-- This way, non-branching unlabelled pathes can be stored
-- more efficiently.
--
-- As a side effect, the notion of height changes
-- which we utilize in the 'truncate' function.
--
-- An 'Agda.Utils.CompressedTrie.Trie' is usually constructed
-- from an 'Agda.Utils.Trie.Trie' via the 'compress' function.

module Agda.Utils.CompressedTrie where

import Prelude hiding (null, truncate)

import Control.DeepSeq

import Data.Map.Strict         (Map)
import Data.Map.Strict         qualified as Map
import Data.Monoid             (Sum(Sum,getSum))

import Agda.Utils.List1        (List1, pattern (:|))
import Agda.Utils.List1        qualified as List1
import Agda.Utils.Map          (isSingleMap)
import Agda.Utils.Maybe.Strict qualified as Strict
import Agda.Utils.Null
import Agda.Utils.Singleton
import Agda.Utils.Trie         qualified as Uncompressed

import Agda.Utils.Impossible (__IMPOSSIBLE__)

-- | Finite map from @[k]@ to @v@.
--
--   With the strict 'Maybe' type, 'Trie' is also strict in 'v'.
--
--   Invariant: only the root of the tree
--   may be of the form @Trie Nothing (Map.singleton ks t)@.
--   A node inside of the trie of that form must be fused
--   into its branch.
data Trie k v = Trie
  { forall k v. Trie k v -> Maybe v
trieValue    :: !(Strict.Maybe v)
  , forall k v. Trie k v -> Map (List1 k) (Trie k v)
trieBranches :: !(Map (List1 k) (Trie k v))
      -- Invariant: The keys are prefix-free, meaning their first letters differ.
  }
  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 (List1 k) (Trie k v)
b) = Maybe v -> ()
forall a. NFData a => a -> ()
rnf Maybe v
a () -> () -> ()
forall a b. a -> b -> b
`seq` Map (List1 k) (Trie k v) -> ()
forall a. NFData a => a -> ()
rnf Map (List1 k) (Trie k v)
b

---------------------------------------------------------------------------
-- * Construction

-- | Empty trie.
instance Null (Trie k v) where
  empty :: Trie k v
empty = Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
forall k v. Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
Trie Maybe v
forall a. Maybe a
Strict.Nothing Map (List1 k) (Trie k v)
forall k a. Map k a
Map.empty
  null :: Trie k v -> Bool
null (Trie Maybe v
v Map (List1 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 (List1 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 (List1 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.

-- | Singleton tree with value at the root.
root :: v -> Trie k v
root :: forall v k. v -> Trie k v
root v
v = Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
forall k v. Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
Trie (v -> Maybe v
forall a. a -> Maybe a
Strict.Just v
v) Map (List1 k) (Trie k v)
forall k a. Map k a
Map.empty

instance Singleton ([k],v) (Trie k v) where
  singleton :: ([k], v) -> Trie k v
singleton = ([k] -> v -> Trie k v) -> ([k], v) -> Trie k v
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [k] -> v -> Trie k v
forall k v. [k] -> v -> Trie k v
Agda.Utils.CompressedTrie.singleton

-- | Singleton trie.
singleton :: [k] -> v -> Trie k v
singleton :: forall k v. [k] -> v -> Trie k v
singleton []       = v -> Trie k v
forall v k. v -> Trie k v
root
singleton (k
k : [k]
ks) = List1 k -> v -> Trie k v
forall k v. List1 k -> v -> Trie k v
singleton1 (k
k k -> [k] -> List1 k
forall a. a -> [a] -> NonEmpty a
:| [k]
ks)

-- | Singleton tree with empty root.
singleton1 :: List1 k -> v -> Trie k v
singleton1 :: forall k v. List1 k -> v -> Trie k v
singleton1 List1 k
ks = Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
forall k v. Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
Trie Maybe v
forall a. Maybe a
Strict.Nothing (Map (List1 k) (Trie k v) -> Trie k v)
-> (v -> Map (List1 k) (Trie k v)) -> v -> Trie k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. List1 k -> Trie k v -> Map (List1 k) (Trie k v)
forall k a. k -> a -> Map k a
Map.singleton List1 k
ks (Trie k v -> Map (List1 k) (Trie k v))
-> (v -> Trie k v) -> v -> Map (List1 k) (Trie k v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> Trie k v
forall v k. v -> Trie k v
root

-- | A single branch with the given subtree.
--   Precondition: the given subtree has a value at the root.
singleBranch :: List1 k -> Trie k v -> (List1 k, Trie k v)
singleBranch :: forall k v. List1 k -> Trie k v -> (List1 k, Trie k v)
singleBranch List1 k
ks Trie k v
t
  | Just (List1 k
ks', Trie k v
t') <- Trie k v -> Maybe (List1 k, Trie k v)
forall k v. Trie k v -> Maybe (List1 k, Trie k v)
isSingleBranch Trie k v
t = (List1 k
ks List1 k -> List1 k -> List1 k
forall a. Semigroup a => a -> a -> a
<> List1 k
ks', Trie k v
t')
  | Bool
otherwise = (List1 k
ks, Trie k v
t)

-- | Is the trie starting with a single branch?
isSingleBranch :: Trie k v -> Maybe (List1 k, Trie k v)
isSingleBranch :: forall k v. Trie k v -> Maybe (List1 k, Trie k v)
isSingleBranch = \case
  Trie Maybe v
Strict.Nothing Map (List1 k) (Trie k v)
ts | Just (List1 k
k, Trie k v
t) <- Map (List1 k) (Trie k v) -> Maybe (List1 k, Trie k v)
forall k v. Map k v -> Maybe (k, v)
isSingleMap Map (List1 k) (Trie k v)
ts
    -> (List1 k, Trie k v) -> Maybe (List1 k, Trie k v)
forall a. a -> Maybe a
Just (List1 k
k, Trie k v
t)
  Trie k v
_ -> Maybe (List1 k, Trie k v)
forall a. Maybe a
Nothing

-- | Turn a trie into a compressed trie.
compress :: Ord k => Uncompressed.Trie k v -> Trie k v
compress :: forall k v. Ord k => Trie k v -> Trie k v
compress (Uncompressed.Trie Maybe v
mv Map k (Trie k v)
ts) =
    Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
forall k v. Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
Trie Maybe v
mv (Map (List1 k) (Trie k v) -> Trie k v)
-> Map (List1 k) (Trie k v) -> Trie k v
forall a b. (a -> b) -> a -> b
$ [(List1 k, Trie k v)] -> Map (List1 k) (Trie k v)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(List1 k, Trie k v)] -> Map (List1 k) (Trie k v))
-> [(List1 k, Trie k v)] -> Map (List1 k) (Trie k v)
forall a b. (a -> b) -> a -> b
$ ((k, Trie k v) -> (List1 k, Trie k v))
-> [(k, Trie k v)] -> [(List1 k, Trie k v)]
forall a b. (a -> b) -> [a] -> [b]
map (k, Trie k v) -> (List1 k, Trie k v)
forall {k} {v}. Ord k => (k, Trie k v) -> (List1 k, Trie k v)
go ([(k, Trie k v)] -> [(List1 k, Trie k v)])
-> [(k, Trie k v)] -> [(List1 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.toList Map k (Trie k v)
ts
  where
    go :: (k, Trie k v) -> (List1 k, Trie k v)
go (k
k, Trie k v
t) = List1 k -> Trie k v -> (List1 k, Trie k v)
forall k v. List1 k -> Trie k v -> (List1 k, Trie k v)
singleBranch (k -> List1 k
forall a. a -> NonEmpty a
List1.singleton k
k) (Trie k v -> Trie k v
forall k v. Ord k => Trie k v -> Trie k v
compress Trie k v
t)

---------------------------------------------------------------------------
-- * Deconstruction

-- | Turn a compressed trie into an uncompressed trie.
uncompress :: forall k v. Ord k => Trie k v -> Uncompressed.Trie k v
uncompress :: forall k v. Ord k => Trie k v -> Trie k v
uncompress (Trie Maybe v
mv Map (List1 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
Uncompressed.Trie Maybe v
mv (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)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(k, Trie k v)] -> Map k (Trie k v))
-> [(k, Trie k v)] -> Map k (Trie k v)
forall a b. (a -> b) -> a -> b
$ ((List1 k, Trie k v) -> (k, Trie k v))
-> [(List1 k, Trie k v)] -> [(k, Trie k v)]
forall a b. (a -> b) -> [a] -> [b]
map (List1 k, Trie k v) -> (k, Trie k v)
go ([(List1 k, Trie k v)] -> [(k, Trie k v)])
-> [(List1 k, Trie k v)] -> [(k, Trie k v)]
forall a b. (a -> b) -> a -> b
$ Map (List1 k) (Trie k v) -> [(List1 k, Trie k v)]
forall k a. Map k a -> [(k, a)]
Map.toList Map (List1 k) (Trie k v)
ts
  where
    go :: (List1 k, Trie k v) -> (k, Uncompressed.Trie k v)
    go :: (List1 k, Trie k v) -> (k, Trie k v)
go (k
k :| [k]
ks, Trie k v
t) = (k
k, [k] -> Trie k v -> Trie k v
forall k v. [k] -> Trie k v -> Trie k v
Uncompressed.prefixBy [k]
ks (Trie k v -> Trie k v) -> Trie k v -> Trie k v
forall a b. (a -> b) -> a -> b
$ Trie k v -> Trie k v
forall k v. Ord k => Trie k v -> Trie k v
uncompress 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 (List1 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]
++
  [ (List1 k -> [Item (List1 k)]
forall l. IsList l => l -> [Item l]
List1.toList List1 k
ks1 [k] -> [k] -> [k]
forall a. [a] -> [a] -> [a]
++ [k]
ks, v
v)
  | (List1 k
ks1, Trie k v
t) <- Map (List1 k) (Trie k v) -> [(List1 k, Trie k v)]
forall k a. Map k a -> [(k, a)]
Map.toAscList Map (List1 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
  ]

-- | The number of labeled nodes.
size :: Trie k v -> Int
size :: forall k a. Trie k a -> Int
size (Trie Maybe v
mv Map (List1 k) (Trie k v)
ts) = Int -> (v -> Int) -> Maybe v -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
Strict.maybe Int
0 (Int -> v -> Int
forall a b. a -> b -> a
const Int
1) Maybe v
mv Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sum Int -> Int
forall a. Sum a -> a
getSum ((Trie k v -> Sum Int) -> Map (List1 k) (Trie k v) -> Sum Int
forall m a. Monoid m => (a -> m) -> Map (List1 k) a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Int -> Sum Int
forall a. a -> Sum a
Sum (Int -> Sum Int) -> (Trie k v -> Int) -> Trie k v -> Sum Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie k v -> Int
forall k a. Trie k a -> Int
size) Map (List1 k) (Trie k v)
ts)

-- | The number of nodes that are labeled or leaves..
sizeWithLeaves :: Trie k v -> Int
sizeWithLeaves :: forall k a. Trie k a -> Int
sizeWithLeaves (Trie Maybe v
mv Map (List1 k) (Trie k v)
ts)
  | Map (List1 k) (Trie k v) -> Bool
forall a. Null a => a -> Bool
null Map (List1 k) (Trie k v)
ts   = Int
1
  | Bool
otherwise = Int -> (v -> Int) -> Maybe v -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
Strict.maybe Int
0 (Int -> v -> Int
forall a b. a -> b -> a
const Int
1) Maybe v
mv Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sum Int -> Int
forall a. Sum a -> a
getSum ((Trie k v -> Sum Int) -> Map (List1 k) (Trie k v) -> Sum Int
forall m a. Monoid m => (a -> m) -> Map (List1 k) a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (Int -> Sum Int
forall a. a -> Sum a
Sum (Int -> Sum Int) -> (Trie k v -> Int) -> Trie k v -> Sum Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie k v -> Int
forall k a. Trie k a -> Int
sizeWithLeaves) Map (List1 k) (Trie k v)
ts)

---------------------------------------------------------------------------
-- * Modification

-- | Modify the leaves (input and output may be non-canonical).

substLeaves :: (Maybe v -> Trie k v) -> Trie k v -> Trie k v
substLeaves :: forall v k. (Maybe v -> Trie k v) -> Trie k v -> Trie k v
substLeaves Maybe v -> Trie k v
f (Trie Maybe v
mv Map (List1 k) (Trie k v)
ts)
  | Map (List1 k) (Trie k v) -> Bool
forall a. Null a => a -> Bool
null Map (List1 k) (Trie k v)
ts   = Maybe v -> Trie k v
f (Maybe v -> Trie k v) -> Maybe v -> Trie k v
forall a b. (a -> b) -> a -> b
$ Maybe v -> Maybe v
forall lazy strict. Strict lazy strict => strict -> lazy
Strict.toLazy Maybe v
mv
  | Bool
otherwise = Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
forall k v. Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
Trie Maybe v
mv ((Trie k v -> Trie k v)
-> Map (List1 k) (Trie k v) -> Map (List1 k) (Trie k v)
forall a b. (a -> b) -> Map (List1 k) a -> Map (List1 k) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Maybe v -> Trie k v) -> Trie k v -> Trie k v
forall v k. (Maybe v -> Trie k v) -> Trie k v -> Trie k v
substLeaves Maybe v -> Trie k v
f) Map (List1 k) (Trie k v)
ts)

-- | Truncate tree at given depth.
--   @truncate 0 t@ returns the root.

truncate :: Int -> Trie k v -> Trie k v
truncate :: forall k v. Int -> Trie k v -> Trie k v
truncate Int
n t :: Trie k v
t@(Trie Maybe v
mv Map (List1 k) (Trie k v)
ts)
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0    = if Map (List1 k) (Trie k v) -> Bool
forall a. Null a => a -> Bool
null Map (List1 k) (Trie k v)
ts then Trie k v
t else Trie k v
forall a. Null a => a
empty
  | Bool
otherwise = Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
forall k v. Maybe v -> Map (List1 k) (Trie k v) -> Trie k v
Trie Maybe v
mv (Map (List1 k) (Trie k v) -> Trie k v)
-> Map (List1 k) (Trie k v) -> Trie k v
forall a b. (a -> b) -> a -> b
$ (Trie k v -> Trie k v)
-> Map (List1 k) (Trie k v) -> Map (List1 k) (Trie k v)
forall a b. (a -> b) -> Map (List1 k) a -> Map (List1 k) b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> Trie k v -> Trie k v
forall k v. Int -> Trie k v -> Trie k v
truncate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Map (List1 k) (Trie k v)
ts

-- | Truncate the tree so that it has at most the given size
--   according to the given sizing function.
--   @truncateSize sizing 0 t@ returns the trie @truncate 0 t@.

truncateSize :: forall k v. (Trie k v -> Int) -> Int -> Trie k v -> Trie k v
truncateSize :: forall k v. (Trie k v -> Int) -> Int -> Trie k v -> Trie k v
truncateSize Trie k v -> Int
size Int
n Trie k v
t =
  -- Drop while still increasing but staying below n.
  case
    (((Int, Trie k v), (Int, Trie k v)) -> Bool)
-> [((Int, Trie k v), (Int, Trie k v))]
-> [((Int, Trie k v), (Int, Trie k v))]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile
      (\ ((Int
n1,Trie k v
_),(Int
n2,Trie k v
_)) -> Int
n1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n2 Bool -> Bool -> Bool
&& Int
n2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n)
      ([(Int, Trie k v)]
-> [(Int, Trie k v)] -> [((Int, Trie k v), (Int, Trie k v))]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Int, Trie k v)]
ts (Int -> [(Int, Trie k v)] -> [(Int, Trie k v)]
forall a. Int -> [a] -> [a]
drop Int
1 [(Int, Trie k v)]
ts))
  of
  -- Return the last trie that was still increasing and stayed below n.
    ((Int
_, Trie k v
t1), (Int, Trie k v)
_) : [((Int, Trie k v), (Int, Trie k v))]
_ -> Trie k v
t1
    [] -> Trie k v
forall a. HasCallStack => a
__IMPOSSIBLE__
  where
    ts :: [(Int, Trie k v)]
    ts :: [(Int, Trie k v)]
ts = [ (Trie k v -> Int
size Trie k v
t', Trie k v
t') | Int
d <- [Int
0..], let t' :: Trie k v
t' = Int -> Trie k v -> Trie k v
forall k v. Int -> Trie k v -> Trie k v
truncate Int
d Trie k v
t ]

-- | @truncateSize sizeWithLeaves@
truncateSizeWithLeaves :: Int -> Trie k v -> Trie k v
truncateSizeWithLeaves :: forall k v. Int -> Trie k v -> Trie k v
truncateSizeWithLeaves = (Trie k v -> Int) -> Int -> Trie k v -> Trie k v
forall k v. (Trie k v -> Int) -> Int -> Trie k v -> Trie k v
truncateSize Trie k v -> Int
forall k a. Trie k a -> Int
sizeWithLeaves