{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}

{- | Tries with 'Bytes' (equiv. 'ByteArray') as keys.
This implementation is optimized for performing queries rather
than updating the structure repeatedly.
-}
module Data.Trie.Word8
  ( -- * Trie Type
    Trie
  , valid

    -- * Query

    -- ** Lookup
  , lookup
  , lookupTrie
  , lookupPrefixes

    -- ** Search
  , multiFindReplace
  , search
  , replace
  , stripPrefix
  , stripPrefixWithKey

    -- ** Size
  , null
  , size

    -- * Construction
  , empty
  , singleton

    -- ** Conversion
  , fromList
  , toList
  , foldl'
  , traverse_

    -- ** Insertion
  , insert
  , insertWith

    -- ** Deletion
  , delete

    -- ** Combine
  , union
  , unionWith
  , append
  , prepend
  ) where

import Prelude hiding (lookup, null)

import Control.Applicative ((<|>))
import Data.Bifunctor (first)
import Data.Bytes (Bytes, fromByteArray, toByteArray)
import Data.Bytes.Chunks (Chunks)
import Data.Map.Word8 (Map)
import Data.Maybe (isNothing)
import Data.Monoid (Any (Any), getAny)
import Data.Primitive.ByteArray (ByteArray, indexByteArray)
import Data.Word (Word8)

import qualified Data.Bytes as Bytes
import qualified Data.Bytes.Builder as Build
import qualified Data.Foldable as Foldable
import qualified Data.Map.Word8 as Map
import qualified Data.Maybe.Unpacked as U

{- | Tries implemented using a 256-entry bitmap as given in
"Data.Map.Word8".
This means that each branch point can be navigated with only
some bit manipulations and adding an offset.
On sparse data, this should save a lot of space relative to holding
a 256-entry pointer array.

This data type has 'Tip', 'Run', and 'Branch' nodes.
Branches always have at least two children,
and Runs always have at least one byte.
Leaves are 'Tip's.
Once the invariants are met (see below),
there is exactly one 'Trie' representation for each trie.

In each constructor, the @U.Maybe a@ is a possible entry;
it comes before any child bytes.

INVARIANT: The Run constructor never has a linear child.
           Linear nodes are those with no value and exactly one child,
           which in this implementation is only valueless runs.
INVARIANT: The Run constructor never has zero bytes.
INVARIANT: The Branch constructor has at least two children.
INVARIANT: No child of a node has size zero. That includes:
             The next node after a run is never null.
             No child of a branch is ever null.
-}
data Trie a
  = Tip {-# UNPACK #-} !(U.Maybe a)
  | -- ByteArray uses more copying on modification,
    -- but the data structures are smaller than with Bytes, making lookup faster
    UnsafeRun {-# UNPACK #-} !(U.Maybe a) {-# UNPACK #-} !ByteArray !(Trie a)
  | UnsafeBranch {-# UNPACK #-} !(U.Maybe a) !(Map (Trie a))
  deriving (Trie a -> Trie a -> Bool
(Trie a -> Trie a -> Bool)
-> (Trie a -> Trie a -> Bool) -> Eq (Trie a)
forall a. Eq a => Trie a -> Trie a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Trie a -> Trie a -> Bool
== :: Trie a -> Trie a -> Bool
$c/= :: forall a. Eq a => Trie a -> Trie a -> Bool
/= :: Trie a -> Trie a -> Bool
Eq, (forall a b. (a -> b) -> Trie a -> Trie b)
-> (forall a b. a -> Trie b -> Trie a) -> Functor Trie
forall a b. a -> Trie b -> Trie a
forall a b. (a -> b) -> Trie a -> Trie b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Trie a -> Trie b
fmap :: forall a b. (a -> b) -> Trie a -> Trie b
$c<$ :: forall a b. a -> Trie b -> Trie a
<$ :: forall a b. a -> Trie b -> Trie a
Functor)

instance (Semigroup a) => Semigroup (Trie a) where <> :: Trie a -> Trie a -> Trie a
(<>) = Trie a -> Trie a -> Trie a
forall a. Semigroup a => Trie a -> Trie a -> Trie a
append
instance (Semigroup a) => Monoid (Trie a) where mempty :: Trie a
mempty = Trie a
forall a. Trie a
empty
instance (Show a) => Show (Trie a) where show :: Trie a -> String
show = [(Bytes, a)] -> String
forall a. Show a => a -> String
show ([(Bytes, a)] -> String)
-> (Trie a -> [(Bytes, a)]) -> Trie a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Trie a -> [(Bytes, a)]
forall a. Trie a -> [(Bytes, a)]
toList

{-# COMPLETE Tip, Run, Branch #-}

pattern Run :: U.Maybe a -> ByteArray -> Trie a -> Trie a
pattern $mRun :: forall {r} {a}.
Trie a
-> (Maybe a -> ByteArray -> Trie a -> r) -> ((# #) -> r) -> r
$bRun :: forall a. Maybe a -> ByteArray -> Trie a -> Trie a
Run v run next <- UnsafeRun v run next
  where
    Run Maybe a
v ByteArray
run Trie a
next
      | Trie a -> Bool
forall a. Trie a -> Bool
null Trie a
next = Maybe a -> Trie a
forall a. Maybe a -> Trie a
Tip Maybe a
v
      | Bytes -> Bool
Bytes.null (ByteArray -> Bytes
fromByteArray ByteArray
run) = Trie a
next -- WARNING: throws away `v`, value/non-value from `next`
      | Just (ByteArray
run', Trie a
next') <- Trie a -> Maybe (ByteArray, Trie a)
forall a. Trie a -> Maybe (ByteArray, Trie a)
fromLinear Trie a
next =
          Maybe a -> ByteArray -> Trie a -> Trie a
forall a. Maybe a -> ByteArray -> Trie a -> Trie a
UnsafeRun Maybe a
v (ByteArray
run ByteArray -> ByteArray -> ByteArray
forall a. Semigroup a => a -> a -> a
<> ByteArray
run') Trie a
next'
      | Bool
otherwise = Maybe a -> ByteArray -> Trie a -> Trie a
forall a. Maybe a -> ByteArray -> Trie a -> Trie a
UnsafeRun Maybe a
v ByteArray
run Trie a
next

pattern Branch :: U.Maybe a -> Map (Trie a) -> Trie a
pattern $mBranch :: forall {r} {a}.
Trie a -> (Maybe a -> Map (Trie a) -> r) -> ((# #) -> r) -> r
$bBranch :: forall a. Maybe a -> Map (Trie a) -> Trie a
Branch v children <- UnsafeBranch v children
  where
    Branch Maybe a
v (Map (Trie a) -> Map (Trie a)
forall a. Map (Trie a) -> Map (Trie a)
removeEmptyChildren -> Map (Trie a)
children)
      | Map (Trie a) -> Bool
forall a. Map a -> Bool
Map.null Map (Trie a)
children = Maybe a -> Trie a
forall a. Maybe a -> Trie a
Tip Maybe a
v
      | Just (Word8
c, Trie a
child) <- Map (Trie a) -> Maybe (Word8, Trie a)
forall a. Map a -> Maybe (Word8, a)
fromSingletonMap Map (Trie a)
children =
          Maybe a -> ByteArray -> Trie a -> Trie a
forall a. Maybe a -> ByteArray -> Trie a -> Trie a
Run Maybe a
v (Bytes -> ByteArray
toByteArray (Bytes -> ByteArray) -> Bytes -> ByteArray
forall a b. (a -> b) -> a -> b
$ Word8 -> Bytes
Bytes.singleton Word8
c) Trie a
child
      | Bool
otherwise = Maybe a -> Map (Trie a) -> Trie a
forall a. Maybe a -> Map (Trie a) -> Trie a
UnsafeBranch Maybe a
v Map (Trie a)
children
removeEmptyChildren :: Map (Trie a) -> Map (Trie a)
removeEmptyChildren :: forall a. Map (Trie a) -> Map (Trie a)
removeEmptyChildren = (Word8 -> Trie a -> Map (Trie a) -> Map (Trie a))
-> Map (Trie a) -> Map (Trie a) -> Map (Trie a)
forall a b. (Word8 -> a -> b -> b) -> b -> Map a -> b
Map.foldrWithKeys Word8 -> Trie a -> Map (Trie a) -> Map (Trie a)
forall {a}. Word8 -> Trie a -> Map (Trie a) -> Map (Trie a)
f Map (Trie a)
forall a. Map a
Map.empty
 where
  f :: Word8 -> Trie a -> Map (Trie a) -> Map (Trie a)
f Word8
k Trie a
v Map (Trie a)
xs = if Trie a -> Bool
forall a. Trie a -> Bool
null Trie a
v then Map (Trie a)
xs else Word8 -> Trie a -> Map (Trie a) -> Map (Trie a)
forall a. Word8 -> a -> Map a -> Map a
Map.insert Word8
k Trie a
v Map (Trie a)
xs

-- Get nodes with no value, and exactly one possible next byte.
-- I.e. it never returns an empty bytes in the tuple.
fromLinear :: Trie a -> Maybe (ByteArray, Trie a)
fromLinear :: forall a. Trie a -> Maybe (ByteArray, Trie a)
fromLinear (Run Maybe a
U.Nothing ByteArray
run Trie a
next) = (ByteArray, Trie a) -> Maybe (ByteArray, Trie a)
forall a. a -> Maybe a
Just (ByteArray
run, Trie a
next)
fromLinear Trie a
_ = Maybe (ByteArray, Trie a)
forall a. Maybe a
Nothing

valid :: Trie a -> Bool
valid :: forall a. Trie a -> Bool
valid (Tip Maybe a
_) = Bool
True
valid (Run Maybe a
_ ByteArray
run Trie a
next) =
  Bool -> Bool
not (Bytes -> Bool
Bytes.null (ByteArray -> Bytes
fromByteArray ByteArray
run))
    Bool -> Bool -> Bool
&& Maybe (ByteArray, Trie a) -> Bool
forall a. Maybe a -> Bool
isNothing (Trie a -> Maybe (ByteArray, Trie a)
forall a. Trie a -> Maybe (ByteArray, Trie a)
fromLinear Trie a
next)
    Bool -> Bool -> Bool
&& Bool -> Bool
not (Trie a -> Bool
forall a. Trie a -> Bool
null Trie a
next)
    Bool -> Bool -> Bool
&& Trie a -> Bool
forall a. Trie a -> Bool
valid Trie a
next
valid (Branch Maybe a
_ Map (Trie a)
children) =
  Map (Trie a) -> Int
forall a. Map a -> Int
Map.size Map (Trie a)
children Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
    Bool -> Bool -> Bool
&& (Word8 -> Trie a -> Bool -> Bool) -> Bool -> Map (Trie a) -> Bool
forall a b. (Word8 -> a -> b -> b) -> b -> Map a -> b
Map.foldrWithKeys Word8 -> Trie a -> Bool -> Bool
forall {p} {a}. p -> Trie a -> Bool -> Bool
nonNullChild Bool
True Map (Trie a)
children
 where
  nonNullChild :: p -> Trie a -> Bool -> Bool
nonNullChild p
_ Trie a
child !Bool
acc = Bool
acc Bool -> Bool -> Bool
&& Bool -> Bool
not (Trie a -> Bool
forall a. Trie a -> Bool
null Trie a
child)

------------ Find/Replace ------------

{- | The raison-d'etre of this library: repeatedly search in a byte string
for the longest of multiple patterns and make replacements.
-}
multiFindReplace ::
  (Semigroup b) =>
  -- | construct a portion of the result from unmatched bytes
  (Bytes -> b) ->
  -- | construct a replacement from the found value
  (a -> b) ->
  -- | the dictionary of all replacements
  Trie a ->
  -- | input to be edited
  Bytes ->
  -- | result of replacement
  b
{-# INLINE multiFindReplace #-}
multiFindReplace :: forall b a.
Semigroup b =>
(Bytes -> b) -> (a -> b) -> Trie a -> Bytes -> b
multiFindReplace Bytes -> b
fromNoMatch a -> b
fromMatch = \Trie a
t ->
  let needles :: Trie a
needles = Bytes -> Trie a -> Trie a
forall a. Bytes -> Trie a -> Trie a
delete Bytes
forall a. Monoid a => a
mempty Trie a
t
      -- `into` counts up until the first index where a replacement is found
      go :: Int -> Bytes -> b
go !Int
into Bytes
rawInp =
        let inp :: Bytes
inp = Int -> Bytes -> Bytes
Bytes.unsafeDrop Int
into Bytes
rawInp
            unMatched :: Bytes
unMatched = Int -> Bytes -> Bytes
Bytes.unsafeTake Int
into Bytes
rawInp
         in if
              | Bytes -> Bool
Bytes.null Bytes
inp -> Bytes -> b
fromNoMatch Bytes
unMatched
              | Just (a
val, Bytes
rest) <- Trie a -> Bytes -> Maybe (a, Bytes)
forall a. Trie a -> Bytes -> Maybe (a, Bytes)
stripPrefix Trie a
needles Bytes
inp ->
                  Bytes -> b
fromNoMatch Bytes
unMatched b -> b -> b
forall a. Semigroup a => a -> a -> a
<> a -> b
fromMatch a
val b -> b -> b
forall a. Semigroup a => a -> a -> a
<> Int -> Bytes -> b
go Int
0 Bytes
rest
              | Bool
otherwise -> Int -> Bytes -> b
go (Int
into Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Bytes
rawInp
   in Int -> Bytes -> b
go Int
0

replace :: Trie Bytes -> Bytes -> Chunks
replace :: Trie Bytes -> Bytes -> Chunks
replace Trie Bytes
t Bytes
inp = Int -> Builder -> Chunks
Build.run Int
4080 (Builder -> Chunks) -> Builder -> Chunks
forall a b. (a -> b) -> a -> b
$ Trie Bytes -> Bytes -> Builder
go Trie Bytes
t Bytes
inp
 where
  go :: Trie Bytes -> Bytes -> Builder
go = (Bytes -> Builder)
-> (Bytes -> Builder) -> Trie Bytes -> Bytes -> Builder
forall b a.
Semigroup b =>
(Bytes -> b) -> (a -> b) -> Trie a -> Bytes -> b
multiFindReplace Bytes -> Builder
Build.bytes Bytes -> Builder
Build.bytes

search :: Trie a -> Bytes -> Bool
search :: forall a. Trie a -> Bytes -> Bool
search Trie a
t Bytes
inp = Any -> Bool
getAny (Any -> Bool) -> Any -> Bool
forall a b. (a -> b) -> a -> b
$ Trie a -> Bytes -> Any
forall {a}. Trie a -> Bytes -> Any
go Trie a
t Bytes
inp
 where
  go :: Trie a -> Bytes -> Any
go = (Bytes -> Any) -> (a -> Any) -> Trie a -> Bytes -> Any
forall b a.
Semigroup b =>
(Bytes -> b) -> (a -> b) -> Trie a -> Bytes -> b
multiFindReplace (Any -> Bytes -> Any
forall a b. a -> b -> a
const Any
forall a. Monoid a => a
mempty) (Any -> a -> Any
forall a b. a -> b -> a
const (Any -> a -> Any) -> Any -> a -> Any
forall a b. (a -> b) -> a -> b
$ Bool -> Any
Any Bool
True)

------------ Construction ------------

-- | The empty trie.
empty :: Trie a
empty :: forall a. Trie a
empty = Maybe a -> Trie a
forall a. Maybe a -> Trie a
Tip Maybe a
forall a. Maybe a
U.Nothing

-- | A trie with a single element.
singleton :: Bytes -> a -> Trie a
singleton :: forall a. Bytes -> a -> Trie a
singleton Bytes
k a
v = Bytes -> Trie a -> Trie a
forall a. Bytes -> Trie a -> Trie a
prepend Bytes
k (Trie a -> Trie a) -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ Maybe a -> Trie a
forall a. Maybe a -> Trie a
Tip (a -> Maybe a
forall a. a -> Maybe a
U.Just a
v)

{- | Prepend every key in the 'Trie' with the given 'Bytes'.

This should be used internally instead of the Run ctor,
thereby ensuring the run length >= 2 invariant is maintained.
It is exported anyway because someone may find it useful.
-}
prepend :: Bytes -> Trie a -> Trie a
prepend :: forall a. Bytes -> Trie a -> Trie a
prepend Bytes
bytes Trie a
next = Maybe a -> ByteArray -> Trie a -> Trie a
forall a. Maybe a -> ByteArray -> Trie a -> Trie a
Run Maybe a
forall a. Maybe a
U.Nothing (Bytes -> ByteArray
toByteArray Bytes
bytes) Trie a
next

{- | Insert a new key/value into the trie.
If the key is already present in the trie, the associated value is
replaced with the new one.
'insert' is equivalent to 'insertWith' 'const'.
-}
insert :: Bytes -> a -> Trie a -> Trie a
insert :: forall a. Bytes -> a -> Trie a -> Trie a
insert = (a -> a -> a) -> Bytes -> a -> Trie a -> Trie a
forall a. (a -> a -> a) -> Bytes -> a -> Trie a -> Trie a
insertWith a -> a -> a
forall a b. a -> b -> a
const

{- | Insert with a function, combining new value and old value.
@'insertWith' f key value trie@ will insert the pair @(key, value)@
into @trie@ if @key@ does not exist in the trie.
If the key does exist, the function will insert the pair
@(key, f new_value old_value)@.
-}
insertWith :: (a -> a -> a) -> Bytes -> a -> Trie a -> Trie a
insertWith :: forall a. (a -> a -> a) -> Bytes -> a -> Trie a -> Trie a
insertWith a -> a -> a
f Bytes
k a
v = (a -> a -> a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> a) -> Trie a -> Trie a -> Trie a
unionWith a -> a -> a
f (Bytes -> a -> Trie a
forall a. Bytes -> a -> Trie a
singleton Bytes
k a
v)

delete :: Bytes -> Trie a -> Trie a
delete :: forall a. Bytes -> Trie a -> Trie a
delete Bytes
k0 Trie a
trie = Bytes -> Trie a -> Trie a
forall a. Bytes -> Trie a -> Trie a
go Bytes
k0 Trie a
trie
 where
  -- `go` is not always tail-recursive.
  -- Instead, each node with exactly one child must be checked after the
  --  deletion to ensure that child is non-empty.
  -- TODO
  -- However, as soon as it is known that the size must be greater than one,
  --  we can throw away all queued normalizations so far.
  -- Therefore, we maintain a delimited continuation as an accumulator,
  --  but I'm not yet sure how to manually store it.
  -- go :: Bytes -> Trie a
  go :: Bytes -> Trie a -> Trie a
go Bytes
key node :: Trie a
node@(Tip Maybe a
v)
    | Bytes -> Bool
Bytes.null Bytes
key
    , U.Just a
_ <- Maybe a
v -- NOTE this is redundant now, but when I use cps, it won't be
      =
        Trie a
forall a. Trie a
empty
    | Bool
otherwise = Trie a
node
  go Bytes
key node :: Trie a
node@(Run Maybe a
v (ByteArray -> Bytes
fromByteArray -> Bytes
run) Trie a
next)
    -- found key, therefore delete
    | Bytes -> Bool
Bytes.null Bytes
key
    , U.Just a
_ <- Maybe a
v -- NOTE this is redundant now, but when I use cps, it won't be
      =
        Bytes -> Trie a -> Trie a
forall a. Bytes -> Trie a -> Trie a
prepend Bytes
run Trie a
next
    -- carry on searching for the key
    | Just Bytes
key' <- Bytes -> Bytes -> Maybe Bytes
Bytes.stripPrefix Bytes
run Bytes
key =
        Maybe a -> ByteArray -> Trie a -> Trie a
forall a. Maybe a -> ByteArray -> Trie a -> Trie a
Run Maybe a
v (Bytes -> ByteArray
toByteArray Bytes
run) (Bytes -> Trie a -> Trie a
go Bytes
key' Trie a
next)
    -- key not present
    | Bool
otherwise = Trie a
node
  go Bytes
key node :: Trie a
node@(Branch Maybe a
v Map (Trie a)
children)
    -- found key, therefore delete
    | Bytes -> Bool
Bytes.null Bytes
key
    , U.Just a
_ <- Maybe a
v =
        Maybe a -> Map (Trie a) -> Trie a
forall a. Maybe a -> Map (Trie a) -> Trie a
UnsafeBranch Maybe a
forall a. Maybe a
U.Nothing Map (Trie a)
children
    -- carry on searching for the key
    | Just (Word8
c, Bytes
key') <- Bytes -> Maybe (Word8, Bytes)
Bytes.uncons Bytes
key
    , Just Trie a
child <- Word8 -> Map (Trie a) -> Maybe (Trie a)
forall a. Word8 -> Map a -> Maybe a
Map.lookup Word8
c Map (Trie a)
children =
        Maybe a -> Map (Trie a) -> Trie a
forall a. Maybe a -> Map (Trie a) -> Trie a
Branch Maybe a
v (Word8 -> Trie a -> Map (Trie a) -> Map (Trie a)
forall a. Word8 -> a -> Map a -> Map a
Map.insert Word8
c (Bytes -> Trie a -> Trie a
go Bytes
key' Trie a
child) Map (Trie a)
children)
    -- key not present
    | Bool
otherwise = Trie a
node

{- | Union of the two tries, but where a key appears in both,
the associated values are combined with '(<>)' to produce the new value,
i.e. @append == unionWith (<>)@.
-}
append :: (Semigroup a) => Trie a -> Trie a -> Trie a
append :: forall a. Semigroup a => Trie a -> Trie a -> Trie a
append = (a -> a -> a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> a) -> Trie a -> Trie a -> Trie a
unionWith a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>)

{- | The left-biased union of the two tries.
It prefers the first when duplicate keys are encountered,
i.e. @union == unionWith const@.
-}
union :: Trie a -> Trie a -> Trie a
union :: forall a. Trie a -> Trie a -> Trie a
union = (a -> a -> a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> a) -> Trie a -> Trie a -> Trie a
unionWith a -> a -> a
forall a b. a -> b -> a
const

-- | Union with a combining function.
unionWith :: (a -> a -> a) -> Trie a -> Trie a -> Trie a
unionWith :: forall a. (a -> a -> a) -> Trie a -> Trie a -> Trie a
unionWith a -> a -> a
f Trie a
trieA Trie a
trieB = case (Trie a
trieA, Trie a
trieB) of
  (Tip Maybe a
a, Tip Maybe a
b) -> Maybe a -> Trie a
forall a. Maybe a -> Trie a
Tip (Maybe a
a Maybe a -> Maybe a -> Maybe a
`mergeValue` Maybe a
b)
  (Tip Maybe a
a, Run Maybe a
b ByteArray
run Trie a
next) -> Maybe a -> ByteArray -> Trie a -> Trie a
forall a. Maybe a -> ByteArray -> Trie a -> Trie a
UnsafeRun (Maybe a
a Maybe a -> Maybe a -> Maybe a
`mergeValue` Maybe a
b) ByteArray
run Trie a
next
  (Run Maybe a
a ByteArray
run Trie a
next, Tip Maybe a
b) -> Maybe a -> ByteArray -> Trie a -> Trie a
forall a. Maybe a -> ByteArray -> Trie a -> Trie a
UnsafeRun (Maybe a
a Maybe a -> Maybe a -> Maybe a
`mergeValue` Maybe a
b) ByteArray
run Trie a
next
  (Tip Maybe a
a, Branch Maybe a
b Map (Trie a)
children) -> Maybe a -> Map (Trie a) -> Trie a
forall a. Maybe a -> Map (Trie a) -> Trie a
UnsafeBranch (Maybe a
a Maybe a -> Maybe a -> Maybe a
`mergeValue` Maybe a
b) Map (Trie a)
children
  (Branch Maybe a
a Map (Trie a)
children, Tip Maybe a
b) -> Maybe a -> Map (Trie a) -> Trie a
forall a. Maybe a -> Map (Trie a) -> Trie a
UnsafeBranch (Maybe a
a Maybe a -> Maybe a -> Maybe a
`mergeValue` Maybe a
b) Map (Trie a)
children
  -- all non-Tip cases
  (Branch Maybe a
a Map (Trie a)
children, Branch Maybe a
b Map (Trie a)
children') ->
    Maybe a -> Map (Trie a) -> Trie a
forall a. Maybe a -> Map (Trie a) -> Trie a
UnsafeBranch (Maybe a
a Maybe a -> Maybe a -> Maybe a
`mergeValue` Maybe a
b) (Map (Trie a) -> Map (Trie a) -> Map (Trie a)
mergeChildren Map (Trie a)
children Map (Trie a)
children')
  (Branch Maybe a
a Map (Trie a)
children, r :: Trie a
r@(Run Maybe a
_ ByteArray
_ Trie a
_)) ->
    Maybe a -> Map (Trie a) -> Trie a
forall a. Maybe a -> Map (Trie a) -> Trie a
UnsafeBranch (Maybe a
a Maybe a -> Maybe a -> Maybe a
`mergeValue` Maybe a
b) (Map (Trie a) -> Map (Trie a) -> Map (Trie a)
mergeChildren Map (Trie a)
children Map (Trie a)
children')
   where
    (Maybe a
b, Word8
c, Trie a
child') = Trie a -> (Maybe a, Word8, Trie a)
forall a. Trie a -> (Maybe a, Word8, Trie a)
unsafeUnconsRun Trie a
r
    children' :: Map (Trie a)
children' = Word8 -> Trie a -> Map (Trie a)
forall a. Word8 -> a -> Map a
Map.singleton Word8
c Trie a
child'
  (r :: Trie a
r@(Run Maybe a
_ ByteArray
_ Trie a
_), Branch Maybe a
b Map (Trie a)
children') ->
    Maybe a -> Map (Trie a) -> Trie a
forall a. Maybe a -> Map (Trie a) -> Trie a
UnsafeBranch (Maybe a
a Maybe a -> Maybe a -> Maybe a
`mergeValue` Maybe a
b) (Map (Trie a) -> Map (Trie a) -> Map (Trie a)
mergeChildren Map (Trie a)
children Map (Trie a)
children')
   where
    (Maybe a
a, Word8
c, Trie a
child') = Trie a -> (Maybe a, Word8, Trie a)
forall a. Trie a -> (Maybe a, Word8, Trie a)
unsafeUnconsRun Trie a
r
    children :: Map (Trie a)
children = Word8 -> Trie a -> Map (Trie a)
forall a. Word8 -> a -> Map a
Map.singleton Word8
c Trie a
child'
  (Run Maybe a
a (ByteArray -> Bytes
fromByteArray -> Bytes
run) Trie a
next, Run Maybe a
b (ByteArray -> Bytes
fromByteArray -> Bytes
run') Trie a
next') ->
    if Bytes -> Bool
Bytes.null Bytes
common
      then
        let mkChild :: Bytes -> Trie a -> Map (Trie a)
mkChild Bytes
bytes Trie a
trie = case Bytes -> Maybe (Word8, Bytes)
Bytes.uncons Bytes
bytes of
              Prelude.Just (Word8
c, Bytes
k) -> Word8 -> Trie a -> Map (Trie a)
forall a. Word8 -> a -> Map a
Map.singleton Word8
c (Bytes -> Trie a -> Trie a
forall a. Bytes -> Trie a -> Trie a
prepend Bytes
k Trie a
trie)
              Maybe (Word8, Bytes)
Prelude.Nothing -> String -> Map (Trie a)
forall a. HasCallStack => String -> a
error String
"invariant violation: empty run bytes"
            child :: Map (Trie a)
child = Bytes -> Trie a -> Map (Trie a)
forall {a}. Bytes -> Trie a -> Map (Trie a)
mkChild Bytes
run Trie a
next
            child' :: Map (Trie a)
child' = Bytes -> Trie a -> Map (Trie a)
forall {a}. Bytes -> Trie a -> Map (Trie a)
mkChild Bytes
run' Trie a
next'
         in Maybe a -> Map (Trie a) -> Trie a
forall a. Maybe a -> Map (Trie a) -> Trie a
UnsafeBranch (Maybe a
a Maybe a -> Maybe a -> Maybe a
`mergeValue` Maybe a
b) (Map (Trie a) -> Trie a) -> Map (Trie a) -> Trie a
forall a b. (a -> b) -> a -> b
$ Map (Trie a)
child Map (Trie a) -> Map (Trie a) -> Map (Trie a)
forall a. Map a -> Map a -> Map a
`Map.union` Map (Trie a)
child'
      else
        let child :: Trie a
child = Bytes -> Trie a -> Trie a
forall a. Bytes -> Trie a -> Trie a
prepend (Bytes -> Bytes
uncommon Bytes
run) Trie a
next
            child' :: Trie a
child' = Bytes -> Trie a -> Trie a
forall a. Bytes -> Trie a -> Trie a
prepend (Bytes -> Bytes
uncommon Bytes
run') Trie a
next'
         in Maybe a -> ByteArray -> Trie a -> Trie a
forall a. Maybe a -> ByteArray -> Trie a -> Trie a
Run (Maybe a
a Maybe a -> Maybe a -> Maybe a
`mergeValue` Maybe a
b) (Bytes -> ByteArray
toByteArray Bytes
common) (Trie a -> Trie a) -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ (a -> a -> a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> a) -> Trie a -> Trie a -> Trie a
unionWith a -> a -> a
f Trie a
child Trie a
child'
   where
    common :: Bytes
common = Bytes -> Bytes -> Bytes
Bytes.longestCommonPrefix Bytes
run Bytes
run'
    uncommon :: Bytes -> Bytes
uncommon Bytes
bytes = Int -> Bytes -> Bytes
Bytes.unsafeDrop (Bytes -> Int
Bytes.length Bytes
common) Bytes
bytes
 where
  mergeChildren :: Map (Trie a) -> Map (Trie a) -> Map (Trie a)
mergeChildren Map (Trie a)
left Map (Trie a)
right = (Trie a -> Trie a -> Trie a)
-> Map (Trie a) -> Map (Trie a) -> Map (Trie a)
forall a. (a -> a -> a) -> Map a -> Map a -> Map a
Map.unionWith ((a -> a -> a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> a) -> Trie a -> Trie a -> Trie a
unionWith a -> a -> a
f) Map (Trie a)
left Map (Trie a)
right
  mergeValue :: Maybe a -> Maybe a -> Maybe a
mergeValue Maybe a
U.Nothing Maybe a
U.Nothing = Maybe a
forall a. Maybe a
U.Nothing
  mergeValue (U.Just a
x) (U.Just a
y) = a -> Maybe a
forall a. a -> Maybe a
U.Just (a -> a -> a
f a
x a
y)
  mergeValue Maybe a
x Maybe a
y = Maybe a
x Maybe a -> Maybe a -> Maybe a
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe a
y

------------ Conversion ------------

{- | Build a trie from a list of key/value pairs.
If more than one value for the same key appears, the last value for that
key is retained.
-}
fromList :: [(Bytes, a)] -> Trie a
fromList :: forall a. [(Bytes, a)] -> Trie a
fromList [(Bytes, a)]
kvs = (Trie a -> (Bytes, a) -> Trie a)
-> Trie a -> [(Bytes, a)] -> Trie a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' (\Trie a
xs (Bytes
k, a
v) -> Bytes -> a -> Trie a -> Trie a
forall a. Bytes -> a -> Trie a -> Trie a
insert Bytes
k a
v Trie a
xs) Trie a
forall a. Trie a
empty [(Bytes, a)]
kvs

{- | Convert the trie to a list of key/value pairs.
The resulting list has its keys sorted in ascending order.
-}
toList :: Trie a -> [(Bytes, a)]
toList :: forall a. Trie a -> [(Bytes, a)]
toList = \case
  Tip Maybe a
valO -> Maybe a -> [(Bytes, a)]
forall {t} {a}. Monoid t => Maybe a -> [(t, a)]
fromValue Maybe a
valO
  Run Maybe a
valO ByteArray
run Trie a
next -> Maybe a -> [(Bytes, a)]
forall {t} {a}. Monoid t => Maybe a -> [(t, a)]
fromValue Maybe a
valO [(Bytes, a)] -> [(Bytes, a)] -> [(Bytes, a)]
forall a. [a] -> [a] -> [a]
++ Bytes -> [(Bytes, a)] -> [(Bytes, a)]
forall {f :: * -> *} {p :: * -> * -> *} {b} {c}.
(Functor f, Bifunctor p, Semigroup b) =>
b -> f (p b c) -> f (p b c)
prependList (ByteArray -> Bytes
fromByteArray ByteArray
run) (Trie a -> [(Bytes, a)]
forall a. Trie a -> [(Bytes, a)]
toList Trie a
next)
  Branch Maybe a
valO Map (Trie a)
children -> Maybe a -> [(Bytes, a)]
forall {t} {a}. Monoid t => Maybe a -> [(t, a)]
fromValue Maybe a
valO [(Bytes, a)] -> [(Bytes, a)] -> [(Bytes, a)]
forall a. [a] -> [a] -> [a]
++ (Word8 -> Trie a -> [(Bytes, a)] -> [(Bytes, a)])
-> [(Bytes, a)] -> Map (Trie a) -> [(Bytes, a)]
forall a b. (Word8 -> a -> b -> b) -> b -> Map a -> b
Map.foldrWithKeys Word8 -> Trie a -> [(Bytes, a)] -> [(Bytes, a)]
forall {c}. Word8 -> Trie c -> [(Bytes, c)] -> [(Bytes, c)]
f [] Map (Trie a)
children
   where
    f :: Word8 -> Trie c -> [(Bytes, c)] -> [(Bytes, c)]
f Word8
k Trie c
v [(Bytes, c)]
acc = Bytes -> [(Bytes, c)] -> [(Bytes, c)]
forall {f :: * -> *} {p :: * -> * -> *} {b} {c}.
(Functor f, Bifunctor p, Semigroup b) =>
b -> f (p b c) -> f (p b c)
prependList (Word8 -> Bytes
Bytes.singleton Word8
k) (Trie c -> [(Bytes, c)]
forall a. Trie a -> [(Bytes, a)]
toList Trie c
v) [(Bytes, c)] -> [(Bytes, c)] -> [(Bytes, c)]
forall a. [a] -> [a] -> [a]
++ [(Bytes, c)]
acc
 where
  fromValue :: Maybe a -> [(t, a)]
fromValue Maybe a
valO = (t
forall a. Monoid a => a
mempty,) (a -> (t, a)) -> [a] -> [(t, a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe a -> [a]
forall a. Maybe a -> [a]
U.maybeToList Maybe a
valO
  prependList :: b -> f (p b c) -> f (p b c)
prependList b
run f (p b c)
list = (b -> b) -> p b c -> p b c
forall a b c. (a -> b) -> p a c -> p b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (b
run b -> b -> b
forall a. Semigroup a => a -> a -> a
<>) (p b c -> p b c) -> f (p b c) -> f (p b c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (p b c)
list

foldl' :: (b -> a -> b) -> b -> Trie a -> b
{-# INLINE foldl' #-}
foldl' :: forall b a. (b -> a -> b) -> b -> Trie a -> b
foldl' b -> a -> b
f !b
b0 Trie a
t0 = b -> Trie a -> b
go b
b0 Trie a
t0
 where
  go :: b -> Trie a -> b
go !b
b Trie a
t = case Trie a
t of
    Tip Maybe a
valO -> case Maybe a
valO of
      U.Just a
x -> b -> a -> b
f b
b a
x
      Maybe a
_ -> b
b
    Run Maybe a
valO ByteArray
_ Trie a
next ->
      let b' :: b
b' = case Maybe a
valO of
            U.Just a
x -> b -> a -> b
f b
b a
x
            Maybe a
_ -> b
b
       in b -> Trie a -> b
go b
b' Trie a
next
    Branch Maybe a
valO Map (Trie a)
children ->
      let b' :: b
b' = case Maybe a
valO of
            U.Just a
x -> b -> a -> b
f b
b a
x
            Maybe a
_ -> b
b
       in (b -> Trie a -> b) -> b -> Map (Trie a) -> b
forall b a. (b -> a -> b) -> b -> Map a -> b
Map.foldl' b -> Trie a -> b
go b
b' Map (Trie a)
children

traverse_ :: (Applicative m) => (a -> m b) -> Trie a -> m ()
{-# INLINE traverse_ #-}
traverse_ :: forall (m :: * -> *) a b.
Applicative m =>
(a -> m b) -> Trie a -> m ()
traverse_ a -> m b
f Trie a
t0 = Trie a -> m ()
go Trie a
t0
 where
  go :: Trie a -> m ()
go Trie a
t = case Trie a
t of
    Tip Maybe a
valO -> case Maybe a
valO of
      U.Just a
x -> a -> m b
f a
x m b -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
      Maybe a
_ -> () -> m ()
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Run Maybe a
valO ByteArray
_ Trie a
next -> case Maybe a
valO of
      U.Just a
x -> a -> m b
f a
x m b -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Trie a -> m ()
go Trie a
next
      Maybe a
_ -> Trie a -> m ()
go Trie a
next
    Branch Maybe a
valO Map (Trie a)
children -> case Maybe a
valO of
      U.Just a
x -> a -> m b
f a
x m b -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> (Trie a -> m ()) -> Map (Trie a) -> m ()
forall (m :: * -> *) a b.
Applicative m =>
(a -> m b) -> Map a -> m ()
Map.traverse_ Trie a -> m ()
go Map (Trie a)
children
      Maybe a
_ -> (Trie a -> m ()) -> Map (Trie a) -> m ()
forall (m :: * -> *) a b.
Applicative m =>
(a -> m b) -> Map a -> m ()
Map.traverse_ Trie a -> m ()
go Map (Trie a)
children

------------ Query ------------

-- | Lookup the value at the 'Bytes' key in the trie.
lookup :: Bytes -> Trie a -> Maybe a
lookup :: forall a. Bytes -> Trie a -> Maybe a
lookup Bytes
k (Tip Maybe a
v)
  | Bytes -> Bool
Bytes.null Bytes
k = Maybe a -> Maybe a
forall a. Maybe a -> Maybe a
U.toBaseMaybe Maybe a
v
  | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing
lookup Bytes
k (Run Maybe a
v (ByteArray -> Bytes
fromByteArray -> Bytes
run) Trie a
next)
  | Bytes -> Bool
Bytes.null Bytes
k = Maybe a -> Maybe a
forall a. Maybe a -> Maybe a
U.toBaseMaybe Maybe a
v
  | Bytes
run Bytes -> Bytes -> Bool
`Bytes.isPrefixOf` Bytes
k =
      let k' :: Bytes
k' = Int -> Bytes -> Bytes
Bytes.unsafeDrop (Bytes -> Int
Bytes.length Bytes
run) Bytes
k
       in Bytes -> Trie a -> Maybe a
forall a. Bytes -> Trie a -> Maybe a
lookup Bytes
k' Trie a
next
  | Bool
otherwise = Maybe a
forall a. Maybe a
Prelude.Nothing
lookup Bytes
k (Branch Maybe a
valO Map (Trie a)
children) = case Bytes -> Maybe (Word8, Bytes)
Bytes.uncons Bytes
k of
  Maybe (Word8, Bytes)
Prelude.Nothing -> Maybe a -> Maybe a
forall a. Maybe a -> Maybe a
U.toBaseMaybe Maybe a
valO
  Prelude.Just (Word8
c, Bytes
k') -> Bytes -> Trie a -> Maybe a
forall a. Bytes -> Trie a -> Maybe a
lookup Bytes
k' (Trie a -> Maybe a) -> Maybe (Trie a) -> Maybe a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Word8 -> Map (Trie a) -> Maybe (Trie a)
forall a. Word8 -> Map a -> Maybe a
Map.lookup Word8
c Map (Trie a)
children

{- | Lookup the value at the 'Bytes' key in the trie. Returns the value
of the exact match and the values for any keys that are prefixes of
the search key. The shortest prefix is first. The exact match (if there
is one) is last.
-}
lookupPrefixes :: Bytes -> Trie a -> [a]
{-# INLINE lookupPrefixes #-}
lookupPrefixes :: forall a. Bytes -> Trie a -> [a]
lookupPrefixes = [a] -> Bytes -> Trie a -> [a]
forall a. [a] -> Bytes -> Trie a -> [a]
lookupPrefixesGo []

lookupPrefixesGo :: [a] -> Bytes -> Trie a -> [a]
lookupPrefixesGo :: forall a. [a] -> Bytes -> Trie a -> [a]
lookupPrefixesGo ![a]
acc !Bytes
_ (Tip Maybe a
v)
  | U.Just a
x <- Maybe a
v = a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc
  | Bool
otherwise = [a]
acc
lookupPrefixesGo ![a]
acc !Bytes
k (UnsafeRun Maybe a
v (ByteArray -> Bytes
fromByteArray -> Bytes
run) Trie a
next) =
  let acc' :: [a]
acc' = case Maybe a
v of
        U.Just a
x -> a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc
        Maybe a
_ -> [a]
acc
   in if
        | Bytes -> Bool
Bytes.null Bytes
k -> [a]
acc'
        | Bytes
run Bytes -> Bytes -> Bool
`Bytes.isPrefixOf` Bytes
k ->
            let k' :: Bytes
k' = Int -> Bytes -> Bytes
Bytes.unsafeDrop (Bytes -> Int
Bytes.length Bytes
run) Bytes
k
             in [a] -> Bytes -> Trie a -> [a]
forall a. [a] -> Bytes -> Trie a -> [a]
lookupPrefixesGo [a]
acc' Bytes
k' Trie a
next
        | Bool
otherwise -> [a]
acc'
lookupPrefixesGo ![a]
acc !Bytes
k (UnsafeBranch Maybe a
valO Map (Trie a)
children) =
  let acc' :: [a]
acc' = case Maybe a
valO of
        U.Just a
x -> a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc
        Maybe a
_ -> [a]
acc
   in case Bytes -> Maybe (Word8, Bytes)
Bytes.uncons Bytes
k of
        Maybe (Word8, Bytes)
Prelude.Nothing -> [a]
acc'
        Prelude.Just (Word8
c, Bytes
k') -> case Word8 -> Map (Trie a) -> Maybe (Trie a)
forall a. Word8 -> Map a -> Maybe a
Map.lookup Word8
c Map (Trie a)
children of
          Maybe (Trie a)
Nothing -> [a]
acc'
          Just Trie a
child -> [a] -> Bytes -> Trie a -> [a]
forall a. [a] -> Bytes -> Trie a -> [a]
lookupPrefixesGo [a]
acc' Bytes
k' Trie a
child

{- | Lookup the trie at the 'Bytes' key in the trie. Returns the subtrie
at this position.

>>> (k1 <> k2 == k) ==> (lookup k v t == lookup k2 (lookupTrie k1 t))
-}
lookupTrie :: Bytes -> Trie a -> Trie a
lookupTrie :: forall a. Bytes -> Trie a -> Trie a
lookupTrie !Bytes
k Trie a
trie
  | Bytes -> Bool
Bytes.null Bytes
k = Trie a
trie
  | Bool
otherwise = case Trie a
trie of
      Tip {} -> Trie a
forall a. Trie a
empty
      Run Maybe a
_ (ByteArray -> Bytes
fromByteArray -> Bytes
run) Trie a
next
        | Bytes
run Bytes -> Bytes -> Bool
`Bytes.isPrefixOf` Bytes
k ->
            let k' :: Bytes
k' = Int -> Bytes -> Bytes
Bytes.unsafeDrop (Bytes -> Int
Bytes.length Bytes
run) Bytes
k
             in Bytes -> Trie a -> Trie a
forall a. Bytes -> Trie a -> Trie a
lookupTrie Bytes
k' Trie a
next
        | Bytes
k Bytes -> Bytes -> Bool
`Bytes.isPrefixOf` Bytes
run ->
            let run' :: Bytes
run' = Int -> Bytes -> Bytes
Bytes.unsafeDrop (Bytes -> Int
Bytes.length Bytes
k) Bytes
run
             in Maybe a -> ByteArray -> Trie a -> Trie a
forall a. Maybe a -> ByteArray -> Trie a -> Trie a
Run Maybe a
forall a. Maybe a
U.Nothing (Bytes -> ByteArray
Bytes.toByteArrayClone Bytes
run') Trie a
next
        | Bool
otherwise -> Trie a
forall a. Trie a
empty
      Branch Maybe a
_ Map (Trie a)
children ->
        let !k' :: Bytes
k' = Int -> Bytes -> Bytes
Bytes.unsafeDrop Int
1 Bytes
k
            !c :: Word8
c = Bytes -> Int -> Word8
Bytes.unsafeIndex Bytes
k Int
0
         in case Word8 -> Map (Trie a) -> Maybe (Trie a)
forall a. Word8 -> Map a -> Maybe a
Map.lookup Word8
c Map (Trie a)
children of
              Maybe (Trie a)
Nothing -> Trie a
forall a. Trie a
empty
              Just Trie a
child -> Bytes -> Trie a -> Trie a
forall a. Bytes -> Trie a -> Trie a
lookupTrie Bytes
k' Trie a
child

{- | Find the longest prefix of the input 'Bytes' which has a value in the trie.
Returns the associated value and the remainder of the input after the prefix.
-}
stripPrefix :: Trie a -> Bytes -> Maybe (a, Bytes)
stripPrefix :: forall a. Trie a -> Bytes -> Maybe (a, Bytes)
stripPrefix Trie a
trie Bytes
inp = ((Bytes, a) -> a) -> ((Bytes, a), Bytes) -> (a, Bytes)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (Bytes, a) -> a
forall a b. (a, b) -> b
snd (((Bytes, a), Bytes) -> (a, Bytes))
-> Maybe ((Bytes, a), Bytes) -> Maybe (a, Bytes)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Trie a -> Bytes -> Maybe ((Bytes, a), Bytes)
forall a. Trie a -> Bytes -> Maybe ((Bytes, a), Bytes)
stripPrefixWithKey Trie a
trie Bytes
inp

{- | Find the longest prefix of the input 'Bytes' which has a value in the trie.
Returns the prefix and associated value found as a key/value tuple,
and also the remainder of the input after the prefix.
-}
stripPrefixWithKey :: forall a. Trie a -> Bytes -> Maybe ((Bytes, a), Bytes)
stripPrefixWithKey :: forall a. Trie a -> Bytes -> Maybe ((Bytes, a), Bytes)
stripPrefixWithKey Trie a
trie0 Bytes
rawInp = Int -> Maybe (Bytes, a) -> Trie a -> Maybe ((Bytes, a), Bytes)
go Int
0 Maybe (Bytes, a)
forall a. Maybe a
Nothing Trie a
trie0
 where
  go :: Int -> Maybe (Bytes, a) -> Trie a -> Maybe ((Bytes, a), Bytes)
  go :: Int -> Maybe (Bytes, a) -> Trie a -> Maybe ((Bytes, a), Bytes)
go !Int
into !Maybe (Bytes, a)
prior Trie a
node =
    let inp :: Bytes
inp = Int -> Bytes -> Bytes
Bytes.unsafeDrop Int
into Bytes
rawInp
        candidate :: Maybe (Bytes, a)
candidate = (Int -> Bytes -> Bytes
Bytes.unsafeTake Int
into Bytes
rawInp,) (a -> (Bytes, a)) -> Maybe a -> Maybe (Bytes, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Trie a -> Maybe a
forall a. Trie a -> Maybe a
topValue Trie a
node
        found :: Maybe (Bytes, a)
found = Maybe (Bytes, a)
candidate Maybe (Bytes, a) -> Maybe (Bytes, a) -> Maybe (Bytes, a)
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Maybe (Bytes, a)
prior
     in if
          | Run Maybe a
_ (ByteArray -> Bytes
fromByteArray -> Bytes
run) Trie a
next <- Trie a
node
          , Bytes
run Bytes -> Bytes -> Bool
`Bytes.isPrefixOf` Bytes
inp ->
              Int -> Maybe (Bytes, a) -> Trie a -> Maybe ((Bytes, a), Bytes)
go (Int
into Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Bytes -> Int
Bytes.length Bytes
run) Maybe (Bytes, a)
found Trie a
next
          | Branch Maybe a
_ Map (Trie a)
children <- Trie a
node
          , Just (Word8
c, Bytes
_) <- Bytes -> Maybe (Word8, Bytes)
Bytes.uncons Bytes
inp
          , Just Trie a
next <- Word8 -> Map (Trie a) -> Maybe (Trie a)
forall a. Word8 -> Map a -> Maybe a
Map.lookup Word8
c Map (Trie a)
children ->
              Int -> Maybe (Bytes, a) -> Trie a -> Maybe ((Bytes, a), Bytes)
go (Int
into Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Maybe (Bytes, a)
found Trie a
next
          | Bool
otherwise -> (Bytes, a) -> ((Bytes, a), Bytes)
forall {b}. (Bytes, b) -> ((Bytes, b), Bytes)
mkReturn ((Bytes, a) -> ((Bytes, a), Bytes))
-> Maybe (Bytes, a) -> Maybe ((Bytes, a), Bytes)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Bytes, a)
found
  mkReturn :: (Bytes, b) -> ((Bytes, b), Bytes)
mkReturn (Bytes
prefix, b
v) =
    let post :: Bytes
post = Int -> Bytes -> Bytes
Bytes.unsafeDrop (Bytes -> Int
Bytes.length Bytes
prefix) Bytes
rawInp
     in ((Bytes
prefix, b
v), Bytes
post)

null :: Trie a -> Bool
null :: forall a. Trie a -> Bool
null (Tip Maybe a
U.Nothing) = Bool
True
null Trie a
_ = Bool
False

size :: Trie a -> Int
size :: forall a. Trie a -> Int
size Trie a
node = Int
here Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
under
 where
  here :: Int
here = Int -> (a -> Int) -> Maybe a -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 (Int -> a -> Int
forall a b. a -> b -> a
const Int
1) (Trie a -> Maybe a
forall a. Trie a -> Maybe a
topValue Trie a
node)
  under :: Int
under = case Trie a
node of
    Tip Maybe a
_ -> Int
0
    Run Maybe a
_ ByteArray
_ Trie a
next -> Trie a -> Int
forall a. Trie a -> Int
size Trie a
next
    Branch Maybe a
_ Map (Trie a)
children -> (Word8 -> Trie a -> Int -> Int) -> Int -> Map (Trie a) -> Int
forall a b. (Word8 -> a -> b -> b) -> b -> Map a -> b
Map.foldrWithKeys (\Word8
_ Trie a
child !Int
acc -> Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Trie a -> Int
forall a. Trie a -> Int
size Trie a
child) Int
0 Map (Trie a)
children

------ Helpers ------

topValue :: Trie a -> Maybe a
topValue :: forall a. Trie a -> Maybe a
topValue = \case
  Tip Maybe a
v -> Maybe a -> Maybe a
forall a. Maybe a -> Maybe a
U.toBaseMaybe Maybe a
v
  Run Maybe a
v ByteArray
_ Trie a
_ -> Maybe a -> Maybe a
forall a. Maybe a -> Maybe a
U.toBaseMaybe Maybe a
v
  Branch Maybe a
v Map (Trie a)
_ -> Maybe a -> Maybe a
forall a. Maybe a -> Maybe a
U.toBaseMaybe Maybe a
v

unsafeUnconsRun :: Trie a -> (U.Maybe a, Word8, Trie a)
unsafeUnconsRun :: forall a. Trie a -> (Maybe a, Word8, Trie a)
unsafeUnconsRun (Run Maybe a
v0 ByteArray
bs Trie a
next) = (Maybe a
v0, Word8
c, Trie a
run')
 where
  c :: Word8
c = ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
bs Int
0
  bs' :: Bytes
bs' = Int -> Bytes -> Bytes
Bytes.unsafeDrop Int
1 (ByteArray -> Bytes
fromByteArray ByteArray
bs)
  run' :: Trie a
run' = Bytes -> Trie a -> Trie a
forall a. Bytes -> Trie a -> Trie a
prepend Bytes
bs' Trie a
next
unsafeUnconsRun (Tip Maybe a
_) = String -> (Maybe a, Word8, Trie a)
forall a. HasCallStack => String -> a
error String
"unsafeUnconsRun on Tip trie"
unsafeUnconsRun (Branch Maybe a
_ Map (Trie a)
_) = String -> (Maybe a, Word8, Trie a)
forall a. HasCallStack => String -> a
error String
"unsafeUnconsRun on Branch trie"

-- TODO is this really a decent way to do this?
fromSingletonMap :: Map a -> Maybe (Word8, a)
fromSingletonMap :: forall a. Map a -> Maybe (Word8, a)
fromSingletonMap Map a
mp = case Map a -> [(Word8, a)]
forall a. Map a -> [(Word8, a)]
Map.toList Map a
mp of
  [(Word8
c, a
v)] -> (Word8, a) -> Maybe (Word8, a)
forall a. a -> Maybe a
Just (Word8
c, a
v)
  [(Word8, a)]
_ -> Maybe (Word8, a)
forall a. Maybe a
Nothing