-- To make GHC stop warning about the Prelude
{-# OPTIONS_GHC -Wall -fwarn-tabs -fno-warn-unused-imports #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- For list fusion on toListBy, and guarding `base` versions.
{-# LANGUAGE CPP #-}

------------------------------------------------------------
--                                              ~ 2019.02.24
-- |
-- Module      :  Data.Trie.Internal
-- Copyright   :  Copyright (c) 2008--2019 wren gayle romano
-- License     :  BSD3
-- Maintainer  :  wren@cpan.org
-- Stability   :  provisional
-- Portability :  portable (with CPP)
--
-- Internal definition of the 'Trie' data type and generic functions
-- for manipulating them. Almost everything here is re-exported
-- from "Data.Trie", which is the preferred API for users. This
-- module is for developers who need deeper (and potentially fragile)
-- access to the abstract type.
------------------------------------------------------------

module Data.Trie.Internal
    (
    -- * Data types
      Trie(), showTrie

    -- * Functions for 'ByteString's
    , breakMaximalPrefix

    -- * Basic functions
    , empty, null, singleton, size

    -- * Conversion and folding functions
    , foldrWithKey, toListBy

    -- * Query functions
    , lookupBy_, submap
    , match_, matches_

    -- * Single-value modification
    , alterBy, alterBy_, adjustBy

    -- * Combining tries
    , mergeBy

    -- * Mapping functions
    , mapBy
    , filterMap
    , contextualMap
    , contextualMap'
    , contextualFilterMap
    , contextualMapBy

    -- * Priority-queue functions
    , minAssoc, maxAssoc
    , updateMinViewBy, updateMaxViewBy
    ) where

import Prelude hiding    (null, lookup)
import qualified Prelude (null, lookup)

import qualified Data.ByteString as S
import Data.Trie.ByteStringInternal
import Data.Trie.BitTwiddle

import Data.Binary         (Binary(..), Get, Word8)
#if MIN_VERSION_base(4,9,0)
import Data.Semigroup      (Semigroup(..))
#endif
import Data.Monoid         (Monoid(..))
import Control.Monad       (liftM, liftM3, liftM4)
import Control.Monad       (ap)
import Control.Applicative (Applicative(..), (<$>))
import Data.Foldable       (Foldable(foldMap))
import Data.Traversable    (Traversable(traverse))

#ifdef __GLASGOW_HASKELL__
import GHC.Exts (build)
#endif
------------------------------------------------------------
------------------------------------------------------------


{-----------------------------------------------------------
-- ByteString Big-endian Patricia Trie datatype
-----------------------------------------------------------}
{-
In our idealized representation, we use a (directed) discrete graph
to represent our finite state machine. To organize the set of
outgoing arcs from a given Node we have ArcSet be a big-endian
patricia tree like Data.IntMap. In order to simplify things we then
go through a series of derivations.

data Node a   = Accept a (ArcSet a)
              | Reject   (Branch a)          -- Invariant: Must be Branch
data Arc a    = Arc    ByteString (Node a)   -- Invariant: never empty string
data ArcSet a = None
              | One    {KeyElem} (Arc a)
              | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a)
data Trie a   = Empty
              | Start  ByteString (Node a)   -- Maybe empty string [1]

[1] If we maintain the invariants on how Nodes recurse, then we
can't simply have Start(Node a) because we may have a shared prefix
where the prefix itself is not Accept'ed.


-- Squash Arc into One:
-- (pure good)
data Node a   = Accept a (ArcSet a)
              | Reject   (Branch a)
data ArcSet a = None
              | Arc    ByteString (Node a)
              | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a)
data Trie a   = Empty
              | Start  ByteString (Node a)


-- Squash Node together:
-- (most likely good)
data Node a   = Node (Maybe a) (ArcSet a)
data ArcSet a = None
              | Arc    ByteString (Node a)
              | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a)
data Trie a   = Empty
              | Start  ByteString (Node a)


-- Squash Empty/None and Arc/Start together:
-- (This complicates invariants about non-empty strings and Node's
-- recursion, but those can be circumvented by using smart
-- constructors.)
data Node a = Node (Maybe a) (ArcSet a)
data Trie a = Empty
            | Arc    ByteString (Node a)
            | Branch {Prefix} {Mask} (Trie a) (Trie a)


-- Squash Node into Arc:
-- (By this point, pure good)
-- Unseen invariants:
-- * ByteString non-empty, unless Arc is absolute root of tree
-- * If (Maybe a) is Nothing, then (Trie a) is Branch
--   * With views, we could re-expand Arc into accepting and
--     nonaccepting variants
--
-- [2] Maybe we shouldn't unpack the ByteString. We could specialize
-- or inline the breakMaximalPrefix function to prevent constructing
-- a new ByteString from the parts...
-}
-- | A map from 'ByteString's to @a@. For all the generic functions,
-- note that tries are strict in the @Maybe@ but not in @a@.
--
-- The 'Monad' instance is strange. If a key @k1@ is a prefix of
-- other keys, then results from binding the value at @k1@ will
-- override values from longer keys when they collide. If this is
-- useful for anything, or if there's a more sensible instance, I'd
-- be curious to know.

data Trie a = Empty
            | Arc    {-# UNPACK #-} !ByteString
                                    !(Maybe a)
                                    !(Trie a)
            | Branch {-# UNPACK #-} !Prefix
                     {-# UNPACK #-} !Mask
                                    !(Trie a)
                                    !(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
/= :: Trie a -> Trie a -> Bool
$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
Eq
    -- Prefix/Mask should be deterministic regardless of insertion order
    -- TODO: prove this is so.


-- TODO? add Ord instance like Data.Map?

{-----------------------------------------------------------
-- Trie instances: serialization et cetera
-----------------------------------------------------------}

-- This instance does not unveil the innards of our abstract type.
-- It doesn't emit truly proper Haskell code though, since ByteStrings
-- are printed as (ASCII) Strings, but that's not our fault. (Also
-- 'fromList' is in "Data.Trie" instead of here.)
instance (Show a) => Show (Trie a) where
    showsPrec :: Int -> Trie a -> ShowS
showsPrec Int
p Trie a
t = Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10)
                  (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ (String
"Data.Trie.fromList "String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(ByteString, a)] -> ShowS
forall a. Show a => a -> ShowS
shows ((ByteString -> a -> (ByteString, a)) -> Trie a -> [(ByteString, a)]
forall a b. (ByteString -> a -> b) -> Trie a -> [b]
toListBy (,) Trie a
t)


-- | Visualization fuction for debugging.
showTrie :: (Show a) => Trie a -> String
showTrie :: Trie a -> String
showTrie Trie a
t = ShowS -> Trie a -> ShowS
forall a. Show a => ShowS -> Trie a -> ShowS
shows' ShowS
forall a. a -> a
id Trie a
t String
""
    where
    spaces :: (String -> [b]) -> String
spaces String -> [b]
f = (b -> Char) -> [b] -> String
forall a b. (a -> b) -> [a] -> [b]
map (Char -> b -> Char
forall a b. a -> b -> a
const Char
' ') (String -> [b]
f String
"")

    shows' :: ShowS -> Trie a -> ShowS
shows' ShowS
_  Trie a
Empty            = (String
".\n"String -> ShowS
forall a. [a] -> [a] -> [a]
++)
    shows' ShowS
ss (Branch Prefix
p Prefix
m Trie a
l Trie a
r) =
        let s' :: ShowS
s'  = (String
"--"String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prefix -> ShowS
forall a. Show a => a -> ShowS
shows Prefix
p ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
","String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prefix -> ShowS
forall a. Show a => a -> ShowS
shows Prefix
m ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"-+"String -> ShowS
forall a. [a] -> [a] -> [a]
++)
            ss' :: ShowS
ss' = ShowS
ss ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ShowS
forall a. [a] -> [a]
tail (ShowS -> String
forall b. (String -> [b]) -> String
spaces ShowS
s') String -> ShowS
forall a. [a] -> [a] -> [a]
++)
        in ShowS
s'              ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS -> Trie a -> ShowS
shows' (ShowS
ss' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"|"String -> ShowS
forall a. [a] -> [a] -> [a]
++)) Trie a
l
           ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
ss' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"|\n"String -> ShowS
forall a. [a] -> [a] -> [a]
++)
           ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
ss' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"`"String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS -> Trie a -> ShowS
shows' (ShowS
ss' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
" "String -> ShowS
forall a. [a] -> [a] -> [a]
++)) Trie a
r
    shows' ShowS
ss (Arc ByteString
k Maybe a
mv Trie a
t') =
        let s' :: ShowS
s' = (String
"--"String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ShowS
forall a. Show a => a -> ShowS
shows ByteString
k
                 ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS -> (a -> ShowS) -> Maybe a -> ShowS
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ShowS
forall a. a -> a
id (\a
v -> (String
"-("String -> ShowS
forall a. [a] -> [a] -> [a]
++) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ShowS
forall a. Show a => a -> ShowS
shows a
v ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
")"String -> ShowS
forall a. [a] -> [a] -> [a]
++)) Maybe a
mv
                 ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String
"--"String -> ShowS
forall a. [a] -> [a] -> [a]
++)
        in  ShowS
s' ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS -> Trie a -> ShowS
shows' (ShowS
ss ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ShowS -> String
forall b. (String -> [b]) -> String
spaces ShowS
s' String -> ShowS
forall a. [a] -> [a] -> [a]
++)) Trie a
t'


-- TODO?? a Read instance? hrm... should I?

-- TODO: consider an instance more like the new one for Data.Map. Better?
instance (Binary a) => Binary (Trie a) where
    put :: Trie a -> Put
put Trie a
Empty            = do Prefix -> Put
forall t. Binary t => t -> Put
put (Prefix
0 :: Word8)
    put (Arc ByteString
k Maybe a
m Trie a
t)      = do Prefix -> Put
forall t. Binary t => t -> Put
put (Prefix
1 :: Word8); ByteString -> Put
forall t. Binary t => t -> Put
put ByteString
k; Maybe a -> Put
forall t. Binary t => t -> Put
put Maybe a
m; Trie a -> Put
forall t. Binary t => t -> Put
put Trie a
t
    put (Branch Prefix
p Prefix
m Trie a
l Trie a
r) = do Prefix -> Put
forall t. Binary t => t -> Put
put (Prefix
2 :: Word8); Prefix -> Put
forall t. Binary t => t -> Put
put Prefix
p; Prefix -> Put
forall t. Binary t => t -> Put
put Prefix
m; Trie a -> Put
forall t. Binary t => t -> Put
put Trie a
l; Trie a -> Put
forall t. Binary t => t -> Put
put Trie a
r

    get :: Get (Trie a)
get = do Prefix
tag <- Get Prefix
forall t. Binary t => Get t
get :: Get Word8
             case Prefix
tag of
                 Prefix
0 -> Trie a -> Get (Trie a)
forall (m :: * -> *) a. Monad m => a -> m a
return Trie a
forall a. Trie a
Empty
                 Prefix
1 -> (ByteString -> Maybe a -> Trie a -> Trie a)
-> Get ByteString -> Get (Maybe a) -> Get (Trie a) -> Get (Trie a)
forall (m :: * -> *) a1 a2 a3 r.
Monad m =>
(a1 -> a2 -> a3 -> r) -> m a1 -> m a2 -> m a3 -> m r
liftM3 ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc Get ByteString
forall t. Binary t => Get t
get Get (Maybe a)
forall t. Binary t => Get t
get Get (Trie a)
forall t. Binary t => Get t
get
                 Prefix
_ -> (Prefix -> Prefix -> Trie a -> Trie a -> Trie a)
-> Get Prefix
-> Get Prefix
-> Get (Trie a)
-> Get (Trie a)
-> Get (Trie a)
forall (m :: * -> *) a1 a2 a3 a4 r.
Monad m =>
(a1 -> a2 -> a3 -> a4 -> r) -> m a1 -> m a2 -> m a3 -> m a4 -> m r
liftM4 Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Get Prefix
forall t. Binary t => Get t
get Get Prefix
forall t. Binary t => Get t
get Get (Trie a)
forall t. Binary t => Get t
get Get (Trie a)
forall t. Binary t => Get t
get


{-----------------------------------------------------------
-- Trie instances: Abstract Nonsense
-----------------------------------------------------------}

instance Functor Trie where
    fmap :: (a -> b) -> Trie a -> Trie b
fmap a -> b
f = Trie a -> Trie b
go
        where
        go :: Trie a -> Trie b
go Trie a
Empty              = Trie b
forall a. Trie a
Empty
        go (Arc ByteString
k Maybe a
Nothing  Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k Maybe b
forall a. Maybe a
Nothing      (Trie a -> Trie b
go Trie a
t)
        go (Arc ByteString
k (Just a
v) Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k (b -> Maybe b
forall a. a -> Maybe a
Just (a -> b
f a
v)) (Trie a -> Trie b
go Trie a
t)
        go (Branch Prefix
p Prefix
m Trie a
l Trie a
r)   = Prefix -> Prefix -> Trie b -> Trie b -> Trie b
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p Prefix
m (Trie a -> Trie b
go Trie a
l) (Trie a -> Trie b
go Trie a
r)


instance Foldable Trie where
    -- If our definition of foldr is so much faster than the Endo
    -- default, then maybe we should remove this and use the default
    -- foldMap based on foldr
    foldMap :: (a -> m) -> Trie a -> m
foldMap a -> m
f = Trie a -> m
go
        where
        go :: Trie a -> m
go Trie a
Empty              = m
forall a. Monoid a => a
mempty
        go (Arc ByteString
_ Maybe a
Nothing  Trie a
t) = Trie a -> m
go Trie a
t
        go (Arc ByteString
_ (Just a
v) Trie a
t) = a -> m
f a
v m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` Trie a -> m
go Trie a
t
        go (Branch Prefix
_ Prefix
_ Trie a
l Trie a
r)   = Trie a -> m
go Trie a
l m -> m -> m
forall a. Monoid a => a -> a -> a
`mappend` Trie a -> m
go Trie a
r

    {- This definition is much faster, but it's also wrong
    -- (or at least different than foldrWithKey)
    foldr f = \z t -> go t id z
        where
        go Empty              k x = k x
        go (Branch _ _ l r)   k x = go r (go l k) x
        go (Arc _ Nothing t)  k x = go t k x
        go (Arc _ (Just v) t) k x = go t k (f v x)

    foldl f = \z t -> go t id z
        where
        go Empty              k x = k x
        go (Branch _ _ l r)   k x = go l (go r k) x
        go (Arc _ Nothing t)  k x = go t k x
        go (Arc _ (Just v) t) k x = go t k (f x v)
    -}

-- TODO: newtype Keys = K Trie  ; instance Foldable Keys
-- TODO: newtype Assoc = A Trie ; instance Foldable Assoc

instance Traversable Trie where
    traverse :: (a -> f b) -> Trie a -> f (Trie b)
traverse a -> f b
f = Trie a -> f (Trie b)
go
        where
        go :: Trie a -> f (Trie b)
go Trie a
Empty              = Trie b -> f (Trie b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Trie b
forall a. Trie a
Empty
        go (Arc ByteString
k Maybe a
Nothing  Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k Maybe b
forall a. Maybe a
Nothing        (Trie b -> Trie b) -> f (Trie b) -> f (Trie b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Trie a -> f (Trie b)
go Trie a
t
        go (Arc ByteString
k (Just a
v) Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k (Maybe b -> Trie b -> Trie b)
-> (b -> Maybe b) -> b -> Trie b -> Trie b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Maybe b
forall a. a -> Maybe a
Just (b -> Trie b -> Trie b) -> f b -> f (Trie b -> Trie b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
v f (Trie b -> Trie b) -> f (Trie b) -> f (Trie b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Trie a -> f (Trie b)
go Trie a
t
        go (Branch Prefix
p Prefix
m Trie a
l Trie a
r)   = Prefix -> Prefix -> Trie b -> Trie b -> Trie b
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p Prefix
m (Trie b -> Trie b -> Trie b) -> f (Trie b) -> f (Trie b -> Trie b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Trie a -> f (Trie b)
go Trie a
l f (Trie b -> Trie b) -> f (Trie b) -> f (Trie b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Trie a -> f (Trie b)
go Trie a
r

instance Applicative Trie where
    pure :: a -> Trie a
pure  = a -> Trie a
forall (m :: * -> *) a. Monad m => a -> m a
return
    <*> :: Trie (a -> b) -> Trie a -> Trie b
(<*>) = Trie (a -> b) -> Trie a -> Trie b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

-- Does this even make sense? It's not nondeterminism like lists
-- and sets. If no keys were prefixes of other keys it'd make sense
-- as a decision-tree; but since keys /can/ prefix, tries formed
-- from shorter keys can shadow the results from longer keys due
-- to the 'unionL'. It does seem to follow the laws though... What
-- computation could this possibly represent?
--
--  1. return x >>= f  == f x
--  2. m >>= return    == m
--  3. (m >>= f) >>= g == m >>= (\x -> f x >>= g)
instance Monad Trie where
    return :: a -> Trie a
return = ByteString -> a -> Trie a
forall a. ByteString -> a -> Trie a
singleton ByteString
S.empty

    >>= :: Trie a -> (a -> Trie b) -> Trie b
(>>=) Trie a
Empty              a -> Trie b
_ = Trie b
forall a. Trie a
empty
    (>>=) (Branch Prefix
p Prefix
m Trie a
l Trie a
r)   a -> Trie b
f = Prefix -> Prefix -> Trie b -> Trie b -> Trie b
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p Prefix
m (Trie a
l Trie a -> (a -> Trie b) -> Trie b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Trie b
f) (Trie a
r Trie a -> (a -> Trie b) -> Trie b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Trie b
f)
    (>>=) (Arc ByteString
k Maybe a
Nothing  Trie a
t) a -> Trie b
f = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k Maybe b
forall a. Maybe a
Nothing (Trie a
t Trie a -> (a -> Trie b) -> Trie b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Trie b
f)
    (>>=) (Arc ByteString
k (Just a
v) Trie a
t) a -> Trie b
f = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k Maybe b
forall a. Maybe a
Nothing (a -> Trie b
f a
v Trie b -> Trie b -> Trie b
forall a. Trie a -> Trie a -> Trie a
`unionL` (Trie a
t Trie a -> (a -> Trie b) -> Trie b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> Trie b
f))
                               where
                               unionL :: Trie a -> Trie a -> Trie a
unionL = (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
mergeBy (\a
x a
_ -> a -> Maybe a
forall a. a -> Maybe a
Just a
x)


#if MIN_VERSION_base(4,9,0)
-- The "Data.Semigroup" module is in base since 4.9.0.0; but having
-- the 'Semigroup' superclass for the 'Monoid' instance only comes
-- into force in base 4.11.0.0.
instance (Semigroup a) => Semigroup (Trie a) where
    <> :: Trie a -> Trie a -> Trie a
(<>) = (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
mergeBy ((a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a)
-> (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ \a
x a
y -> a -> Maybe a
forall a. a -> Maybe a
Just (a
x a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
y)
    -- TODO: optimized implementations of:
    -- sconcat :: NonEmpty a -> a
    -- stimes :: Integral b => b -> a -> a
#endif

-- This instance is more sensible than Data.IntMap and Data.Map's
instance (Monoid a) => Monoid (Trie a) where
    mempty :: Trie a
mempty  = Trie a
forall a. Trie a
empty
    mappend :: Trie a -> Trie a -> Trie a
mappend = (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
forall a. (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
mergeBy ((a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a)
-> (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
forall a b. (a -> b) -> a -> b
$ \a
x a
y -> a -> Maybe a
forall a. a -> Maybe a
Just (a
x a -> a -> a
forall a. Monoid a => a -> a -> a
`mappend` a
y)


-- Since the Monoid instance isn't natural in @a@, I can't think
-- of any other sensible instance for MonadPlus. It's as specious
-- as Maybe, IO, and STM's instances though.
--
-- MonadPlus laws: <http://www.haskell.org/haskellwiki/MonadPlus>
--  1. <Trie a, mzero, mplus> forms a monoid
--  2. mzero >>= f        === mzero
--  3. m >> mzero         === mzero
--  4. mplus m n >>= k    === mplus (m >>= k) (n >>= k)
--  4' mplus (return a) n === return a
{-
-- Follows #1, #1, and #3. But it does something like 4' instead
-- of actually doing #4 (since we'd merge the trees generated by
-- @k@ for conflicting values)
--
-- TODO: cf Control.Applicative.Alternative (base-4, but not Hugs).
-- But (<*>) gets odd when the function is not 'pure'... maybe
-- helpful though.
instance MonadPlus Trie where
    mzero = empty
    mplus = unionL where unionL = mergeBy (\x _ -> Just x)
-}


{-----------------------------------------------------------
-- Extra mapping functions
-----------------------------------------------------------}

-- | Apply a function to all values, potentially removing them.
filterMap :: (a -> Maybe b) -> Trie a -> Trie b
filterMap :: (a -> Maybe b) -> Trie a -> Trie b
filterMap a -> Maybe b
f = Trie a -> Trie b
go
    where
    go :: Trie a -> Trie b
go Trie a
Empty              = Trie b
forall a. Trie a
empty
    go (Arc ByteString
k Maybe a
Nothing  Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k Maybe b
forall a. Maybe a
Nothing (Trie a -> Trie b
go Trie a
t)
    go (Arc ByteString
k (Just a
v) Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k (a -> Maybe b
f a
v)   (Trie a -> Trie b
go Trie a
t)
    go (Branch Prefix
p Prefix
m Trie a
l Trie a
r)   = Prefix -> Prefix -> Trie b -> Trie b -> Trie b
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p Prefix
m (Trie a -> Trie b
go Trie a
l) (Trie a -> Trie b
go Trie a
r)


-- | Generic version of 'fmap'. This function is notably more
-- expensive than 'fmap' or 'filterMap' because we have to reconstruct
-- the keys.
mapBy :: (ByteString -> a -> Maybe b) -> Trie a -> Trie b
mapBy :: (ByteString -> a -> Maybe b) -> Trie a -> Trie b
mapBy ByteString -> a -> Maybe b
f = ByteString -> Trie a -> Trie b
go ByteString
S.empty
    where
    go :: ByteString -> Trie a -> Trie b
go ByteString
_ Trie a
Empty              = Trie b
forall a. Trie a
empty
    go ByteString
q (Arc ByteString
k Maybe a
Nothing  Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k Maybe b
forall a. Maybe a
Nothing  (ByteString -> Trie a -> Trie b
go ByteString
q' Trie a
t) where q' :: ByteString
q' = ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k
    go ByteString
q (Arc ByteString
k (Just a
v) Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k (ByteString -> a -> Maybe b
f ByteString
q' a
v) (ByteString -> Trie a -> Trie b
go ByteString
q' Trie a
t) where q' :: ByteString
q' = ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k
    go ByteString
q (Branch Prefix
p Prefix
m Trie a
l Trie a
r)   = Prefix -> Prefix -> Trie b -> Trie b -> Trie b
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p Prefix
m (ByteString -> Trie a -> Trie b
go ByteString
q Trie a
l) (ByteString -> Trie a -> Trie b
go ByteString
q Trie a
r)


-- | A variant of 'fmap' which provides access to the subtrie rooted
-- at each value.
contextualMap :: (a -> Trie a -> b) -> Trie a -> Trie b
contextualMap :: (a -> Trie a -> b) -> Trie a -> Trie b
contextualMap a -> Trie a -> b
f = Trie a -> Trie b
go
    where
    go :: Trie a -> Trie b
go Trie a
Empty              = Trie b
forall a. Trie a
Empty
    go (Arc ByteString
k Maybe a
Nothing  Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k Maybe b
forall a. Maybe a
Nothing        (Trie a -> Trie b
go Trie a
t)
    go (Arc ByteString
k (Just a
v) Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k (b -> Maybe b
forall a. a -> Maybe a
Just (a -> Trie a -> b
f a
v Trie a
t)) (Trie a -> Trie b
go Trie a
t)
    go (Branch Prefix
p Prefix
m Trie a
l Trie a
r)   = Prefix -> Prefix -> Trie b -> Trie b -> Trie b
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p Prefix
m (Trie a -> Trie b
go Trie a
l) (Trie a -> Trie b
go Trie a
r)


-- | A variant of 'contextualMap' which applies the function strictly.
contextualMap' :: (a -> Trie a -> b) -> Trie a -> Trie b
contextualMap' :: (a -> Trie a -> b) -> Trie a -> Trie b
contextualMap' a -> Trie a -> b
f = Trie a -> Trie b
go
    where
    go :: Trie a -> Trie b
go Trie a
Empty              = Trie b
forall a. Trie a
Empty
    go (Arc ByteString
k Maybe a
Nothing  Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k Maybe b
forall a. Maybe a
Nothing         (Trie a -> Trie b
go Trie a
t)
    go (Arc ByteString
k (Just a
v) Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k (b -> Maybe b
forall a. a -> Maybe a
Just (b -> Maybe b) -> b -> Maybe b
forall a b. (a -> b) -> a -> b
$! a -> Trie a -> b
f a
v Trie a
t) (Trie a -> Trie b
go Trie a
t)
    go (Branch Prefix
p Prefix
m Trie a
l Trie a
r)   = Prefix -> Prefix -> Trie b -> Trie b -> Trie b
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p Prefix
m (Trie a -> Trie b
go Trie a
l) (Trie a -> Trie b
go Trie a
r)


-- | A contextual variant of 'filterMap'.
contextualFilterMap :: (a -> Trie a -> Maybe b) -> Trie a -> Trie b
contextualFilterMap :: (a -> Trie a -> Maybe b) -> Trie a -> Trie b
contextualFilterMap a -> Trie a -> Maybe b
f = Trie a -> Trie b
go
    where
    go :: Trie a -> Trie b
go Trie a
Empty              = Trie b
forall a. Trie a
empty
    go (Arc ByteString
k Maybe a
Nothing  Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k Maybe b
forall a. Maybe a
Nothing (Trie a -> Trie b
go Trie a
t)
    go (Arc ByteString
k (Just a
v) Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k (a -> Trie a -> Maybe b
f a
v Trie a
t) (Trie a -> Trie b
go Trie a
t)
    go (Branch Prefix
p Prefix
m Trie a
l Trie a
r)   = Prefix -> Prefix -> Trie b -> Trie b -> Trie b
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p Prefix
m (Trie a -> Trie b
go Trie a
l) (Trie a -> Trie b
go Trie a
r)


-- | A contextual variant of 'mapBy'. Again note that this is
-- expensive since we must reconstruct the keys.
contextualMapBy :: (ByteString -> a -> Trie a -> Maybe b) -> Trie a -> Trie b
contextualMapBy :: (ByteString -> a -> Trie a -> Maybe b) -> Trie a -> Trie b
contextualMapBy ByteString -> a -> Trie a -> Maybe b
f = ByteString -> Trie a -> Trie b
go ByteString
S.empty
    where
    go :: ByteString -> Trie a -> Trie b
go ByteString
_ Trie a
Empty              = Trie b
forall a. Trie a
empty
    go ByteString
q (Arc ByteString
k Maybe a
Nothing  Trie a
t) = ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k Maybe b
forall a. Maybe a
Nothing (ByteString -> Trie a -> Trie b
go (ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k) Trie a
t)
    go ByteString
q (Arc ByteString
k (Just a
v) Trie a
t) = let q' :: ByteString
q' = ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k
                              in ByteString -> Maybe b -> Trie b -> Trie b
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k (ByteString -> a -> Trie a -> Maybe b
f ByteString
q' a
v Trie a
t) (ByteString -> Trie a -> Trie b
go ByteString
q' Trie a
t)
    go ByteString
q (Branch Prefix
p Prefix
m Trie a
l Trie a
r)   = Prefix -> Prefix -> Trie b -> Trie b -> Trie b
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p Prefix
m (ByteString -> Trie a -> Trie b
go ByteString
q Trie a
l) (ByteString -> Trie a -> Trie b
go ByteString
q Trie a
r)


{-----------------------------------------------------------
-- Smart constructors and helper functions for building tries
-----------------------------------------------------------}

-- | Smart constructor to prune @Empty@ from @Branch@es.
branch :: Prefix -> Mask -> Trie a -> Trie a -> Trie a
{-# INLINE branch #-}
branch :: Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
_ Prefix
_ Trie a
Empty Trie a
r     = Trie a
r
branch Prefix
_ Prefix
_ Trie a
l     Trie a
Empty = Trie a
l
branch Prefix
p Prefix
m Trie a
l     Trie a
r     = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p Prefix
m Trie a
l Trie a
r


-- | Smart constructor to prune @Arc@s that lead nowhere.
-- N.B if mv=Just then doesn't check whether t=epsilon. It's up to callers to ensure that invariant isn't broken.
arc :: ByteString -> Maybe a -> Trie a -> Trie a
{-# INLINE arc #-}
arc :: ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k mv :: Maybe a
mv@(Just a
_)   Trie a
t                            = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k Maybe a
mv Trie a
t
arc ByteString
_    Maybe a
Nothing    Trie a
Empty                        = Trie a
forall a. Trie a
Empty
arc ByteString
k    Maybe a
Nothing  t :: Trie a
t@(Branch Prefix
_ Prefix
_ Trie a
_ Trie a
_) | ByteString -> Bool
S.null ByteString
k  = Trie a
t
                                     | Bool
otherwise = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k Maybe a
forall a. Maybe a
Nothing Trie a
t
arc ByteString
k    Maybe a
Nothing    (Arc ByteString
k' Maybe a
mv' Trie a
t')              = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc (ByteString -> ByteString -> ByteString
S.append ByteString
k ByteString
k') Maybe a
mv' Trie a
t'


-- | Smart constructor to join two tries into a @Branch@ with maximal
-- prefix sharing. Requires knowing the prefixes, but can combine
-- either @Branch@es or @Arc@s.
--
-- N.B. /do not/ use if prefixes could match entirely!
branchMerge :: Prefix -> Trie a -> Prefix -> Trie a -> Trie a
{-# INLINE branchMerge #-}
branchMerge :: Prefix -> Trie a -> Prefix -> Trie a -> Trie a
branchMerge Prefix
_ Trie a
Empty Prefix
_ Trie a
t2    = Trie a
t2
branchMerge Prefix
_  Trie a
t1   Prefix
_ Trie a
Empty = Trie a
t1
branchMerge Prefix
p1 Trie a
t1  Prefix
p2 Trie a
t2
    | Prefix -> Prefix -> Bool
zero Prefix
p1 Prefix
m             = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p Prefix
m Trie a
t1 Trie a
t2
    | Bool
otherwise             = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p Prefix
m Trie a
t2 Trie a
t1
    where
    m :: Prefix
m = Prefix -> Prefix -> Prefix
branchMask Prefix
p1 Prefix
p2
    p :: Prefix
p = Prefix -> Prefix -> Prefix
mask Prefix
p1 Prefix
m


-- It would be better if Arc used
-- Data.ByteString.TrieInternal.wordHead somehow, that way
-- we can see 4/8/?*Word8 at a time instead of just one.
-- But that makes maintaining invariants ...difficult :(
getPrefix :: Trie a -> Prefix
{-# INLINE getPrefix #-}
getPrefix :: Trie a -> Prefix
getPrefix (Branch Prefix
p Prefix
_ Trie a
_ Trie a
_)        = Prefix
p
getPrefix (Arc ByteString
k Maybe a
_ Trie a
_) | ByteString -> Bool
S.null ByteString
k  = Prefix
0 -- for lack of a better value
                      | Bool
otherwise = ByteString -> Prefix
S.head ByteString
k
getPrefix Trie a
Empty                   = String -> Prefix
forall a. HasCallStack => String -> a
error String
"getPrefix: no Prefix of Empty"


{-----------------------------------------------------------
-- Error messages
-----------------------------------------------------------}

-- TODO: shouldn't we inline the logic and just NOINLINE the string constant? There are only three use sites, which themselves aren't inlined...
errorLogHead :: String -> ByteString -> ByteStringElem
{-# NOINLINE errorLogHead #-}
errorLogHead :: String -> ByteString -> Prefix
errorLogHead String
fn ByteString
q
    | ByteString -> Bool
S.null ByteString
q  = String -> Prefix
forall a. HasCallStack => String -> a
error (String -> Prefix) -> String -> Prefix
forall a b. (a -> b) -> a -> b
$ String
"Data.Trie.Internal." String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
fn String -> ShowS
forall a. [a] -> [a] -> [a]
++String
": found null subquery"
    | Bool
otherwise = ByteString -> Prefix
S.head ByteString
q


------------------------------------------------------------
------------------------------------------------------------

{-----------------------------------------------------------
-- Basic functions
-----------------------------------------------------------}

-- | /O(1)/, Construct the empty trie.
empty :: Trie a
{-# INLINE empty #-}
empty :: Trie a
empty = Trie a
forall a. Trie a
Empty


-- | /O(1)/, Is the trie empty?
null :: Trie a -> Bool
{-# INLINE null #-}
null :: Trie a -> Bool
null Trie a
Empty = Bool
True
null Trie a
_     = Bool
False


-- | /O(1)/, Construct a singleton trie.
singleton :: ByteString -> a -> Trie a
{-# INLINE singleton #-}
singleton :: ByteString -> a -> Trie a
singleton ByteString
k a
v = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k (a -> Maybe a
forall a. a -> Maybe a
Just a
v) Trie a
forall a. Trie a
Empty
-- For singletons, don't need to verify invariant on arc length >0


-- | /O(n)/, Get count of elements in trie.
size  :: Trie a -> Int
{-# INLINE size #-}
size :: Trie a -> Int
size Trie a
t = Trie a -> (Int -> Int) -> Int -> Int
forall a. Trie a -> (Int -> Int) -> Int -> Int
size' Trie a
t Int -> Int
forall a. a -> a
id Int
0

-- | /O(n)/, CPS accumulator helper for calculating 'size'.
size' :: Trie a -> (Int -> Int) -> Int -> Int
size' :: Trie a -> (Int -> Int) -> Int -> Int
size' Trie a
Empty              Int -> Int
f Int
n = Int -> Int
f Int
n
size' (Branch Prefix
_ Prefix
_ Trie a
l Trie a
r)   Int -> Int
f Int
n = Trie a -> (Int -> Int) -> Int -> Int
forall a. Trie a -> (Int -> Int) -> Int -> Int
size' Trie a
l (Trie a -> (Int -> Int) -> Int -> Int
forall a. Trie a -> (Int -> Int) -> Int -> Int
size' Trie a
r Int -> Int
f) Int
n
size' (Arc ByteString
_ Maybe a
Nothing Trie a
t)  Int -> Int
f Int
n = Trie a -> (Int -> Int) -> Int -> Int
forall a. Trie a -> (Int -> Int) -> Int -> Int
size' Trie a
t Int -> Int
f Int
n
size' (Arc ByteString
_ (Just a
_) Trie a
t) Int -> Int
f Int
n = Trie a -> (Int -> Int) -> Int -> Int
forall a. Trie a -> (Int -> Int) -> Int -> Int
size' Trie a
t Int -> Int
f (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$! Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1


{-----------------------------------------------------------
-- Conversion functions
-----------------------------------------------------------}

-- Still rather inefficient
--
-- TODO: rewrite list-catenation to be lazier (real CPS instead of
-- function building? is the function building really better than
-- (++) anyways?)
-- N.B. If our manual definition of foldr/foldl (using function
-- application) is so much faster than the default Endo definition
-- (using function composition), then we should make this use
-- application instead too.
--
-- TODO: the @q@ accumulator should be lazy ByteString and only
-- forced by @fcons@. It's already non-strict, but we should ensure
-- O(n) not O(n^2) when it's forced.
--
-- BUG: not safe for deep strict @fcons@, only for WHNF-strict like (:)
-- Where to put the strictness to amortize it?
--
-- | Convert a trie into a list (in key-sorted order) using a
-- function, folding the list as we go.
foldrWithKey :: (ByteString -> a -> b -> b) -> b -> Trie a -> b
foldrWithKey :: (ByteString -> a -> b -> b) -> b -> Trie a -> b
foldrWithKey ByteString -> a -> b -> b
fcons b
nil = \Trie a
t -> ByteString -> Trie a -> b -> b
go ByteString
S.empty Trie a
t b
nil
    where
    go :: ByteString -> Trie a -> b -> b
go ByteString
_ Trie a
Empty            = b -> b
forall a. a -> a
id
    go ByteString
q (Branch Prefix
_ Prefix
_ Trie a
l Trie a
r) = ByteString -> Trie a -> b -> b
go ByteString
q Trie a
l (b -> b) -> (b -> b) -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Trie a -> b -> b
go ByteString
q Trie a
r
    go ByteString
q (Arc ByteString
k Maybe a
mv Trie a
t)     =
        case Maybe a
mv of
        Maybe a
Nothing -> b -> b
rest
        Just a
v  -> ByteString -> a -> b -> b
fcons ByteString
k' a
v (b -> b) -> (b -> b) -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> b
rest
        where
        rest :: b -> b
rest = ByteString -> Trie a -> b -> b
go ByteString
k' Trie a
t
        k' :: ByteString
k'   = ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k


-- cf Data.ByteString.unpack
-- <http://hackage.haskell.org/packages/archive/bytestring/0.9.1.4/doc/html/src/Data-ByteString.html>
--
-- | Convert a trie into a list using a function. Resulting values
-- are in key-sorted order.
toListBy :: (ByteString -> a -> b) -> Trie a -> [b]
{-# INLINE toListBy #-}
#if !defined(__GLASGOW_HASKELL__)
-- TODO: should probably inline foldrWithKey
-- TODO: compare performance of that vs both this and the GHC version
toListBy f t = foldrWithKey (((:) .) . f) [] t
#else
-- Written with 'build' to enable the build\/foldr fusion rules.
toListBy :: (ByteString -> a -> b) -> Trie a -> [b]
toListBy ByteString -> a -> b
f Trie a
t = (forall b. (b -> b -> b) -> b -> b) -> [b]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build ((ByteString -> a -> b) -> Trie a -> (b -> b -> b) -> b -> b
forall a b c.
(ByteString -> a -> b) -> Trie a -> (b -> c -> c) -> c -> c
toListByFB ByteString -> a -> b
f Trie a
t)

-- TODO: should probably have a specialized version for strictness,
-- and a rule to rewrite generic lazy version into it. As per
-- Data.ByteString.unpack and the comments there about strictness
-- and fusion.
toListByFB :: (ByteString -> a -> b) -> Trie a -> (b -> c -> c) -> c -> c
{-# INLINE [0] toListByFB #-}
toListByFB :: (ByteString -> a -> b) -> Trie a -> (b -> c -> c) -> c -> c
toListByFB ByteString -> a -> b
f Trie a
t b -> c -> c
cons c
nil = (ByteString -> a -> c -> c) -> c -> Trie a -> c
forall a b. (ByteString -> a -> b -> b) -> b -> Trie a -> b
foldrWithKey ((b -> c -> c
cons (b -> c -> c) -> (a -> b) -> a -> c -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((a -> b) -> a -> c -> c)
-> (ByteString -> a -> b) -> ByteString -> a -> c -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> a -> b
f) c
nil Trie a
t
#endif


{-----------------------------------------------------------
-- Query functions (just recurse)
-----------------------------------------------------------}

-- | Generic function to find a value (if it exists) and the subtrie
-- rooted at the prefix. The first function argument is called if and
-- only if a node is exactly reachable by the query; if no node is
-- exactly reachable the default value is used; if the middle of
-- an arc is reached, the second function argument is used.
--
-- This function is intended for internal use. For the public-facing
-- version, see @lookupBy@ in "Data.Trie".
lookupBy_ :: (Maybe a -> Trie a -> b) -> b -> (Trie a -> b)
          -> ByteString -> Trie a -> b
lookupBy_ :: (Maybe a -> Trie a -> b)
-> b -> (Trie a -> b) -> ByteString -> Trie a -> b
lookupBy_ Maybe a -> Trie a -> b
f b
z Trie a -> b
a = ByteString -> Trie a -> b
lookupBy_'
    where
    -- | Deal with epsilon query (when there is no epsilon value)
    lookupBy_' :: ByteString -> Trie a -> b
lookupBy_' ByteString
q t :: Trie a
t@(Branch Prefix
_ Prefix
_ Trie a
_ Trie a
_) | ByteString -> Bool
S.null ByteString
q = Maybe a -> Trie a -> b
f Maybe a
forall a. Maybe a
Nothing Trie a
t
    lookupBy_' ByteString
q Trie a
t                             = ByteString -> Trie a -> b
go ByteString
q Trie a
t

    -- | The main recursion
    go :: ByteString -> Trie a -> b
go ByteString
_    Trie a
Empty       = b
z

    go ByteString
q   (Arc ByteString
k Maybe a
mv Trie a
t) =
        let (ByteString
_,ByteString
k',ByteString
q')   = ByteString -> ByteString -> (ByteString, ByteString, ByteString)
breakMaximalPrefix ByteString
k ByteString
q
        in case (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ByteString -> Bool
S.null ByteString
k', ByteString -> Bool
S.null ByteString
q') of
                (Bool
True,  Bool
True)  -> Trie a -> b
a (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k' Maybe a
mv Trie a
t)
                (Bool
True,  Bool
False) -> b
z
                (Bool
False, Bool
True)  -> Maybe a -> Trie a -> b
f Maybe a
mv Trie a
t
                (Bool
False, Bool
False) -> ByteString -> Trie a -> b
go ByteString
q' Trie a
t

    go ByteString
q t_ :: Trie a
t_@(Branch Prefix
_ Prefix
_ Trie a
_ Trie a
_) = Trie a -> b
findArc Trie a
t_
        where
        qh :: Prefix
qh = String -> ByteString -> Prefix
errorLogHead String
"lookupBy_" ByteString
q

        -- | /O(min(m,W))/, where /m/ is number of @Arc@s in this
        -- branching, and /W/ is the word size of the Prefix,Mask type.
        findArc :: Trie a -> b
findArc (Branch Prefix
p Prefix
m Trie a
l Trie a
r)
            | Prefix -> Prefix -> Prefix -> Bool
nomatch Prefix
qh Prefix
p Prefix
m  = b
z
            | Prefix -> Prefix -> Bool
zero Prefix
qh Prefix
m       = Trie a -> b
findArc Trie a
l
            | Bool
otherwise       = Trie a -> b
findArc Trie a
r
        findArc t :: Trie a
t@(Arc ByteString
_ Maybe a
_ Trie a
_) = ByteString -> Trie a -> b
go ByteString
q Trie a
t
        findArc Trie a
Empty         = b
z


-- This function needs to be here, not in "Data.Trie", because of
-- 'arc' which isn't exported. We could use the monad instance
-- instead, though it'd be far more circuitous.
--     arc k Nothing  t === singleton k () >> t
--     arc k (Just v) t === singleton k v  >>= unionR t . singleton S.empty
--         (...except 'arc' doesn't do the invariant correction
--           of (>>=) for epsilon`elem`t)
--
-- | Return the subtrie containing all keys beginning with a prefix.
submap :: ByteString -> Trie a -> Trie a
{-# INLINE submap #-}
submap :: ByteString -> Trie a -> Trie a
submap ByteString
q = (Maybe a -> Trie a -> Trie a)
-> Trie a -> (Trie a -> Trie a) -> ByteString -> Trie a -> Trie a
forall a b.
(Maybe a -> Trie a -> b)
-> b -> (Trie a -> b) -> ByteString -> Trie a -> b
lookupBy_ (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
q) Trie a
forall a. Trie a
empty (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
q Maybe a
forall a. Maybe a
Nothing) ByteString
q
{-  -- Disable superfluous error checking.
    -- @submap'@ would replace the first argument to @lookupBy_@
    where
    submap' Nothing Empty       = errorEmptyAfterNothing "submap"
    submap' Nothing (Arc _ _ _) = errorArcAfterNothing   "submap"
    submap' mx      t           = Arc q mx t

errorInvariantBroken :: String -> String -> a
{-# NOINLINE errorInvariantBroken #-}
errorInvariantBroken s e =  error (s ++ ": Invariant was broken" ++ e')
    where
    e' = if Prelude.null e then e else ", found: " ++ e

errorArcAfterNothing    :: String -> a
{-# NOINLINE errorArcAfterNothing #-}
errorArcAfterNothing   s = errorInvariantBroken s "Arc after Nothing"

errorEmptyAfterNothing  :: String -> a
{-# NOINLINE errorEmptyAfterNothing #-}
errorEmptyAfterNothing s = errorInvariantBroken s "Empty after Nothing"
-- -}



-- TODO: would it be worth it to have a variant like 'lookupBy_' which takes the three continuations?

-- | Given a query, find the longest prefix with an associated value
-- in the trie, returning the length of that prefix and the associated
-- value.
--
-- This function may not have the most useful return type. For a
-- version that returns the prefix itself as well as the remaining
-- string, see @match@ in "Data.Trie".
match_ :: Trie a -> ByteString -> Maybe (Int, a)
match_ :: Trie a -> ByteString -> Maybe (Int, a)
match_ = (ByteString -> Trie a -> Maybe (Int, a))
-> Trie a -> ByteString -> Maybe (Int, a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip ByteString -> Trie a -> Maybe (Int, a)
forall b. ByteString -> Trie b -> Maybe (Int, b)
start
    where
    -- | Deal with epsilon query (when there is no epsilon value)
    start :: ByteString -> Trie b -> Maybe (Int, b)
start ByteString
q (Branch Prefix
_ Prefix
_ Trie b
_ Trie b
_) | ByteString -> Bool
S.null ByteString
q = Maybe (Int, b)
forall a. Maybe a
Nothing
    start ByteString
q Trie b
t                           = Int -> ByteString -> Trie b -> Maybe (Int, b)
forall b. Int -> ByteString -> Trie b -> Maybe (Int, b)
goNothing Int
0 ByteString
q Trie b
t

    -- | The initial recursion
    goNothing :: Int -> ByteString -> Trie b -> Maybe (Int, b)
goNothing Int
_ ByteString
_    Trie b
Empty       = Maybe (Int, b)
forall a. Maybe a
Nothing

    goNothing Int
n ByteString
q   (Arc ByteString
k Maybe b
mv Trie b
t) =
        let (ByteString
p,ByteString
k',ByteString
q') = ByteString -> ByteString -> (ByteString, ByteString, ByteString)
breakMaximalPrefix ByteString
k ByteString
q
            n' :: Int
n'        = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ByteString -> Int
S.length ByteString
p
        in Int
n' Int -> Maybe (Int, b) -> Maybe (Int, b)
`seq`
            if ByteString -> Bool
S.null ByteString
k'
            then
                if ByteString -> Bool
S.null ByteString
q'
                then (,) Int
n' (b -> (Int, b)) -> Maybe b -> Maybe (Int, b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe b
mv
                else
                    case Maybe b
mv of
                    Maybe b
Nothing -> Int -> ByteString -> Trie b -> Maybe (Int, b)
goNothing   Int
n' ByteString
q' Trie b
t
                    Just b
v  -> Int -> b -> Int -> ByteString -> Trie b -> Maybe (Int, b)
forall b. Int -> b -> Int -> ByteString -> Trie b -> Maybe (Int, b)
goJust Int
n' b
v Int
n' ByteString
q' Trie b
t
            else Maybe (Int, b)
forall a. Maybe a
Nothing

    goNothing Int
n ByteString
q t_ :: Trie b
t_@(Branch Prefix
_ Prefix
_ Trie b
_ Trie b
_) = Trie b -> Maybe (Int, b)
findArc Trie b
t_
        where
        qh :: Prefix
qh = String -> ByteString -> Prefix
errorLogHead String
"match_" ByteString
q

        -- | /O(min(m,W))/, where /m/ is number of @Arc@s in this
        -- branching, and /W/ is the word size of the Prefix,Mask type.
        findArc :: Trie b -> Maybe (Int, b)
findArc (Branch Prefix
p Prefix
m Trie b
l Trie b
r)
            | Prefix -> Prefix -> Prefix -> Bool
nomatch Prefix
qh Prefix
p Prefix
m  = Maybe (Int, b)
forall a. Maybe a
Nothing
            | Prefix -> Prefix -> Bool
zero Prefix
qh Prefix
m       = Trie b -> Maybe (Int, b)
findArc Trie b
l
            | Bool
otherwise       = Trie b -> Maybe (Int, b)
findArc Trie b
r
        findArc t :: Trie b
t@(Arc ByteString
_ Maybe b
_ Trie b
_) = Int -> ByteString -> Trie b -> Maybe (Int, b)
goNothing Int
n ByteString
q Trie b
t
        findArc Trie b
Empty         = Maybe (Int, b)
forall a. Maybe a
Nothing

    -- | The main recursion
    goJust :: Int -> b -> Int -> ByteString -> Trie b -> Maybe (Int, b)
goJust Int
n0 b
v0 Int
_ ByteString
_    Trie b
Empty       = (Int, b) -> Maybe (Int, b)
forall a. a -> Maybe a
Just (Int
n0,b
v0)

    goJust Int
n0 b
v0 Int
n ByteString
q   (Arc ByteString
k Maybe b
mv Trie b
t) =
        let (ByteString
p,ByteString
k',ByteString
q') = ByteString -> ByteString -> (ByteString, ByteString, ByteString)
breakMaximalPrefix ByteString
k ByteString
q
            n' :: Int
n'        = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ByteString -> Int
S.length ByteString
p
        in Int
n' Int -> Maybe (Int, b) -> Maybe (Int, b)
`seq`
            if ByteString -> Bool
S.null ByteString
k'
            then
                if ByteString -> Bool
S.null ByteString
q'
                then
                    case Maybe b
mv of
                    Maybe b
Nothing -> (Int, b) -> Maybe (Int, b)
forall a. a -> Maybe a
Just (Int
n0,b
v0)
                    Just b
v  -> (Int, b) -> Maybe (Int, b)
forall a. a -> Maybe a
Just (Int
n',b
v)
                else
                    case Maybe b
mv of
                    Maybe b
Nothing -> Int -> b -> Int -> ByteString -> Trie b -> Maybe (Int, b)
goJust Int
n0 b
v0 Int
n' ByteString
q' Trie b
t
                    Just b
v  -> Int -> b -> Int -> ByteString -> Trie b -> Maybe (Int, b)
goJust Int
n' b
v  Int
n' ByteString
q' Trie b
t
            else (Int, b) -> Maybe (Int, b)
forall a. a -> Maybe a
Just (Int
n0,b
v0)

    goJust Int
n0 b
v0 Int
n ByteString
q t_ :: Trie b
t_@(Branch Prefix
_ Prefix
_ Trie b
_ Trie b
_) = Trie b -> Maybe (Int, b)
findArc Trie b
t_
        where
        qh :: Prefix
qh = String -> ByteString -> Prefix
errorLogHead String
"match_" ByteString
q

        -- | /O(min(m,W))/, where /m/ is number of @Arc@s in this
        -- branching, and /W/ is the word size of the Prefix,Mask type.
        findArc :: Trie b -> Maybe (Int, b)
findArc (Branch Prefix
p Prefix
m Trie b
l Trie b
r)
            | Prefix -> Prefix -> Prefix -> Bool
nomatch Prefix
qh Prefix
p Prefix
m  = (Int, b) -> Maybe (Int, b)
forall a. a -> Maybe a
Just (Int
n0,b
v0)
            | Prefix -> Prefix -> Bool
zero Prefix
qh Prefix
m       = Trie b -> Maybe (Int, b)
findArc Trie b
l
            | Bool
otherwise       = Trie b -> Maybe (Int, b)
findArc Trie b
r
        findArc t :: Trie b
t@(Arc ByteString
_ Maybe b
_ Trie b
_) = Int -> b -> Int -> ByteString -> Trie b -> Maybe (Int, b)
goJust Int
n0 b
v0 Int
n ByteString
q Trie b
t
        findArc Trie b
Empty         = (Int, b) -> Maybe (Int, b)
forall a. a -> Maybe a
Just (Int
n0,b
v0)


-- | Given a query, find all prefixes with associated values in the
-- trie, returning their lengths and values. This function is a
-- good producer for list fusion.
--
-- This function may not have the most useful return type. For a
-- version that returns the prefix itself as well as the remaining
-- string, see @matches@ in "Data.Trie".
matches_ :: Trie a -> ByteString -> [(Int,a)]
matches_ :: Trie a -> ByteString -> [(Int, a)]
matches_ Trie a
t ByteString
q =
#if !defined(__GLASGOW_HASKELL__)
    matchFB_ t q (((:) .) . (,)) []
#else
    (forall b. ((Int, a) -> b -> b) -> b -> b) -> [(Int, a)]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\(Int, a) -> b -> b
cons b
nil -> Trie a -> ByteString -> (Int -> a -> b -> b) -> b -> b
forall a r. Trie a -> ByteString -> (Int -> a -> r -> r) -> r -> r
matchFB_ Trie a
t ByteString
q (((Int, a) -> b -> b
cons ((Int, a) -> b -> b) -> (a -> (Int, a)) -> a -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
.) ((a -> (Int, a)) -> a -> b -> b)
-> (Int -> a -> (Int, a)) -> Int -> a -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (,)) b
nil)
{-# INLINE matches_ #-}
#endif

matchFB_ :: Trie a -> ByteString -> (Int -> a -> r -> r) -> r -> r
matchFB_ :: Trie a -> ByteString -> (Int -> a -> r -> r) -> r -> r
matchFB_ = \Trie a
t ByteString
q Int -> a -> r -> r
cons r
nil -> (Int -> a -> r -> r) -> ByteString -> Trie a -> r -> r
forall t a. (Int -> t -> a -> a) -> ByteString -> Trie t -> a -> a
matchFB_' Int -> a -> r -> r
cons ByteString
q Trie a
t r
nil
    where
    matchFB_' :: (Int -> t -> a -> a) -> ByteString -> Trie t -> a -> a
matchFB_' Int -> t -> a -> a
cons = ByteString -> Trie t -> a -> a
start
        where
        -- | Deal with epsilon query (when there is no epsilon value)
        start :: ByteString -> Trie t -> a -> a
start ByteString
q (Branch Prefix
_ Prefix
_ Trie t
_ Trie t
_) | ByteString -> Bool
S.null ByteString
q = a -> a
forall a. a -> a
id
        start ByteString
q Trie t
t                           = Int -> ByteString -> Trie t -> a -> a
go Int
0 ByteString
q Trie t
t

        -- | The main recursion
        go :: Int -> ByteString -> Trie t -> a -> a
go Int
_ ByteString
_    Trie t
Empty       = a -> a
forall a. a -> a
id

        go Int
n ByteString
q   (Arc ByteString
k Maybe t
mv Trie t
t) =
            let (ByteString
p,ByteString
k',ByteString
q') = ByteString -> ByteString -> (ByteString, ByteString, ByteString)
breakMaximalPrefix ByteString
k ByteString
q
                n' :: Int
n'        = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ByteString -> Int
S.length ByteString
p
            in Int
n' Int -> (a -> a) -> a -> a
`seq`
                if ByteString -> Bool
S.null ByteString
k'
                then
                    case Maybe t
mv of { Maybe t
Nothing -> a -> a
forall a. a -> a
id; Just t
v  -> Int -> t -> a -> a
cons Int
n' t
v}
                    (a -> a) -> (a -> a) -> a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
                    if ByteString -> Bool
S.null ByteString
q' then a -> a
forall a. a -> a
id else Int -> ByteString -> Trie t -> a -> a
go Int
n' ByteString
q' Trie t
t
                else a -> a
forall a. a -> a
id

        go Int
n ByteString
q t_ :: Trie t
t_@(Branch Prefix
_ Prefix
_ Trie t
_ Trie t
_) = Trie t -> a -> a
findArc Trie t
t_
            where
            qh :: Prefix
qh = String -> ByteString -> Prefix
errorLogHead String
"matches_" ByteString
q

            -- | /O(min(m,W))/, where /m/ is number of @Arc@s in this
            -- branching, and /W/ is the word size of the Prefix,Mask type.
            findArc :: Trie t -> a -> a
findArc (Branch Prefix
p Prefix
m Trie t
l Trie t
r)
                | Prefix -> Prefix -> Prefix -> Bool
nomatch Prefix
qh Prefix
p Prefix
m  = a -> a
forall a. a -> a
id
                | Prefix -> Prefix -> Bool
zero Prefix
qh Prefix
m       = Trie t -> a -> a
findArc Trie t
l
                | Bool
otherwise       = Trie t -> a -> a
findArc Trie t
r
            findArc t :: Trie t
t@(Arc ByteString
_ Maybe t
_ Trie t
_) = Int -> ByteString -> Trie t -> a -> a
go Int
n ByteString
q Trie t
t
            findArc Trie t
Empty         = a -> a
forall a. a -> a
id


{-----------------------------------------------------------
-- Single-value modification functions (recurse and clone spine)
-----------------------------------------------------------}

-- TODO: We should CPS on Empty to avoid cloning spine if no change.
-- Difficulties arise with the calls to 'branch' and 'arc'. Will
-- have to create a continuation chain, so no savings on memory
-- allocation; but would have savings on held memory, if they're
-- still holding the old one...
--
-- | Generic function to alter a trie by one element with a function
-- to resolve conflicts (or non-conflicts).
alterBy :: (ByteString -> a -> Maybe a -> Maybe a)
         -> ByteString -> a -> Trie a -> Trie a
alterBy :: (ByteString -> a -> Maybe a -> Maybe a)
-> ByteString -> a -> Trie a -> Trie a
alterBy ByteString -> a -> Maybe a -> Maybe a
f = (ByteString -> a -> Maybe a -> Trie a -> (Maybe a, Trie a))
-> ByteString -> a -> Trie a -> Trie a
forall a.
(ByteString -> a -> Maybe a -> Trie a -> (Maybe a, Trie a))
-> ByteString -> a -> Trie a -> Trie a
alterBy_ (\ByteString
k a
v Maybe a
mv Trie a
t -> (ByteString -> a -> Maybe a -> Maybe a
f ByteString
k a
v Maybe a
mv, Trie a
t))
-- TODO: use GHC's 'inline' function so that this gets specialized away.
-- TODO: benchmark to be sure that this doesn't introduce unforseen performance costs because of the uncurrying etc.


-- | A variant of 'alterBy' which also allows modifying the sub-trie.
alterBy_ :: (ByteString -> a -> Maybe a -> Trie a -> (Maybe a, Trie a))
         -> ByteString -> a -> Trie a -> Trie a
alterBy_ :: (ByteString -> a -> Maybe a -> Trie a -> (Maybe a, Trie a))
-> ByteString -> a -> Trie a -> Trie a
alterBy_ ByteString -> a -> Maybe a -> Trie a -> (Maybe a, Trie a)
f_ ByteString
q_ a
x_
    | ByteString -> Bool
S.null ByteString
q_ = Trie a -> Trie a
alterEpsilon
    | Bool
otherwise = ByteString -> Trie a -> Trie a
go ByteString
q_
    where
    f :: Maybe a -> Trie a -> (Maybe a, Trie a)
f         = ByteString -> a -> Maybe a -> Trie a -> (Maybe a, Trie a)
f_ ByteString
q_ a
x_
    nothing :: ByteString -> Trie a
nothing ByteString
q = (Maybe a -> Trie a -> Trie a) -> (Maybe a, Trie a) -> Trie a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
q) (Maybe a -> Trie a -> (Maybe a, Trie a)
f Maybe a
forall a. Maybe a
Nothing Trie a
forall a. Trie a
Empty)

    alterEpsilon :: Trie a -> Trie a
alterEpsilon t_ :: Trie a
t_@Trie a
Empty                    = (Maybe a -> Trie a -> Trie a) -> (Maybe a, Trie a) -> Trie a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
q_) (Maybe a -> Trie a -> (Maybe a, Trie a)
f Maybe a
forall a. Maybe a
Nothing Trie a
t_)
    alterEpsilon t_ :: Trie a
t_@(Branch Prefix
_ Prefix
_ Trie a
_ Trie a
_)         = (Maybe a -> Trie a -> Trie a) -> (Maybe a, Trie a) -> Trie a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
q_) (Maybe a -> Trie a -> (Maybe a, Trie a)
f Maybe a
forall a. Maybe a
Nothing Trie a
t_)
    alterEpsilon t_ :: Trie a
t_@(Arc ByteString
k Maybe a
mv Trie a
t) | ByteString -> Bool
S.null ByteString
k  = (Maybe a -> Trie a -> Trie a) -> (Maybe a, Trie a) -> Trie a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
q_) (Maybe a -> Trie a -> (Maybe a, Trie a)
f Maybe a
mv      Trie a
t)
                                 | Bool
otherwise = (Maybe a -> Trie a -> Trie a) -> (Maybe a, Trie a) -> Trie a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
q_) (Maybe a -> Trie a -> (Maybe a, Trie a)
f Maybe a
forall a. Maybe a
Nothing Trie a
t_)


    go :: ByteString -> Trie a -> Trie a
go ByteString
q Trie a
Empty            = ByteString -> Trie a
nothing ByteString
q

    go ByteString
q t :: Trie a
t@(Branch Prefix
p Prefix
m Trie a
l Trie a
r)
        | Prefix -> Prefix -> Prefix -> Bool
nomatch Prefix
qh Prefix
p Prefix
m  = Prefix -> Trie a -> Prefix -> Trie a -> Trie a
forall a. Prefix -> Trie a -> Prefix -> Trie a -> Trie a
branchMerge Prefix
p Trie a
t  Prefix
qh (ByteString -> Trie a
nothing ByteString
q)
        | Prefix -> Prefix -> Bool
zero Prefix
qh Prefix
m       = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p Prefix
m (ByteString -> Trie a -> Trie a
go ByteString
q Trie a
l) Trie a
r
        | Bool
otherwise       = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p Prefix
m Trie a
l (ByteString -> Trie a -> Trie a
go ByteString
q Trie a
r)
        where
        qh :: Prefix
qh = String -> ByteString -> Prefix
errorLogHead String
"alterBy" ByteString
q

    go ByteString
q t_ :: Trie a
t_@(Arc ByteString
k Maybe a
mv Trie a
t) =
        let (ByteString
p,ByteString
k',ByteString
q') = ByteString -> ByteString -> (ByteString, ByteString, ByteString)
breakMaximalPrefix ByteString
k ByteString
q in
        case (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ByteString -> Bool
S.null ByteString
k', ByteString -> Bool
S.null ByteString
q') of
        (Bool
True,  Bool
True)  -> -- add node to middle of arc
                          (Maybe a -> Trie a -> Trie a) -> (Maybe a, Trie a) -> Trie a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
p) (Maybe a -> Trie a -> (Maybe a, Trie a)
f Maybe a
forall a. Maybe a
Nothing (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k' Maybe a
mv Trie a
t))
        (Bool
True,  Bool
False) ->
            case ByteString -> Trie a
nothing ByteString
q' of
            Trie a
Empty -> Trie a
t_ -- Nothing to add, reuse old arc
            Trie a
l     -> Trie a -> Trie a
forall a. Trie a -> Trie a
arc' (Prefix -> Trie a -> Prefix -> Trie a -> Trie a
forall a. Prefix -> Trie a -> Prefix -> Trie a -> Trie a
branchMerge (Trie a -> Prefix
forall a. Trie a -> Prefix
getPrefix Trie a
l) Trie a
l (Trie a -> Prefix
forall a. Trie a -> Prefix
getPrefix Trie a
r) Trie a
r)
                    where
                    r :: Trie a
r = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k' Maybe a
mv Trie a
t

                    -- inlined version of 'arc'
                    arc' :: Trie a -> Trie a
arc' | ByteString -> Bool
S.null ByteString
p  = Trie a -> Trie a
forall a. a -> a
id
                         | Bool
otherwise = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
p Maybe a
forall a. Maybe a
Nothing

        (Bool
False, Bool
True)  -> (Maybe a -> Trie a -> Trie a) -> (Maybe a, Trie a) -> Trie a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k) (Maybe a -> Trie a -> (Maybe a, Trie a)
f Maybe a
mv Trie a
t)
        (Bool
False, Bool
False) -> ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k Maybe a
mv (ByteString -> Trie a -> Trie a
go ByteString
q' Trie a
t)


-- | Alter the value associated with a given key. If the key is not
-- present, then the trie is returned unaltered. See 'alterBy' if
-- you are interested in inserting new keys or deleting old keys.
-- Because this function does not need to worry about changing the
-- trie structure, it is somewhat faster than 'alterBy'.
adjustBy :: (ByteString -> a -> a -> a)
         -> ByteString -> a -> Trie a -> Trie a
adjustBy :: (ByteString -> a -> a -> a) -> ByteString -> a -> Trie a -> Trie a
adjustBy ByteString -> a -> a -> a
f_ ByteString
q_ a
x_
    | ByteString -> Bool
S.null ByteString
q_ = Trie a -> Trie a
adjustEpsilon
    | Bool
otherwise = ByteString -> Trie a -> Trie a
go ByteString
q_
    where
    f :: a -> a
f = ByteString -> a -> a -> a
f_ ByteString
q_ a
x_

    adjustEpsilon :: Trie a -> Trie a
adjustEpsilon (Arc ByteString
k (Just a
v) Trie a
t) | ByteString -> Bool
S.null ByteString
k = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k (a -> Maybe a
forall a. a -> Maybe a
Just (a -> a
f a
v)) Trie a
t
    adjustEpsilon Trie a
t_                            = Trie a
t_

    go :: ByteString -> Trie a -> Trie a
go ByteString
_ Trie a
Empty            = Trie a
forall a. Trie a
Empty

    go ByteString
q t :: Trie a
t@(Branch Prefix
p Prefix
m Trie a
l Trie a
r)
        | Prefix -> Prefix -> Prefix -> Bool
nomatch Prefix
qh Prefix
p Prefix
m  = Trie a
t
        | Prefix -> Prefix -> Bool
zero Prefix
qh Prefix
m       = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p Prefix
m (ByteString -> Trie a -> Trie a
go ByteString
q Trie a
l) Trie a
r
        | Bool
otherwise       = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p Prefix
m Trie a
l (ByteString -> Trie a -> Trie a
go ByteString
q Trie a
r)
        where
        qh :: Prefix
qh = String -> ByteString -> Prefix
errorLogHead String
"adjustBy" ByteString
q

    go ByteString
q t_ :: Trie a
t_@(Arc ByteString
k Maybe a
mv Trie a
t) =
        let (ByteString
_,ByteString
k',ByteString
q') = ByteString -> ByteString -> (ByteString, ByteString, ByteString)
breakMaximalPrefix ByteString
k ByteString
q in
        case (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ ByteString -> Bool
S.null ByteString
k', ByteString -> Bool
S.null ByteString
q') of
        (Bool
True,  Bool
True)  -> Trie a
t_ -- don't break arc inline
        (Bool
True,  Bool
False) -> Trie a
t_ -- don't break arc branching
        (Bool
False, Bool
True)  -> ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k ((a -> a) -> Maybe a -> Maybe a
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM a -> a
f Maybe a
mv) Trie a
t
        (Bool
False, Bool
False) -> ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k Maybe a
mv (ByteString -> Trie a -> Trie a
go ByteString
q' Trie a
t)


{-----------------------------------------------------------
-- Trie-combining functions
-----------------------------------------------------------}

-- TEST CASES: foldr (unionL . uncurry singleton) empty t
--             foldr (uncurry insert) empty t
--    where t = map (\s -> (pk s, 0))
--                  ["heat","hello","hoi","apple","appa","hell","appb","appc"]
--
-- | Combine two tries, using a function to resolve collisions.
-- This can only define the space of functions between union and
-- symmetric difference but, with those two, all set operations can
-- be defined (albeit inefficiently).
mergeBy :: (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
mergeBy :: (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a
mergeBy a -> a -> Maybe a
f = Trie a -> Trie a -> Trie a
mergeBy'
    where
    -- | Deals with epsilon entries, before recursing into @go@
    mergeBy' :: Trie a -> Trie a -> Trie a
mergeBy'
        t0_ :: Trie a
t0_@(Arc ByteString
k0 Maybe a
mv0 Trie a
t0)
        t1_ :: Trie a
t1_@(Arc ByteString
k1 Maybe a
mv1 Trie a
t1)
        | ByteString -> Bool
S.null ByteString
k0 Bool -> Bool -> Bool
&& ByteString -> Bool
S.null ByteString
k1 = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k0 ((a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a
forall a. (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a
mergeMaybe a -> a -> Maybe a
f Maybe a
mv0 Maybe a
mv1) (Trie a -> Trie a -> Trie a
go Trie a
t0 Trie a
t1)
        | ByteString -> Bool
S.null ByteString
k0              = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k0 Maybe a
mv0 (Trie a -> Trie a -> Trie a
go Trie a
t0 Trie a
t1_)
        |              ByteString -> Bool
S.null ByteString
k1 = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k1 Maybe a
mv1 (Trie a -> Trie a -> Trie a
go Trie a
t1 Trie a
t0_)
    mergeBy'
        (Arc ByteString
k0 mv0 :: Maybe a
mv0@(Just a
_) Trie a
t0)
        t1_ :: Trie a
t1_@(Branch Prefix
_ Prefix
_ Trie a
_ Trie a
_)
        | ByteString -> Bool
S.null ByteString
k0              = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k0 Maybe a
mv0 (Trie a -> Trie a -> Trie a
go Trie a
t0 Trie a
t1_)
    mergeBy'
        t0_ :: Trie a
t0_@(Branch Prefix
_ Prefix
_ Trie a
_ Trie a
_)
        (Arc ByteString
k1 mv1 :: Maybe a
mv1@(Just a
_) Trie a
t1)
        | ByteString -> Bool
S.null ByteString
k1              = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k1 Maybe a
mv1 (Trie a -> Trie a -> Trie a
go Trie a
t1 Trie a
t0_)
    mergeBy' Trie a
t0_ Trie a
t1_             = Trie a -> Trie a -> Trie a
go Trie a
t0_ Trie a
t1_


    -- | The main recursion
    go :: Trie a -> Trie a -> Trie a
go Trie a
Empty Trie a
t1    = Trie a
t1
    go Trie a
t0    Trie a
Empty = Trie a
t0

    -- /O(n+m)/ for this part where /n/ and /m/ are sizes of the branchings
    go  t0 :: Trie a
t0@(Branch Prefix
p0 Prefix
m0 Trie a
l0 Trie a
r0)
        t1 :: Trie a
t1@(Branch Prefix
p1 Prefix
m1 Trie a
l1 Trie a
r1)
        | Prefix -> Prefix -> Bool
shorter Prefix
m0 Prefix
m1  = Trie a
union0
        | Prefix -> Prefix -> Bool
shorter Prefix
m1 Prefix
m0  = Trie a
union1
        | Prefix
p0 Prefix -> Prefix -> Bool
forall a. Eq a => a -> a -> Bool
== Prefix
p1       = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p0 Prefix
m0 (Trie a -> Trie a -> Trie a
go Trie a
l0 Trie a
l1) (Trie a -> Trie a -> Trie a
go Trie a
r0 Trie a
r1)
        | Bool
otherwise      = Prefix -> Trie a -> Prefix -> Trie a -> Trie a
forall a. Prefix -> Trie a -> Prefix -> Trie a -> Trie a
branchMerge Prefix
p0 Trie a
t0 Prefix
p1 Trie a
t1
        where
        union0 :: Trie a
union0  | Prefix -> Prefix -> Prefix -> Bool
nomatch Prefix
p1 Prefix
p0 Prefix
m0  = Prefix -> Trie a -> Prefix -> Trie a -> Trie a
forall a. Prefix -> Trie a -> Prefix -> Trie a -> Trie a
branchMerge Prefix
p0 Trie a
t0 Prefix
p1 Trie a
t1
                | Prefix -> Prefix -> Bool
zero Prefix
p1 Prefix
m0        = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p0 Prefix
m0 (Trie a -> Trie a -> Trie a
go Trie a
l0 Trie a
t1) Trie a
r0
                | Bool
otherwise         = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p0 Prefix
m0 Trie a
l0 (Trie a -> Trie a -> Trie a
go Trie a
r0 Trie a
t1)

        union1 :: Trie a
union1  | Prefix -> Prefix -> Prefix -> Bool
nomatch Prefix
p0 Prefix
p1 Prefix
m1  = Prefix -> Trie a -> Prefix -> Trie a -> Trie a
forall a. Prefix -> Trie a -> Prefix -> Trie a -> Trie a
branchMerge Prefix
p0 Trie a
t0 Prefix
p1 Trie a
t1
                | Prefix -> Prefix -> Bool
zero Prefix
p0 Prefix
m1        = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p1 Prefix
m1 (Trie a -> Trie a -> Trie a
go Trie a
t0 Trie a
l1) Trie a
r1
                | Bool
otherwise         = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p1 Prefix
m1 Trie a
l1 (Trie a -> Trie a -> Trie a
go Trie a
t0 Trie a
r1)

    -- We combine these branches of 'go' in order to clarify where the definitions of 'p0', 'p1', 'm'', 'p'' are relevant. However, this may introduce inefficiency in the pattern matching automaton...
    -- TODO: check. And get rid of 'go'' if it does.
    go Trie a
t0_ Trie a
t1_ = Trie a -> Trie a -> Trie a
go' Trie a
t0_ Trie a
t1_
        where
        p0 :: Prefix
p0 = Trie a -> Prefix
forall a. Trie a -> Prefix
getPrefix Trie a
t0_
        p1 :: Prefix
p1 = Trie a -> Prefix
forall a. Trie a -> Prefix
getPrefix Trie a
t1_
        m' :: Prefix
m' = Prefix -> Prefix -> Prefix
branchMask Prefix
p0 Prefix
p1
        p' :: Prefix
p' = Prefix -> Prefix -> Prefix
mask Prefix
p0 Prefix
m'

        go' :: Trie a -> Trie a -> Trie a
go' (Arc ByteString
k0 Maybe a
mv0 Trie a
t0)
            (Arc ByteString
k1 Maybe a
mv1 Trie a
t1)
            | Prefix
m' Prefix -> Prefix -> Bool
forall a. Eq a => a -> a -> Bool
== Prefix
0 =
                let (ByteString
pre,ByteString
k0',ByteString
k1') = ByteString -> ByteString -> (ByteString, ByteString, ByteString)
breakMaximalPrefix ByteString
k0 ByteString
k1 in
                if ByteString -> Bool
S.null ByteString
pre
                then String -> Trie a
forall a. HasCallStack => String -> a
error String
"mergeBy: no mask, but no prefix string"
                else let {-# INLINE arcMerge #-}
                         arcMerge :: Maybe a -> Trie a -> Trie a -> Trie a
arcMerge Maybe a
mv' Trie a
t1' Trie a
t2' = ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
pre Maybe a
mv' (Trie a -> Trie a -> Trie a
go Trie a
t1' Trie a
t2')
                     in case (ByteString -> Bool
S.null ByteString
k0', ByteString -> Bool
S.null ByteString
k1') of
                         (Bool
True, Bool
True)  -> Maybe a -> Trie a -> Trie a -> Trie a
arcMerge ((a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a
forall a. (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a
mergeMaybe a -> a -> Maybe a
f Maybe a
mv0 Maybe a
mv1) Trie a
t0 Trie a
t1
                         (Bool
True, Bool
False) -> Maybe a -> Trie a -> Trie a -> Trie a
arcMerge Maybe a
mv0 Trie a
t0 (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k1' Maybe a
mv1 Trie a
t1)
                         (Bool
False,Bool
True)  -> Maybe a -> Trie a -> Trie a -> Trie a
arcMerge Maybe a
mv1 (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k0' Maybe a
mv0 Trie a
t0) Trie a
t1
                         (Bool
False,Bool
False) -> Maybe a -> Trie a -> Trie a -> Trie a
arcMerge Maybe a
forall a. Maybe a
Nothing (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k0' Maybe a
mv0 Trie a
t0)
                                                           (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
Arc ByteString
k1' Maybe a
mv1 Trie a
t1)
        go' (Arc ByteString
_ Maybe a
_ Trie a
_)
            (Branch Prefix
_p1 Prefix
m1 Trie a
l Trie a
r)
            | Prefix -> Prefix -> Prefix -> Bool
nomatch Prefix
p0 Prefix
p1 Prefix
m1 = Prefix -> Trie a -> Prefix -> Trie a -> Trie a
forall a. Prefix -> Trie a -> Prefix -> Trie a -> Trie a
branchMerge Prefix
p1 Trie a
t1_  Prefix
p0 Trie a
t0_
            | Prefix -> Prefix -> Bool
zero Prefix
p0 Prefix
m1       = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p1 Prefix
m1 (Trie a -> Trie a -> Trie a
go Trie a
t0_ Trie a
l) Trie a
r
            | Bool
otherwise        = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p1 Prefix
m1 Trie a
l (Trie a -> Trie a -> Trie a
go Trie a
t0_ Trie a
r)
        go' (Branch Prefix
_p0 Prefix
m0 Trie a
l Trie a
r)
            (Arc ByteString
_ Maybe a
_ Trie a
_)
            | Prefix -> Prefix -> Prefix -> Bool
nomatch Prefix
p1 Prefix
p0 Prefix
m0 = Prefix -> Trie a -> Prefix -> Trie a -> Trie a
forall a. Prefix -> Trie a -> Prefix -> Trie a -> Trie a
branchMerge Prefix
p0 Trie a
t0_  Prefix
p1 Trie a
t1_
            | Prefix -> Prefix -> Bool
zero Prefix
p1 Prefix
m0       = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p0 Prefix
m0 (Trie a -> Trie a -> Trie a
go Trie a
l Trie a
t1_) Trie a
r
            | Bool
otherwise        = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p0 Prefix
m0 Trie a
l (Trie a -> Trie a -> Trie a
go Trie a
r Trie a
t1_)

        -- Inlined branchMerge. Both tries are disjoint @Arc@s now.
        go' Trie a
_ Trie a
_ | Prefix -> Prefix -> Bool
zero Prefix
p0 Prefix
m'   = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p' Prefix
m' Trie a
t0_ Trie a
t1_
        go' Trie a
_ Trie a
_                = Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
Branch Prefix
p' Prefix
m' Trie a
t1_ Trie a
t0_


mergeMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a
{-# INLINE mergeMaybe #-}
mergeMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a
mergeMaybe a -> a -> Maybe a
_ Maybe a
Nothing      Maybe a
Nothing  = Maybe a
forall a. Maybe a
Nothing
mergeMaybe a -> a -> Maybe a
_ Maybe a
Nothing mv1 :: Maybe a
mv1@(Just a
_)  = Maybe a
mv1
mergeMaybe a -> a -> Maybe a
_ mv0 :: Maybe a
mv0@(Just a
_) Maybe a
Nothing  = Maybe a
mv0
mergeMaybe a -> a -> Maybe a
f (Just a
v0)   (Just a
v1) = a -> a -> Maybe a
f a
v0 a
v1


{-----------------------------------------------------------
-- Priority-queue functions
-----------------------------------------------------------}

-- | Return the lexicographically smallest 'ByteString' and the
-- value it's mapped to; or 'Nothing' for the empty trie.  When one
-- entry is a prefix of another, the prefix will be returned.
minAssoc :: Trie a -> Maybe (ByteString, a)
minAssoc :: Trie a -> Maybe (ByteString, a)
minAssoc = ByteString -> Trie a -> Maybe (ByteString, a)
forall b. ByteString -> Trie b -> Maybe (ByteString, b)
go ByteString
S.empty
    where
    go :: ByteString -> Trie b -> Maybe (ByteString, b)
go ByteString
_ Trie b
Empty              = Maybe (ByteString, b)
forall a. Maybe a
Nothing
    go ByteString
q (Arc ByteString
k (Just b
v) Trie b
_) = (ByteString, b) -> Maybe (ByteString, b)
forall a. a -> Maybe a
Just (ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k,b
v)
    go ByteString
q (Arc ByteString
k Maybe b
Nothing  Trie b
t) = ByteString -> Trie b -> Maybe (ByteString, b)
go   (ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k) Trie b
t
    go ByteString
q (Branch Prefix
_ Prefix
_ Trie b
l Trie b
_)   = ByteString -> Trie b -> Maybe (ByteString, b)
go ByteString
q Trie b
l


-- | Return the lexicographically largest 'ByteString' and the
-- value it's mapped to; or 'Nothing' for the empty trie.  When one
-- entry is a prefix of another, the longer one will be returned.
maxAssoc :: Trie a -> Maybe (ByteString, a)
maxAssoc :: Trie a -> Maybe (ByteString, a)
maxAssoc = ByteString -> Trie a -> Maybe (ByteString, a)
forall b. ByteString -> Trie b -> Maybe (ByteString, b)
go ByteString
S.empty
    where
    go :: ByteString -> Trie b -> Maybe (ByteString, b)
go ByteString
_ Trie b
Empty                  = Maybe (ByteString, b)
forall a. Maybe a
Nothing
    go ByteString
q (Arc ByteString
k (Just b
v) Trie b
Empty) = (ByteString, b) -> Maybe (ByteString, b)
forall a. a -> Maybe a
Just (ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k,b
v)
    go ByteString
q (Arc ByteString
k Maybe b
_        Trie b
t)     = ByteString -> Trie b -> Maybe (ByteString, b)
go   (ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k) Trie b
t
    go ByteString
q (Branch Prefix
_ Prefix
_ Trie b
_ Trie b
r)       = ByteString -> Trie b -> Maybe (ByteString, b)
go ByteString
q Trie b
r


mapView :: (Trie a -> Trie a)
        -> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
mapView :: (Trie a -> Trie a)
-> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
mapView Trie a -> Trie a
_ Maybe (ByteString, a, Trie a)
Nothing        = Maybe (ByteString, a, Trie a)
forall a. Maybe a
Nothing
mapView Trie a -> Trie a
f (Just (ByteString
k,a
v,Trie a
t)) = (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
forall a. a -> Maybe a
Just (ByteString
k,a
v, Trie a -> Trie a
f Trie a
t)


-- | Update the 'minAssoc' and return the old 'minAssoc'.
updateMinViewBy :: (ByteString -> a -> Maybe a)
                -> Trie a -> Maybe (ByteString, a, Trie a)
updateMinViewBy :: (ByteString -> a -> Maybe a)
-> Trie a -> Maybe (ByteString, a, Trie a)
updateMinViewBy ByteString -> a -> Maybe a
f = ByteString -> Trie a -> Maybe (ByteString, a, Trie a)
go ByteString
S.empty
    where
    go :: ByteString -> Trie a -> Maybe (ByteString, a, Trie a)
go ByteString
_ Trie a
Empty              = Maybe (ByteString, a, Trie a)
forall a. Maybe a
Nothing
    go ByteString
q (Arc ByteString
k (Just a
v) Trie a
t) = let q' :: ByteString
q' = ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k
                              in (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
forall a. a -> Maybe a
Just (ByteString
q',a
v, ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k (ByteString -> a -> Maybe a
f ByteString
q' a
v) Trie a
t)
    go ByteString
q (Arc ByteString
k Maybe a
Nothing  Trie a
t) = (Trie a -> Trie a)
-> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
forall a.
(Trie a -> Trie a)
-> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
mapView (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k Maybe a
forall a. Maybe a
Nothing) (ByteString -> Trie a -> Maybe (ByteString, a, Trie a)
go (ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k) Trie a
t)
    go ByteString
q (Branch Prefix
p Prefix
m Trie a
l Trie a
r)   = (Trie a -> Trie a)
-> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
forall a.
(Trie a -> Trie a)
-> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
mapView (\Trie a
l' -> Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p Prefix
m Trie a
l' Trie a
r) (ByteString -> Trie a -> Maybe (ByteString, a, Trie a)
go ByteString
q Trie a
l)


-- | Update the 'maxAssoc' and return the old 'maxAssoc'.
updateMaxViewBy :: (ByteString -> a -> Maybe a)
                -> Trie a -> Maybe (ByteString, a, Trie a)
updateMaxViewBy :: (ByteString -> a -> Maybe a)
-> Trie a -> Maybe (ByteString, a, Trie a)
updateMaxViewBy ByteString -> a -> Maybe a
f = ByteString -> Trie a -> Maybe (ByteString, a, Trie a)
go ByteString
S.empty
    where
    go :: ByteString -> Trie a -> Maybe (ByteString, a, Trie a)
go ByteString
_ Trie a
Empty                  = Maybe (ByteString, a, Trie a)
forall a. Maybe a
Nothing
    go ByteString
q (Arc ByteString
k (Just a
v) Trie a
Empty) = let q' :: ByteString
q' = ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k
                                  in (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
forall a. a -> Maybe a
Just (ByteString
q',a
v, ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k (ByteString -> a -> Maybe a
f ByteString
q' a
v) Trie a
forall a. Trie a
Empty)
    go ByteString
q (Arc ByteString
k Maybe a
mv       Trie a
t)     = (Trie a -> Trie a)
-> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
forall a.
(Trie a -> Trie a)
-> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
mapView (ByteString -> Maybe a -> Trie a -> Trie a
forall a. ByteString -> Maybe a -> Trie a -> Trie a
arc ByteString
k Maybe a
mv) (ByteString -> Trie a -> Maybe (ByteString, a, Trie a)
go (ByteString -> ByteString -> ByteString
S.append ByteString
q ByteString
k) Trie a
t)
    go ByteString
q (Branch Prefix
p Prefix
m Trie a
l Trie a
r)       = (Trie a -> Trie a)
-> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
forall a.
(Trie a -> Trie a)
-> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a)
mapView (Prefix -> Prefix -> Trie a -> Trie a -> Trie a
forall a. Prefix -> Prefix -> Trie a -> Trie a -> Trie a
branch Prefix
p Prefix
m Trie a
l) (ByteString -> Trie a -> Maybe (ByteString, a, Trie a)
go ByteString
q Trie a
r)

------------------------------------------------------------
------------------------------------------------------- fin.