----------------------------------------------------------------------------
-- |
-- Module      :  Data.RadixTree.Internal
-- Copyright   :  (c) Sergey Vinokurov 2018
-- License     :  BSD3-style (see LICENSE)
-- Maintainer  :  serg.foo@gmail.com
--
-- This is an internal module that exposes innards of the 'RadixTree'
-- data structure. This API may change in any new release, even in a
-- patch release - depend on it at your own risk.
----------------------------------------------------------------------------

{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE DeriveFoldable      #-}
{-# LANGUAGE DeriveFunctor       #-}
{-# LANGUAGE DeriveGeneric       #-}
{-# LANGUAGE DeriveTraversable   #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE ScopedTypeVariables #-}

{-# OPTIONS_HADDOCK not-home #-}

module Data.RadixTree.Internal
  ( RadixTree(..)
  , empty
  , null
  , size
  , insert
  , insertWith
  , lookup
  , fromList
  , toList
  , toAscList
  , keys
  , keysSet
  , elems
  , mapMaybe
  , union
  , unionWith
  ) where

import Prelude hiding (lookup, null)

import Control.Arrow (first)
import Control.DeepSeq
import Control.Monad.ST
import Control.Monad.ST.Unsafe

import Data.ByteString.Short (ShortByteString)
import qualified Data.ByteString.Short as BSS
import qualified Data.ByteString.Short.Internal as BSSI
import qualified Data.Foldable as Foldable
import Data.IntMap (IntMap)
import qualified Data.IntMap.Strict as IM
import qualified Data.List as L
import Data.Maybe (fromMaybe)
import Data.Primitive.ByteArray
import Data.Semigroup as Semigroup
import Data.Set (Set)
import qualified Data.Set as S
import Data.Word
import GHC.Generics (Generic)

-- | A tree data structure that efficiently indexes values by string keys.
--
-- This type can be more memory-efficient than 'Data.Map' because it combines
-- common prefixes of all keys. Specific savings will vary depending on
-- concrete data set.
data RadixTree a
  = RadixNode
      !(Maybe a)
      !(IntMap (RadixTree a)) -- ^ Either has 0 or 2 or more children, never 1.
  | RadixStr
      !(Maybe a)
      {-# UNPACK #-} !ShortByteString -- ^ Non-empty
      !(RadixTree a)
  deriving (Int -> RadixTree a -> ShowS
[RadixTree a] -> ShowS
RadixTree a -> String
(Int -> RadixTree a -> ShowS)
-> (RadixTree a -> String)
-> ([RadixTree a] -> ShowS)
-> Show (RadixTree a)
forall a. Show a => Int -> RadixTree a -> ShowS
forall a. Show a => [RadixTree a] -> ShowS
forall a. Show a => RadixTree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> RadixTree a -> ShowS
showsPrec :: Int -> RadixTree a -> ShowS
$cshow :: forall a. Show a => RadixTree a -> String
show :: RadixTree a -> String
$cshowList :: forall a. Show a => [RadixTree a] -> ShowS
showList :: [RadixTree a] -> ShowS
Show, (forall a b. (a -> b) -> RadixTree a -> RadixTree b)
-> (forall a b. a -> RadixTree b -> RadixTree a)
-> Functor RadixTree
forall a b. a -> RadixTree b -> RadixTree a
forall a b. (a -> b) -> RadixTree a -> RadixTree 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) -> RadixTree a -> RadixTree b
fmap :: forall a b. (a -> b) -> RadixTree a -> RadixTree b
$c<$ :: forall a b. a -> RadixTree b -> RadixTree a
<$ :: forall a b. a -> RadixTree b -> RadixTree a
Functor, (forall m. Monoid m => RadixTree m -> m)
-> (forall m a. Monoid m => (a -> m) -> RadixTree a -> m)
-> (forall m a. Monoid m => (a -> m) -> RadixTree a -> m)
-> (forall a b. (a -> b -> b) -> b -> RadixTree a -> b)
-> (forall a b. (a -> b -> b) -> b -> RadixTree a -> b)
-> (forall b a. (b -> a -> b) -> b -> RadixTree a -> b)
-> (forall b a. (b -> a -> b) -> b -> RadixTree a -> b)
-> (forall a. (a -> a -> a) -> RadixTree a -> a)
-> (forall a. (a -> a -> a) -> RadixTree a -> a)
-> (forall a. RadixTree a -> [a])
-> (forall a. RadixTree a -> Bool)
-> (forall a. RadixTree a -> Int)
-> (forall a. Eq a => a -> RadixTree a -> Bool)
-> (forall a. Ord a => RadixTree a -> a)
-> (forall a. Ord a => RadixTree a -> a)
-> (forall a. Num a => RadixTree a -> a)
-> (forall a. Num a => RadixTree a -> a)
-> Foldable RadixTree
forall a. Eq a => a -> RadixTree a -> Bool
forall a. Num a => RadixTree a -> a
forall a. Ord a => RadixTree a -> a
forall m. Monoid m => RadixTree m -> m
forall a. RadixTree a -> Bool
forall a. RadixTree a -> Int
forall a. RadixTree a -> [a]
forall a. (a -> a -> a) -> RadixTree a -> a
forall m a. Monoid m => (a -> m) -> RadixTree a -> m
forall b a. (b -> a -> b) -> b -> RadixTree a -> b
forall a b. (a -> b -> b) -> b -> RadixTree a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall m. Monoid m => RadixTree m -> m
fold :: forall m. Monoid m => RadixTree m -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> RadixTree a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> RadixTree a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> RadixTree a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> RadixTree a -> m
$cfoldr :: forall a b. (a -> b -> b) -> b -> RadixTree a -> b
foldr :: forall a b. (a -> b -> b) -> b -> RadixTree a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> RadixTree a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> RadixTree a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> RadixTree a -> b
foldl :: forall b a. (b -> a -> b) -> b -> RadixTree a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> RadixTree a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> RadixTree a -> b
$cfoldr1 :: forall a. (a -> a -> a) -> RadixTree a -> a
foldr1 :: forall a. (a -> a -> a) -> RadixTree a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> RadixTree a -> a
foldl1 :: forall a. (a -> a -> a) -> RadixTree a -> a
$ctoList :: forall a. RadixTree a -> [a]
toList :: forall a. RadixTree a -> [a]
$cnull :: forall a. RadixTree a -> Bool
null :: forall a. RadixTree a -> Bool
$clength :: forall a. RadixTree a -> Int
length :: forall a. RadixTree a -> Int
$celem :: forall a. Eq a => a -> RadixTree a -> Bool
elem :: forall a. Eq a => a -> RadixTree a -> Bool
$cmaximum :: forall a. Ord a => RadixTree a -> a
maximum :: forall a. Ord a => RadixTree a -> a
$cminimum :: forall a. Ord a => RadixTree a -> a
minimum :: forall a. Ord a => RadixTree a -> a
$csum :: forall a. Num a => RadixTree a -> a
sum :: forall a. Num a => RadixTree a -> a
$cproduct :: forall a. Num a => RadixTree a -> a
product :: forall a. Num a => RadixTree a -> a
Foldable, Functor RadixTree
Foldable RadixTree
(Functor RadixTree, Foldable RadixTree) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> RadixTree a -> f (RadixTree b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    RadixTree (f a) -> f (RadixTree a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> RadixTree a -> m (RadixTree b))
-> (forall (m :: * -> *) a.
    Monad m =>
    RadixTree (m a) -> m (RadixTree a))
-> Traversable RadixTree
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a.
Monad m =>
RadixTree (m a) -> m (RadixTree a)
forall (f :: * -> *) a.
Applicative f =>
RadixTree (f a) -> f (RadixTree a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> RadixTree a -> m (RadixTree b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> RadixTree a -> f (RadixTree b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> RadixTree a -> f (RadixTree b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> RadixTree a -> f (RadixTree b)
$csequenceA :: forall (f :: * -> *) a.
Applicative f =>
RadixTree (f a) -> f (RadixTree a)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
RadixTree (f a) -> f (RadixTree a)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> RadixTree a -> m (RadixTree b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> RadixTree a -> m (RadixTree b)
$csequence :: forall (m :: * -> *) a.
Monad m =>
RadixTree (m a) -> m (RadixTree a)
sequence :: forall (m :: * -> *) a.
Monad m =>
RadixTree (m a) -> m (RadixTree a)
Traversable, (forall x. RadixTree a -> Rep (RadixTree a) x)
-> (forall x. Rep (RadixTree a) x -> RadixTree a)
-> Generic (RadixTree a)
forall x. Rep (RadixTree a) x -> RadixTree a
forall x. RadixTree a -> Rep (RadixTree a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (RadixTree a) x -> RadixTree a
forall a x. RadixTree a -> Rep (RadixTree a) x
$cfrom :: forall a x. RadixTree a -> Rep (RadixTree a) x
from :: forall x. RadixTree a -> Rep (RadixTree a) x
$cto :: forall a x. Rep (RadixTree a) x -> RadixTree a
to :: forall x. Rep (RadixTree a) x -> RadixTree a
Generic)

instance NFData a => NFData (RadixTree a)

-- | Radix tree with no elements.
empty :: RadixTree a
empty :: forall a. RadixTree a
empty = Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode Maybe a
forall a. Maybe a
Nothing IntMap (RadixTree a)
forall a. IntMap a
IM.empty

{-# INLINE interleaveST #-}
interleaveST :: ST s a -> ST s a
interleaveST :: forall s a. ST s a -> ST s a
interleaveST =
#if MIN_VERSION_base(4, 10, 0)
    ST s a -> ST s a
forall s a. ST s a -> ST s a
unsafeDupableInterleaveST
#else
    unsafeInterleaveST
#endif

splitShortByteString :: Int -> ShortByteString -> (ShortByteString, ShortByteString, Word8, ShortByteString)
splitShortByteString :: Int
-> ShortByteString
-> (ShortByteString, ShortByteString, Word8, ShortByteString)
splitShortByteString Int
n (BSSI.SBS ByteArray#
source) = (forall s.
 ST s (ShortByteString, ShortByteString, Word8, ShortByteString))
-> (ShortByteString, ShortByteString, Word8, ShortByteString)
forall a. (forall s. ST s a) -> a
runST ((forall s.
  ST s (ShortByteString, ShortByteString, Word8, ShortByteString))
 -> (ShortByteString, ShortByteString, Word8, ShortByteString))
-> (forall s.
    ST s (ShortByteString, ShortByteString, Word8, ShortByteString))
-> (ShortByteString, ShortByteString, Word8, ShortByteString)
forall a b. (a -> b) -> a -> b
$ do
  MutableByteArray s
prefix <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
prefixSize
  MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
prefix Int
0 ByteArray
source' Int
0 Int
prefixSize
  ByteArray ByteArray#
prefix# <- MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
prefix
  ByteArray
midSuffix         <- ST s ByteArray -> ST s ByteArray
forall s a. ST s a -> ST s a
interleaveST (ST s ByteArray -> ST s ByteArray)
-> ST s ByteArray -> ST s ByteArray
forall a b. (a -> b) -> a -> b
$ do
    MutableByteArray s
midSuffix <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
midSuffixSize
    MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
midSuffix Int
0 ByteArray
source' Int
n Int
midSuffixSize
    MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
midSuffix
  ByteArray
suffix            <- ST s ByteArray -> ST s ByteArray
forall s a. ST s a -> ST s a
interleaveST (ST s ByteArray -> ST s ByteArray)
-> ST s ByteArray -> ST s ByteArray
forall a b. (a -> b) -> a -> b
$ do
    MutableByteArray s
suffix <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
suffixSize
    MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
suffix Int
0 ByteArray
source' (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
suffixSize
    MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
suffix
  (ShortByteString, ShortByteString, Word8, ShortByteString)
-> ST s (ShortByteString, ShortByteString, Word8, ShortByteString)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteArray# -> ShortByteString
BSSI.SBS ByteArray#
prefix#, ByteArray -> ShortByteString
byteArrayToBSS ByteArray
midSuffix, ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
source' Int
n, ByteArray -> ShortByteString
byteArrayToBSS ByteArray
suffix)
  where
    source' :: ByteArray
source' = ByteArray# -> ByteArray
ByteArray ByteArray#
source
    prefixSize :: Int
prefixSize = Int
n
    midSuffixSize :: Int
midSuffixSize = ByteArray -> Int
sizeofByteArray ByteArray
source' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
prefixSize
    suffixSize :: Int
suffixSize = Int
midSuffixSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

{-# INLINE byteArrayToBSS #-}
byteArrayToBSS :: ByteArray -> BSS.ShortByteString
byteArrayToBSS :: ByteArray -> ShortByteString
byteArrayToBSS (ByteArray ByteArray#
xs) = ByteArray# -> ShortByteString
BSSI.SBS ByteArray#
xs

dropShortByteString :: Int -> ShortByteString -> ShortByteString
dropShortByteString :: Int -> ShortByteString -> ShortByteString
dropShortByteString Int
0  ShortByteString
src = ShortByteString
src
dropShortByteString !Int
n (BSSI.SBS ByteArray#
source) = (forall s. ST s ShortByteString) -> ShortByteString
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s ShortByteString) -> ShortByteString)
-> (forall s. ST s ShortByteString) -> ShortByteString
forall a b. (a -> b) -> a -> b
$ do
  MutableByteArray s
dest <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
sz
  MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dest Int
0 ByteArray
source' Int
n Int
sz
  ByteArray -> ShortByteString
byteArrayToBSS (ByteArray -> ShortByteString)
-> ST s ByteArray -> ST s ShortByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dest
  where
    source' :: ByteArray
source' = ByteArray# -> ByteArray
ByteArray ByteArray#
source
    !sz :: Int
sz = ByteArray -> Int
sizeofByteArray ByteArray
source' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n

singletonShortByteString :: Word8 -> ShortByteString
singletonShortByteString :: Word8 -> ShortByteString
singletonShortByteString !Word8
c = (forall s. ST s ShortByteString) -> ShortByteString
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s ShortByteString) -> ShortByteString)
-> (forall s. ST s ShortByteString) -> ShortByteString
forall a b. (a -> b) -> a -> b
$ do
  MutableByteArray s
dest <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray Int
1
  MutableByteArray (PrimState (ST s)) -> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dest Int
0 Word8
c
  ByteArray -> ShortByteString
byteArrayToBSS (ByteArray -> ShortByteString)
-> ST s ByteArray -> ST s ShortByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dest

{-# INLINE unsafeHeadeShortByteString #-}
unsafeHeadeShortByteString :: ShortByteString -> Word8
unsafeHeadeShortByteString :: ShortByteString -> Word8
unsafeHeadeShortByteString = (ShortByteString -> Int -> Word8
`BSSI.unsafeIndex` Int
0)

data Mismatch
  = IsPrefix
  | CommonPrefixThenMismatch
      !ShortByteString -- ^ Prefix of node contents common with the key
      ShortByteString  -- ^ Suffix with the first mismatching byte
      Word8            -- ^ First byte of the suffix that caused mismatch
      ShortByteString  -- ^ Rest of node contents, suffix
  deriving (Int -> Mismatch -> ShowS
[Mismatch] -> ShowS
Mismatch -> String
(Int -> Mismatch -> ShowS)
-> (Mismatch -> String) -> ([Mismatch] -> ShowS) -> Show Mismatch
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Mismatch -> ShowS
showsPrec :: Int -> Mismatch -> ShowS
$cshow :: Mismatch -> String
show :: Mismatch -> String
$cshowList :: [Mismatch] -> ShowS
showList :: [Mismatch] -> ShowS
Show, (forall x. Mismatch -> Rep Mismatch x)
-> (forall x. Rep Mismatch x -> Mismatch) -> Generic Mismatch
forall x. Rep Mismatch x -> Mismatch
forall x. Mismatch -> Rep Mismatch x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Mismatch -> Rep Mismatch x
from :: forall x. Mismatch -> Rep Mismatch x
$cto :: forall x. Rep Mismatch x -> Mismatch
to :: forall x. Rep Mismatch x -> Mismatch
Generic)

analyseMismatch
  :: ShortByteString -- ^ Key
  -> Int             -- ^ Key offset
  -> ShortByteString -- ^ Node contents
  -> Mismatch
analyseMismatch :: ShortByteString -> Int -> ShortByteString -> Mismatch
analyseMismatch (BSSI.SBS ByteArray#
key) !Int
keyOffset nodeContentsBS :: ShortByteString
nodeContentsBS@(BSSI.SBS ByteArray#
nodeContents) =
  case Int -> Maybe Int
findMismatch Int
0 of
    Maybe Int
Nothing          -> Mismatch
IsPrefix
    Just Int
mismatchIdx ->
      case Int
-> ShortByteString
-> (ShortByteString, ShortByteString, Word8, ShortByteString)
splitShortByteString Int
mismatchIdx ShortByteString
nodeContentsBS of
        (ShortByteString
prefix, ShortByteString
midSuffix, Word8
mid, ShortByteString
suffix) -> ShortByteString
-> ShortByteString -> Word8 -> ShortByteString -> Mismatch
CommonPrefixThenMismatch ShortByteString
prefix ShortByteString
midSuffix Word8
mid ShortByteString
suffix
  where
    keySize :: Int
keySize      = ByteArray -> Int
sizeofByteArray ByteArray
key'
    keyLeft :: Int
keyLeft      = Int
keySize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
keyOffset
    contentsSize :: Int
contentsSize = ByteArray -> Int
sizeofByteArray ByteArray
nodeContents'

    key' :: ByteArray
key'          = ByteArray# -> ByteArray
ByteArray ByteArray#
key
    nodeContents' :: ByteArray
nodeContents' = ByteArray# -> ByteArray
ByteArray ByteArray#
nodeContents

    limit :: Int
    limit :: Int
limit = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
keyLeft Int
contentsSize

    findMismatch :: Int -> Maybe Int
    findMismatch :: Int -> Maybe Int
findMismatch !Int
i
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
limit
      = if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
contentsSize
        then Maybe Int
forall a. Maybe a
Nothing
        else Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i -- Key ended in the middle of node's packed key.
      | (ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
key' (Int
keyOffset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i) :: Word8) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
nodeContents' Int
i
      = Int -> Maybe Int
findMismatch (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      | Bool
otherwise
      = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i

mkRadixNodeFuse :: Maybe a -> IntMap (RadixTree a) -> Maybe (RadixTree a)
mkRadixNodeFuse :: forall a. Maybe a -> IntMap (RadixTree a) -> Maybe (RadixTree a)
mkRadixNodeFuse Maybe a
val IntMap (RadixTree a)
children =
  case Maybe a
val of
    Maybe a
Nothing | IntMap (RadixTree a) -> Bool
forall a. IntMap a -> Bool
IM.null IntMap (RadixTree a)
children
      -> Maybe (RadixTree a)
forall a. Maybe a
Nothing
    Maybe a
val'    | [(Int
c, RadixTree a
child)] <- IntMap (RadixTree a) -> [(Int, RadixTree a)]
forall a. IntMap a -> [(Int, a)]
IM.toList IntMap (RadixTree a)
children
      -> RadixTree a -> Maybe (RadixTree a)
forall a. a -> Maybe a
Just (RadixTree a -> Maybe (RadixTree a))
-> RadixTree a -> Maybe (RadixTree a)
forall a b. (a -> b) -> a -> b
$ Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr Maybe a
val' (Word8 -> ShortByteString
singletonShortByteString (Word8 -> ShortByteString) -> Word8 -> ShortByteString
forall a b. (a -> b) -> a -> b
$ Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c) RadixTree a
child
    Maybe a
_ -> RadixTree a -> Maybe (RadixTree a)
forall a. a -> Maybe a
Just (RadixTree a -> Maybe (RadixTree a))
-> RadixTree a -> Maybe (RadixTree a)
forall a b. (a -> b) -> a -> b
$ Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode Maybe a
val IntMap (RadixTree a)
children

-- Precondition: input string is non-empty
mkRadixStrFuse :: Maybe a -> ShortByteString -> RadixTree a -> Maybe (RadixTree a)
mkRadixStrFuse :: forall a.
Maybe a -> ShortByteString -> RadixTree a -> Maybe (RadixTree a)
mkRadixStrFuse Maybe a
val ShortByteString
str RadixTree a
rest =
  case (Maybe a
val, RadixTree a
rest) of
    (Maybe a
val',    RadixStr Maybe a
Nothing ShortByteString
str' RadixTree a
rest') ->
      RadixTree a -> Maybe (RadixTree a)
forall a. a -> Maybe a
Just (RadixTree a -> Maybe (RadixTree a))
-> RadixTree a -> Maybe (RadixTree a)
forall a b. (a -> b) -> a -> b
$ Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr Maybe a
val' (ShortByteString
str ShortByteString -> ShortByteString -> ShortByteString
forall a. Semigroup a => a -> a -> a
Semigroup.<> ShortByteString
str') RadixTree a
rest'
    (Maybe a
Nothing, RadixTree a
node)
      | RadixTree a -> Bool
forall a. RadixTree a -> Bool
null RadixTree a
node -> Maybe (RadixTree a)
forall a. Maybe a
Nothing
    (Maybe a
val', RadixTree a
rest') ->
      RadixTree a -> Maybe (RadixTree a)
forall a. a -> Maybe a
Just (RadixTree a -> Maybe (RadixTree a))
-> RadixTree a -> Maybe (RadixTree a)
forall a b. (a -> b) -> a -> b
$ Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr Maybe a
val' ShortByteString
str RadixTree a
rest'

mkRadixStr :: ShortByteString -> RadixTree a -> RadixTree a
mkRadixStr :: forall a. ShortByteString -> RadixTree a -> RadixTree a
mkRadixStr ShortByteString
str RadixTree a
rest
  | ShortByteString -> Bool
BSS.null ShortByteString
str = RadixTree a
rest
  | Bool
otherwise    = Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr Maybe a
forall a. Maybe a
Nothing ShortByteString
str RadixTree a
rest

-- TODO: prove following function correct.

-- | Check whether radix tree is empty
null :: RadixTree a -> Bool
null :: forall a. RadixTree a -> Bool
null = \case
  RadixNode Maybe a
Nothing IntMap (RadixTree a)
children -> IntMap (RadixTree a) -> Bool
forall a. IntMap a -> Bool
IM.null IntMap (RadixTree a)
children
  RadixStr Maybe a
Nothing ShortByteString
_ RadixTree a
rest    -> RadixTree a -> Bool
forall a. RadixTree a -> Bool
null RadixTree a
rest
  RadixTree a
_                          -> Bool
False

-- | O(n) Get number of elements in a radix tree.
size :: RadixTree a -> Int
size :: forall a. RadixTree a -> Int
size = RadixTree a -> Int
forall a. RadixTree a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length

-- | Add new element to a radix tree.
insert :: forall a. ShortByteString -> a -> RadixTree a -> RadixTree a
insert :: forall a. ShortByteString -> a -> RadixTree a -> RadixTree a
insert = (a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a
forall a.
(a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a
insertWith a -> a -> a
forall a b. a -> b -> a
const

-- | Add new element to a radix tree. If an element was already present for
-- the given key, use supplied funciton @f@ to produce a new value. The
-- function will be called like this @f newValue oldValue@.
insertWith :: forall a. (a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a
insertWith :: forall a.
(a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a
insertWith = (a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a
forall a.
(a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a
insert'

{-# INLINE insert' #-}
insert' :: forall a. (a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a
insert' :: forall a.
(a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a
insert' a -> a -> a
f ShortByteString
key a
value = Int -> RadixTree a -> RadixTree a
go Int
0
  where
    len :: Int
len = ShortByteString -> Int
BSS.length ShortByteString
key

    readKey :: Int -> Int
    readKey :: Int -> Int
readKey = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Int) -> (Int -> Word8) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShortByteString -> Int -> Word8
BSSI.unsafeIndex ShortByteString
key

    go :: Int -> RadixTree a -> RadixTree a
    go :: Int -> RadixTree a -> RadixTree a
go Int
i
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len
      = \case
        RadixNode Maybe a
oldValue IntMap (RadixTree a)
children
          | IntMap (RadixTree a) -> Bool
forall a. IntMap a -> Bool
IM.null IntMap (RadixTree a)
children ->
            Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr Maybe a
oldValue (Int -> ShortByteString -> ShortByteString
dropShortByteString Int
i ShortByteString
key) (RadixTree a -> RadixTree a) -> RadixTree a -> RadixTree a
forall a b. (a -> b) -> a -> b
$ Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode (a -> Maybe a
forall a. a -> Maybe a
Just a
value) IntMap (RadixTree a)
forall a. IntMap a
IM.empty
          | Bool
otherwise ->
            Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode Maybe a
oldValue (IntMap (RadixTree a) -> RadixTree a)
-> IntMap (RadixTree a) -> RadixTree a
forall a b. (a -> b) -> a -> b
$
            (Maybe (RadixTree a) -> Maybe (RadixTree a))
-> Int -> IntMap (RadixTree a) -> IntMap (RadixTree a)
forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
IM.alter (RadixTree a -> Maybe (RadixTree a)
forall a. a -> Maybe a
Just (RadixTree a -> Maybe (RadixTree a))
-> (Maybe (RadixTree a) -> RadixTree a)
-> Maybe (RadixTree a)
-> Maybe (RadixTree a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RadixTree a
-> (RadixTree a -> RadixTree a)
-> Maybe (RadixTree a)
-> RadixTree a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe RadixTree a
optNode (Int -> RadixTree a -> RadixTree a
go Int
i')) Int
c IntMap (RadixTree a)
children
          where
            c :: Int
            c :: Int
c = Int -> Int
readKey Int
i
            i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            optNode :: RadixTree a
optNode =
              ShortByteString -> RadixTree a -> RadixTree a
forall a. ShortByteString -> RadixTree a -> RadixTree a
mkRadixStr (Int -> ShortByteString -> ShortByteString
dropShortByteString Int
i' ShortByteString
key) (RadixTree a -> RadixTree a) -> RadixTree a -> RadixTree a
forall a b. (a -> b) -> a -> b
$ Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode (a -> Maybe a
forall a. a -> Maybe a
Just a
value) IntMap (RadixTree a)
forall a. IntMap a
IM.empty
        RadixStr Maybe a
oldValue ShortByteString
packedKey RadixTree a
rest ->
          case ShortByteString -> Int -> ShortByteString -> Mismatch
analyseMismatch ShortByteString
key Int
i ShortByteString
packedKey of
            Mismatch
IsPrefix ->
              Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr Maybe a
oldValue ShortByteString
packedKey (RadixTree a -> RadixTree a) -> RadixTree a -> RadixTree a
forall a b. (a -> b) -> a -> b
$ Int -> RadixTree a -> RadixTree a
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ShortByteString -> Int
BSS.length ShortByteString
packedKey) RadixTree a
rest
            CommonPrefixThenMismatch ShortByteString
prefix ShortByteString
midSuffix Word8
mid ShortByteString
suffix ->
              (if ShortByteString -> Bool
BSS.null ShortByteString
prefix then RadixTree a -> RadixTree a
forall a. a -> a
id else Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr Maybe a
oldValue ShortByteString
prefix) (RadixTree a -> RadixTree a) -> RadixTree a -> RadixTree a
forall a b. (a -> b) -> a -> b
$
                if Bool
isKeyEnded
                then
                  Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr (a -> Maybe a
forall a. a -> Maybe a
Just a
value) ShortByteString
midSuffix RadixTree a
rest
                else
                  Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode (if ShortByteString -> Bool
BSS.null ShortByteString
prefix then Maybe a
oldValue else Maybe a
forall a. Maybe a
Nothing) (IntMap (RadixTree a) -> RadixTree a)
-> IntMap (RadixTree a) -> RadixTree a
forall a b. (a -> b) -> a -> b
$
                  [(Int, RadixTree a)] -> IntMap (RadixTree a)
forall a. [(Int, a)] -> IntMap a
IM.fromList
                    [ ( Int
mid'
                      , ShortByteString -> RadixTree a -> RadixTree a
forall a. ShortByteString -> RadixTree a -> RadixTree a
mkRadixStr ShortByteString
suffix RadixTree a
rest
                      )
                    , ( Int -> Int
readKey Int
i'
                      , ShortByteString -> RadixTree a -> RadixTree a
forall a. ShortByteString -> RadixTree a -> RadixTree a
mkRadixStr (Int -> ShortByteString -> ShortByteString
dropShortByteString (Int
i' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ShortByteString
key) (RadixTree a -> RadixTree a) -> RadixTree a -> RadixTree a
forall a b. (a -> b) -> a -> b
$ Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode (a -> Maybe a
forall a. a -> Maybe a
Just a
value) IntMap (RadixTree a)
forall a. IntMap a
IM.empty
                      )
                    ]
              where
                i' :: Int
i'         = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ShortByteString -> Int
BSS.length ShortByteString
prefix
                isKeyEnded :: Bool
isKeyEnded = Int
i' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len
                mid' :: Int
mid'       = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
mid
      | Bool
otherwise
      = \case
        RadixNode Maybe a
oldValue IntMap (RadixTree a)
children ->
          Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode (a -> Maybe a
forall a. a -> Maybe a
Just (a -> (a -> a) -> Maybe a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a
value (a -> a -> a
f a
value) Maybe a
oldValue)) IntMap (RadixTree a)
children
        RadixStr Maybe a
oldValue ShortByteString
key' RadixTree a
rest ->
          Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr (a -> Maybe a
forall a. a -> Maybe a
Just (a -> (a -> a) -> Maybe a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a
value (a -> a -> a
f a
value) Maybe a
oldValue)) ShortByteString
key' RadixTree a
rest

canStripPrefixFromShortByteString
  :: Int -> ShortByteString -> ShortByteString -> Bool
canStripPrefixFromShortByteString :: Int -> ShortByteString -> ShortByteString -> Bool
canStripPrefixFromShortByteString Int
bigStart (BSSI.SBS ByteArray#
small) (BSSI.SBS ByteArray#
big)
  | Int
bigStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
smallSize Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
bigSize = Bool
False
  | Bool
otherwise                      = Int -> Bool
findMismatch Int
0
  where
    small' :: ByteArray
small' = ByteArray# -> ByteArray
ByteArray ByteArray#
small
    big' :: ByteArray
big'   = ByteArray# -> ByteArray
ByteArray ByteArray#
big

    smallSize :: Int
smallSize = ByteArray -> Int
sizeofByteArray ByteArray
small'
    bigSize :: Int
bigSize   = ByteArray -> Int
sizeofByteArray ByteArray
big'

    findMismatch :: Int -> Bool
    findMismatch :: Int -> Bool
findMismatch !Int
i
      | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
smallSize
      = Bool
True
      | (ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
small' Int
i :: Word8) Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== ByteArray -> Int -> Word8
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
big' (Int
bigStart Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i)
      = Int -> Bool
findMismatch (Int -> Bool) -> Int -> Bool
forall a b. (a -> b) -> a -> b
$ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
      | Bool
otherwise
      = Bool
False

-- | O(length(key)) Try to find a value associated with the given key.
lookup :: forall a. ShortByteString -> RadixTree a -> Maybe a
lookup :: forall a. ShortByteString -> RadixTree a -> Maybe a
lookup ShortByteString
key = Int -> RadixTree a -> Maybe a
go Int
0
  where
    len :: Int
len = ShortByteString -> Int
BSS.length ShortByteString
key

    readKey :: Int -> Int
    readKey :: Int -> Int
readKey = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Int) -> (Int -> Word8) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShortByteString -> Int -> Word8
BSSI.unsafeIndex ShortByteString
key

    go :: Int -> RadixTree a -> Maybe a
    go :: Int -> RadixTree a -> Maybe a
go !Int
n RadixTree a
tree
      | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len
      = case RadixTree a
tree of
        RadixNode Maybe a
val IntMap (RadixTree a)
_  -> Maybe a
val
        RadixStr Maybe a
val ShortByteString
_ RadixTree a
_ -> Maybe a
val
      | Bool
otherwise
      = case RadixTree a
tree of
      RadixNode Maybe a
_ IntMap (RadixTree a)
children      ->
        Int -> IntMap (RadixTree a) -> Maybe (RadixTree a)
forall a. Int -> IntMap a -> Maybe a
IM.lookup (Int -> Int
readKey Int
n) IntMap (RadixTree a)
children Maybe (RadixTree a) -> (RadixTree a -> Maybe a) -> Maybe a
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> RadixTree a -> Maybe a
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      RadixStr Maybe a
_ ShortByteString
packedKey RadixTree a
rest
        | Int -> ShortByteString -> ShortByteString -> Bool
canStripPrefixFromShortByteString Int
n ShortByteString
packedKey ShortByteString
key
        -> Int -> RadixTree a -> Maybe a
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ShortByteString -> Int
BSS.length ShortByteString
packedKey) RadixTree a
rest
        | Bool
otherwise
        -> Maybe a
forall a. Maybe a
Nothing

-- | Construct a radix tree from list of key-value pairs. If some key
-- appears twice in the input list, later occurrences will override
-- earlier ones.
fromList :: [(ShortByteString, a)] -> RadixTree a
fromList :: forall a. [(ShortByteString, a)] -> RadixTree a
fromList =
  (RadixTree a -> (ShortByteString, a) -> RadixTree a)
-> RadixTree a -> [(ShortByteString, a)] -> RadixTree a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' (\RadixTree a
acc (ShortByteString
k, a
v) -> (a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a
forall a.
(a -> a -> a) -> ShortByteString -> a -> RadixTree a -> RadixTree a
insert' a -> a -> a
forall a b. a -> b -> a
const ShortByteString
k a
v RadixTree a
acc) RadixTree a
forall a. RadixTree a
empty

-- | O(n) Convert a radix tree to a list of key-value pairs.
toList :: RadixTree a -> [(ShortByteString, a)]
toList :: forall a. RadixTree a -> [(ShortByteString, a)]
toList = RadixTree a -> [(ShortByteString, a)]
forall a. RadixTree a -> [(ShortByteString, a)]
toAscList

-- | O(n) Convert a radix tree to an ascending list of key-value pairs.
toAscList :: forall a. RadixTree a -> [(ShortByteString, a)]
toAscList :: forall a. RadixTree a -> [(ShortByteString, a)]
toAscList = (([Word8], a) -> (ShortByteString, a))
-> [([Word8], a)] -> [(ShortByteString, a)]
forall a b. (a -> b) -> [a] -> [b]
map (([Word8] -> ShortByteString)
-> ([Word8], a) -> (ShortByteString, a)
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first [Word8] -> ShortByteString
BSS.pack) ([([Word8], a)] -> [(ShortByteString, a)])
-> (RadixTree a -> [([Word8], a)])
-> RadixTree a
-> [(ShortByteString, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RadixTree a -> [([Word8], a)]
go
  where
    go :: RadixTree a -> [([Word8], a)]
    go :: RadixTree a -> [([Word8], a)]
go = \case
      RadixNode Maybe a
val IntMap (RadixTree a)
children ->
        ([([Word8], a)] -> [([Word8], a)])
-> (a -> [([Word8], a)] -> [([Word8], a)])
-> Maybe a
-> [([Word8], a)]
-> [([Word8], a)]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [([Word8], a)] -> [([Word8], a)]
forall a. a -> a
id (\a
val' [([Word8], a)]
ys -> ([], a
val') ([Word8], a) -> [([Word8], a)] -> [([Word8], a)]
forall a. a -> [a] -> [a]
: [([Word8], a)]
ys) Maybe a
val ([([Word8], a)] -> [([Word8], a)])
-> [([Word8], a)] -> [([Word8], a)]
forall a b. (a -> b) -> a -> b
$
        (Int -> RadixTree a -> [([Word8], a)])
-> IntMap (RadixTree a) -> [([Word8], a)]
forall m a. Monoid m => (Int -> a -> m) -> IntMap a -> m
IM.foldMapWithKey (\Int
c RadixTree a
child -> (([Word8], a) -> ([Word8], a)) -> [([Word8], a)] -> [([Word8], a)]
forall a b. (a -> b) -> [a] -> [b]
map (([Word8] -> [Word8]) -> ([Word8], a) -> ([Word8], a)
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c Word8 -> [Word8] -> [Word8]
forall a. a -> [a] -> [a]
:)) ([([Word8], a)] -> [([Word8], a)])
-> [([Word8], a)] -> [([Word8], a)]
forall a b. (a -> b) -> a -> b
$ RadixTree a -> [([Word8], a)]
go RadixTree a
child) IntMap (RadixTree a)
children
      RadixStr Maybe a
val ShortByteString
packedKey RadixTree a
rest ->
        ([([Word8], a)] -> [([Word8], a)])
-> (a -> [([Word8], a)] -> [([Word8], a)])
-> Maybe a
-> [([Word8], a)]
-> [([Word8], a)]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [([Word8], a)] -> [([Word8], a)]
forall a. a -> a
id (\a
val' [([Word8], a)]
ys -> ([], a
val') ([Word8], a) -> [([Word8], a)] -> [([Word8], a)]
forall a. a -> [a] -> [a]
: [([Word8], a)]
ys) Maybe a
val ([([Word8], a)] -> [([Word8], a)])
-> [([Word8], a)] -> [([Word8], a)]
forall a b. (a -> b) -> a -> b
$
        (([Word8], a) -> ([Word8], a)) -> [([Word8], a)] -> [([Word8], a)]
forall a b. (a -> b) -> [a] -> [b]
map (([Word8] -> [Word8]) -> ([Word8], a) -> ([Word8], a)
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (ShortByteString -> [Word8]
BSS.unpack ShortByteString
packedKey [Word8] -> [Word8] -> [Word8]
forall a. [a] -> [a] -> [a]
++)) ([([Word8], a)] -> [([Word8], a)])
-> [([Word8], a)] -> [([Word8], a)]
forall a b. (a -> b) -> a -> b
$
        RadixTree a -> [([Word8], a)]
go RadixTree a
rest

-- | O(n) Get all keys stored in a radix tree.
keys :: RadixTree a -> [ShortByteString]
keys :: forall a. RadixTree a -> [ShortByteString]
keys = ([Word8] -> ShortByteString) -> [[Word8]] -> [ShortByteString]
forall a b. (a -> b) -> [a] -> [b]
map [Word8] -> ShortByteString
BSS.pack ([[Word8]] -> [ShortByteString])
-> (RadixTree a -> [[Word8]]) -> RadixTree a -> [ShortByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RadixTree a -> [[Word8]]
forall a. RadixTree a -> [[Word8]]
go
  where
    go :: RadixTree a -> [[Word8]]
    go :: forall a. RadixTree a -> [[Word8]]
go = \case
      RadixNode Maybe a
val IntMap (RadixTree a)
children ->
        ([[Word8]] -> [[Word8]])
-> (a -> [[Word8]] -> [[Word8]])
-> Maybe a
-> [[Word8]]
-> [[Word8]]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [[Word8]] -> [[Word8]]
forall a. a -> a
id (\a
_ [[Word8]]
ys -> [] [Word8] -> [[Word8]] -> [[Word8]]
forall a. a -> [a] -> [a]
: [[Word8]]
ys) Maybe a
val ([[Word8]] -> [[Word8]]) -> [[Word8]] -> [[Word8]]
forall a b. (a -> b) -> a -> b
$
        (Int -> RadixTree a -> [[Word8]])
-> IntMap (RadixTree a) -> [[Word8]]
forall m a. Monoid m => (Int -> a -> m) -> IntMap a -> m
IM.foldMapWithKey (\Int
c RadixTree a
child -> ([Word8] -> [Word8]) -> [[Word8]] -> [[Word8]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
c Word8 -> [Word8] -> [Word8]
forall a. a -> [a] -> [a]
:) ([[Word8]] -> [[Word8]]) -> [[Word8]] -> [[Word8]]
forall a b. (a -> b) -> a -> b
$ RadixTree a -> [[Word8]]
forall a. RadixTree a -> [[Word8]]
go RadixTree a
child) IntMap (RadixTree a)
children
      RadixStr Maybe a
val ShortByteString
packedKey RadixTree a
rest ->
        ([[Word8]] -> [[Word8]])
-> (a -> [[Word8]] -> [[Word8]])
-> Maybe a
-> [[Word8]]
-> [[Word8]]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [[Word8]] -> [[Word8]]
forall a. a -> a
id (\a
_ [[Word8]]
ys -> [] [Word8] -> [[Word8]] -> [[Word8]]
forall a. a -> [a] -> [a]
: [[Word8]]
ys) Maybe a
val ([[Word8]] -> [[Word8]]) -> [[Word8]] -> [[Word8]]
forall a b. (a -> b) -> a -> b
$
        ([Word8] -> [Word8]) -> [[Word8]] -> [[Word8]]
forall a b. (a -> b) -> [a] -> [b]
map (ShortByteString -> [Word8]
BSS.unpack ShortByteString
packedKey [Word8] -> [Word8] -> [Word8]
forall a. Semigroup a => a -> a -> a
<>) ([[Word8]] -> [[Word8]]) -> [[Word8]] -> [[Word8]]
forall a b. (a -> b) -> a -> b
$
        RadixTree a -> [[Word8]]
forall a. RadixTree a -> [[Word8]]
go RadixTree a
rest

-- | O(n) Get set of all keys stored in a radix tree.
keysSet :: RadixTree a -> Set ShortByteString
keysSet :: forall a. RadixTree a -> Set ShortByteString
keysSet = [ShortByteString] -> Set ShortByteString
forall a. [a] -> Set a
S.fromDistinctAscList ([ShortByteString] -> Set ShortByteString)
-> (RadixTree a -> [ShortByteString])
-> RadixTree a
-> Set ShortByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RadixTree a -> [ShortByteString]
forall a. RadixTree a -> [ShortByteString]
keys

-- | O(n) Get all values stored in a radix tree.
elems :: RadixTree a -> [a]
elems :: forall a. RadixTree a -> [a]
elems = RadixTree a -> [a]
forall a. RadixTree a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Foldable.toList

-- | O(n) Map a function that can remove some existing elements over a
-- radix tree.
mapMaybe :: forall a b. (a -> Maybe b) -> RadixTree a -> RadixTree b
mapMaybe :: forall a b. (a -> Maybe b) -> RadixTree a -> RadixTree b
mapMaybe a -> Maybe b
f = RadixTree b -> Maybe (RadixTree b) -> RadixTree b
forall a. a -> Maybe a -> a
fromMaybe RadixTree b
forall a. RadixTree a
empty (Maybe (RadixTree b) -> RadixTree b)
-> (RadixTree a -> Maybe (RadixTree b))
-> RadixTree a
-> RadixTree b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RadixTree a -> Maybe (RadixTree b)
go
  where
    go :: RadixTree a -> Maybe (RadixTree b)
    go :: RadixTree a -> Maybe (RadixTree b)
go = \case
      RadixNode Maybe a
val IntMap (RadixTree a)
children ->
        Maybe b -> IntMap (RadixTree b) -> Maybe (RadixTree b)
forall a. Maybe a -> IntMap (RadixTree a) -> Maybe (RadixTree a)
mkRadixNodeFuse (a -> Maybe b
f (a -> Maybe b) -> Maybe a -> Maybe b
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Maybe a
val) (IntMap (RadixTree b) -> Maybe (RadixTree b))
-> IntMap (RadixTree b) -> Maybe (RadixTree b)
forall a b. (a -> b) -> a -> b
$ (RadixTree a -> Maybe (RadixTree b))
-> IntMap (RadixTree a) -> IntMap (RadixTree b)
forall a b. (a -> Maybe b) -> IntMap a -> IntMap b
IM.mapMaybe RadixTree a -> Maybe (RadixTree b)
go IntMap (RadixTree a)
children
      RadixStr Maybe a
val ShortByteString
str RadixTree a
rest ->
        Maybe b -> ShortByteString -> RadixTree b -> Maybe (RadixTree b)
forall a.
Maybe a -> ShortByteString -> RadixTree a -> Maybe (RadixTree a)
mkRadixStrFuse (a -> Maybe b
f (a -> Maybe b) -> Maybe a -> Maybe b
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Maybe a
val) ShortByteString
str (RadixTree b -> Maybe (RadixTree b))
-> RadixTree b -> Maybe (RadixTree b)
forall a b. (a -> b) -> a -> b
$ RadixTree b -> Maybe (RadixTree b) -> RadixTree b
forall a. a -> Maybe a -> a
fromMaybe RadixTree b
forall a. RadixTree a
empty (Maybe (RadixTree b) -> RadixTree b)
-> Maybe (RadixTree b) -> RadixTree b
forall a b. (a -> b) -> a -> b
$ RadixTree a -> Maybe (RadixTree b)
go RadixTree a
rest

-- | O(n + m) Combine two radix trees trees. If a key is present in both
-- trees then the value from left one will be retained.
union :: RadixTree a -> RadixTree a -> RadixTree a
union :: forall a. RadixTree a -> RadixTree a -> RadixTree a
union = (a -> a -> a) -> RadixTree a -> RadixTree a -> RadixTree a
forall a.
(a -> a -> a) -> RadixTree a -> RadixTree a -> RadixTree a
unionWith a -> a -> a
forall a b. a -> b -> a
const

-- | O(n + m) Combine two trees using supplied function to resolve
-- values that have the same key in both trees.
unionWith :: forall a. (a -> a -> a) -> RadixTree a -> RadixTree a -> RadixTree a
unionWith :: forall a.
(a -> a -> a) -> RadixTree a -> RadixTree a -> RadixTree a
unionWith a -> a -> a
f = RadixTree a -> RadixTree a -> RadixTree a
go
  where
    combineVals :: Maybe a -> Maybe a -> Maybe a
    combineVals :: Maybe a -> Maybe a -> Maybe a
combineVals Maybe a
x Maybe a
y = case (Maybe a
x, Maybe a
y) of
      (Maybe a
Nothing,   Maybe a
Nothing)   -> Maybe a
forall a. Maybe a
Nothing
      (Maybe a
Nothing,   y' :: Maybe a
y'@Just{}) -> Maybe a
y'
      (x' :: Maybe a
x'@Just{}, Maybe a
Nothing)   -> Maybe a
x'
      (Just a
x',   Just a
y')   -> a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ a -> a -> a
f a
x' a
y'

    go :: RadixTree a -> RadixTree a -> RadixTree a
    go :: RadixTree a -> RadixTree a -> RadixTree a
go RadixTree a
x RadixTree a
y = case (RadixTree a
x, RadixTree a
y) of
      (RadixNode Maybe a
val IntMap (RadixTree a)
children, RadixNode Maybe a
val' IntMap (RadixTree a)
children') ->
        Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode (Maybe a -> Maybe a -> Maybe a
combineVals Maybe a
val Maybe a
val') ((RadixTree a -> RadixTree a -> RadixTree a)
-> IntMap (RadixTree a)
-> IntMap (RadixTree a)
-> IntMap (RadixTree a)
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
IM.unionWith RadixTree a -> RadixTree a -> RadixTree a
go IntMap (RadixTree a)
children IntMap (RadixTree a)
children')
      (RadixNode Maybe a
val IntMap (RadixTree a)
children, RadixStr Maybe a
val' ShortByteString
str' RadixTree a
rest') ->
        Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode (Maybe a -> Maybe a -> Maybe a
combineVals Maybe a
val Maybe a
val') (IntMap (RadixTree a) -> RadixTree a)
-> IntMap (RadixTree a) -> RadixTree a
forall a b. (a -> b) -> a -> b
$
          (\Maybe (RadixTree a) -> Maybe (RadixTree a)
g -> (Maybe (RadixTree a) -> Maybe (RadixTree a))
-> Int -> IntMap (RadixTree a) -> IntMap (RadixTree a)
forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
IM.alter Maybe (RadixTree a) -> Maybe (RadixTree a)
g Int
h IntMap (RadixTree a)
children) ((Maybe (RadixTree a) -> Maybe (RadixTree a))
 -> IntMap (RadixTree a))
-> (Maybe (RadixTree a) -> Maybe (RadixTree a))
-> IntMap (RadixTree a)
forall a b. (a -> b) -> a -> b
$ \Maybe (RadixTree a)
child ->
            RadixTree a -> Maybe (RadixTree a)
forall a. a -> Maybe a
Just (RadixTree a -> Maybe (RadixTree a))
-> RadixTree a -> Maybe (RadixTree a)
forall a b. (a -> b) -> a -> b
$!
            let rest'' :: RadixTree a
rest'' = ShortByteString -> RadixTree a -> RadixTree a
forall a. ShortByteString -> RadixTree a -> RadixTree a
mkRadixStr (Int -> ShortByteString -> ShortByteString
dropShortByteString Int
1 ShortByteString
str') RadixTree a
rest' in
            case Maybe (RadixTree a)
child of
              Maybe (RadixTree a)
Nothing     -> RadixTree a
rest''
              Just RadixTree a
child' -> RadixTree a -> RadixTree a -> RadixTree a
go RadixTree a
child' RadixTree a
rest''
        where
          h :: Int
h = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Int) -> Word8 -> Int
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Word8
unsafeHeadeShortByteString ShortByteString
str'
      (RadixStr Maybe a
val ShortByteString
str RadixTree a
rest, RadixNode Maybe a
val' IntMap (RadixTree a)
children') ->
        Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode (Maybe a -> Maybe a -> Maybe a
combineVals Maybe a
val Maybe a
val') (IntMap (RadixTree a) -> RadixTree a)
-> IntMap (RadixTree a) -> RadixTree a
forall a b. (a -> b) -> a -> b
$
          (\Maybe (RadixTree a) -> Maybe (RadixTree a)
g -> (Maybe (RadixTree a) -> Maybe (RadixTree a))
-> Int -> IntMap (RadixTree a) -> IntMap (RadixTree a)
forall a. (Maybe a -> Maybe a) -> Int -> IntMap a -> IntMap a
IM.alter Maybe (RadixTree a) -> Maybe (RadixTree a)
g Int
h IntMap (RadixTree a)
children') ((Maybe (RadixTree a) -> Maybe (RadixTree a))
 -> IntMap (RadixTree a))
-> (Maybe (RadixTree a) -> Maybe (RadixTree a))
-> IntMap (RadixTree a)
forall a b. (a -> b) -> a -> b
$ \Maybe (RadixTree a)
child ->
            RadixTree a -> Maybe (RadixTree a)
forall a. a -> Maybe a
Just (RadixTree a -> Maybe (RadixTree a))
-> RadixTree a -> Maybe (RadixTree a)
forall a b. (a -> b) -> a -> b
$!
            let rest' :: RadixTree a
rest' = ShortByteString -> RadixTree a -> RadixTree a
forall a. ShortByteString -> RadixTree a -> RadixTree a
mkRadixStr (Int -> ShortByteString -> ShortByteString
dropShortByteString Int
1 ShortByteString
str) RadixTree a
rest in
            case Maybe (RadixTree a)
child of
              Maybe (RadixTree a)
Nothing     -> RadixTree a
rest'
              Just RadixTree a
child' -> RadixTree a -> RadixTree a -> RadixTree a
go RadixTree a
rest' RadixTree a
child'
        where
          h :: Int
h = Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Int) -> Word8 -> Int
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Word8
unsafeHeadeShortByteString ShortByteString
str
      (RadixStr Maybe a
val ShortByteString
str RadixTree a
rest, RadixStr Maybe a
val' ShortByteString
str' RadixTree a
rest') ->
        case ShortByteString -> Int -> ShortByteString -> Mismatch
analyseMismatch ShortByteString
str Int
0 ShortByteString
str' of
          -- str' is a prefix of str
          Mismatch
IsPrefix ->
            Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr (Maybe a -> Maybe a -> Maybe a
combineVals Maybe a
val Maybe a
val') ShortByteString
str' (RadixTree a -> RadixTree a) -> RadixTree a -> RadixTree a
forall a b. (a -> b) -> a -> b
$
              RadixTree a -> RadixTree a -> RadixTree a
go (ShortByteString -> RadixTree a -> RadixTree a
forall a. ShortByteString -> RadixTree a -> RadixTree a
mkRadixStr (Int -> ShortByteString -> ShortByteString
dropShortByteString (ShortByteString -> Int
BSS.length ShortByteString
str') ShortByteString
str) RadixTree a
rest) RadixTree a
rest'
          -- str' = prefix + firstMismatchStr' + suffixStr'
          --      = prefix + midSuffixStr'
          CommonPrefixThenMismatch ShortByteString
prefix ShortByteString
midSuffixStr' Word8
firstMismatchStr' ShortByteString
suffixStr' ->
            (if ShortByteString -> Bool
BSS.null ShortByteString
prefix then RadixTree a -> RadixTree a
forall a. a -> a
id else Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr (Maybe a -> Maybe a -> Maybe a
combineVals Maybe a
val Maybe a
val') ShortByteString
prefix) (RadixTree a -> RadixTree a) -> RadixTree a -> RadixTree a
forall a b. (a -> b) -> a -> b
$
              if ShortByteString -> Int
BSS.length ShortByteString
prefix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ShortByteString -> Int
BSS.length ShortByteString
str
              then
                RadixTree a -> RadixTree a -> RadixTree a
go RadixTree a
rest (RadixTree a -> RadixTree a) -> RadixTree a -> RadixTree a
forall a b. (a -> b) -> a -> b
$ Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
forall a. Maybe a -> ShortByteString -> RadixTree a -> RadixTree a
RadixStr
                  (if ShortByteString -> Bool
BSS.null ShortByteString
prefix then Maybe a -> Maybe a -> Maybe a
combineVals Maybe a
val Maybe a
val' else Maybe a
forall a. Maybe a
Nothing)
                  ShortByteString
midSuffixStr'
                  RadixTree a
rest'
              else Maybe a -> IntMap (RadixTree a) -> RadixTree a
forall a. Maybe a -> IntMap (RadixTree a) -> RadixTree a
RadixNode (if ShortByteString -> Bool
BSS.null ShortByteString
prefix then Maybe a -> Maybe a -> Maybe a
combineVals Maybe a
val Maybe a
val' else Maybe a
forall a. Maybe a
Nothing) (IntMap (RadixTree a) -> RadixTree a)
-> IntMap (RadixTree a) -> RadixTree a
forall a b. (a -> b) -> a -> b
$ [(Int, RadixTree a)] -> IntMap (RadixTree a)
forall a. [(Int, a)] -> IntMap a
IM.fromList
                [ ( Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
firstMismatchStr'
                  , ShortByteString -> RadixTree a -> RadixTree a
forall a. ShortByteString -> RadixTree a -> RadixTree a
mkRadixStr ShortByteString
suffixStr' RadixTree a
rest'
                  )
                , ( Word8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> Int) -> Word8 -> Int
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Int -> Word8
BSSI.unsafeIndex ShortByteString
str (Int -> Word8) -> Int -> Word8
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Int
BSS.length ShortByteString
prefix
                  , ShortByteString -> RadixTree a -> RadixTree a
forall a. ShortByteString -> RadixTree a -> RadixTree a
mkRadixStr (Int -> ShortByteString -> ShortByteString
dropShortByteString (ShortByteString -> Int
BSSI.length ShortByteString
prefix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ShortByteString
str) RadixTree a
rest
                  )
                ]