{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

module Data.RRBVector.Internal
    ( Vector(..)
    , Tree(..)
    -- * Internal
    , blockShift, blockSize, treeSize, computeSizes, up
    -- * Construction
    , empty, singleton, fromList
    -- ** Concatenation
    , (<|), (|>), (><)
    -- * Deconstruction
    , viewl, viewr
    -- * Indexing
    , lookup, index
    , (!?), (!)
    , update
    , adjust, adjust'
    , take, drop, splitAt
    , insertAt, deleteAt
    -- * Transformations
    , map, reverse
    -- * Zipping and unzipping
    , zip, zipWith, unzip
    ) where

import Control.Applicative (Alternative, liftA2)
import qualified Control.Applicative
import Control.DeepSeq
import Control.Monad (when, MonadPlus)
import Control.Monad.ST (runST)
#if !(MIN_VERSION_base(4,13,0))
import Control.Monad.Fail (MonadFail(..))
#endif
import Control.Monad.Fix (MonadFix(..))
import Control.Monad.Zip (MonadZip(..))

import Data.Bits
import Data.Foldable (Foldable(..), for_)
import Data.Functor.Classes
import Data.Functor.Identity (Identity(..))
import Data.Maybe (fromMaybe)
import qualified Data.List as List
import qualified GHC.Exts as Exts
import GHC.Stack (HasCallStack)
import Text.Read
import Prelude hiding (lookup, map, take, drop, splitAt, head, last, reverse, zip, zipWith, unzip)

import Data.Functor.WithIndex
import Data.Foldable.WithIndex
import Data.Traversable.WithIndex

import Data.Primitive.PrimArray
import qualified Data.RRBVector.Internal.Array as A
import qualified Data.RRBVector.Internal.Buffer as Buffer
import Data.RRBVector.Internal.Indexed

infixr 5 ><
infixr 5 <|
infixl 5 |>

-- Invariant: Children of a Balanced node are always balanced.
-- A Leaf node is considered balanced.
-- Nodes are always non-empty.
data Tree a
    = Balanced !(A.Array (Tree a))
    | Unbalanced !(A.Array (Tree a)) !(PrimArray Int)
    | Leaf !(A.Array a)

-- | A vector.
--
-- The instances are based on those of @Seq@s, which are in turn based on those of lists.
data Vector a
    = Empty
    | Root
        !Int -- size
        !Int -- shift (blockShift * height)
        !(Tree a)

-- The number of bits used per level.
blockShift :: Int
blockShift :: Int
blockShift = Int
4
{-# INLINE blockShift #-}

-- The maximum size of a block.
blockSize :: Int
blockSize :: Int
blockSize = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
blockShift

-- The mask used to extract the index into the array.
blockMask :: Int
blockMask :: Int
blockMask = Int
blockSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

up :: Int -> Int
up :: Int -> Int
up Int
sh = Int
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
blockShift
{-# INLINE up #-}

down :: Int -> Int
down :: Int -> Int
down Int
sh = Int
sh Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
blockShift
{-# INLINE down #-}

radixIndex :: Int -> Int -> Int
radixIndex :: Int -> Int -> Int
radixIndex Int
i Int
sh = Int
i Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
sh Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
blockMask
{-# INLINE radixIndex #-}

relaxedRadixIndex :: PrimArray Int -> Int -> Int -> (Int, Int)
relaxedRadixIndex :: PrimArray Int -> Int -> Int -> (Int, Int)
relaxedRadixIndex PrimArray Int
sizes Int
i Int
sh =
    let guess :: Int
guess = Int -> Int -> Int
radixIndex Int
i Int
sh -- guess <= idx
        idx :: Int
idx = Int -> Int
loop Int
guess
        subIdx :: Int
subIdx = if Int
idx Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Int
i else Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    in (Int
idx, Int
subIdx)
  where
    loop :: Int -> Int
loop Int
idx =
        let current :: Int
current = PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes Int
idx -- idx will always be in range for a well-formed tree
        in if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
current then Int
idx else Int -> Int
loop (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
{-# INLINE relaxedRadixIndex #-}

treeToArray :: Tree a -> A.Array (Tree a)
treeToArray :: Tree a -> Array (Tree a)
treeToArray (Balanced Array (Tree a)
arr) = Array (Tree a)
arr
treeToArray (Unbalanced Array (Tree a)
arr PrimArray Int
_) = Array (Tree a)
arr
treeToArray (Leaf Array a
_) = [Char] -> Array (Tree a)
forall a. HasCallStack => [Char] -> a
error [Char]
"treeToArray: leaf"

treeBalanced :: Tree a -> Bool
treeBalanced :: Tree a -> Bool
treeBalanced (Balanced Array (Tree a)
_) = Bool
True
treeBalanced (Unbalanced Array (Tree a)
_ PrimArray Int
_) = Bool
False
treeBalanced (Leaf Array a
_) = Bool
True

-- @treeSize sh@ is the size of a tree with shift @sh@.
treeSize :: Int -> Tree a -> Int
treeSize :: Int -> Tree a -> Int
treeSize = Int -> Int -> Tree a -> Int
forall a. Int -> Int -> Tree a -> Int
go Int
0
  where
    go :: Int -> Int -> Tree a -> Int
go Int
acc Int
_ (Leaf Array a
arr) = Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr
    go Int
acc Int
_ (Unbalanced Array (Tree a)
_ PrimArray Int
sizes) = Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (PrimArray Int -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
sizes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    go Int
acc Int
sh (Balanced Array (Tree a)
arr) =
        let i :: Int
i = Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        in Int -> Int -> Tree a -> Int
go (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
sh)) (Int -> Int
down Int
sh) (Array (Tree a) -> Int -> Tree a
forall a. Array a -> Int -> a
A.index Array (Tree a)
arr Int
i)
{-# INLINE treeSize #-}

-- @computeSizes sh@ turns an array into a tree node by computing the sizes of its subtrees.
-- @sh@ is the shift of the resulting tree.
computeSizes :: Int -> A.Array (Tree a) -> Tree a
computeSizes :: Int -> Array (Tree a) -> Tree a
computeSizes Int
sh Array (Tree a)
arr = (forall s. ST s (Tree a)) -> Tree a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Tree a)) -> Tree a)
-> (forall s. ST s (Tree a)) -> Tree a
forall a b. (a -> b) -> a -> b
$ do
    let len :: Int
len = Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr
        maxSize :: Int
maxSize = Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftL` Int
sh -- the maximum size of a subtree
    MutablePrimArray s Int
sizes <- Int -> ST s (MutablePrimArray (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
    let loop :: Int -> Bool -> Int -> ST s Bool
loop Int
acc Bool
isBalanced Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len =
                let subtree :: Tree a
subtree = Array (Tree a) -> Int -> Tree a
forall a. Array a -> Int -> a
A.index Array (Tree a)
arr Int
i
                    size :: Int
size = Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh) Tree a
subtree
                    acc' :: Int
acc' = Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
size
                    isBalanced' :: Bool
isBalanced' = Bool
isBalanced Bool -> Bool -> Bool
&& if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 then Tree a -> Bool
forall a. Tree a -> Bool
treeBalanced Tree a
subtree else Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maxSize
                in MutablePrimArray (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
sizes Int
i Int
acc' ST s () -> ST s Bool -> ST s Bool
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> Bool -> Int -> ST s Bool
loop Int
acc' Bool
isBalanced' (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            | Bool
otherwise = Bool -> ST s Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
isBalanced
    Bool
isBalanced <- Int -> Bool -> Int -> ST s Bool
loop Int
0 Bool
True Int
0
    if Bool
isBalanced then
        Tree a -> ST s (Tree a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tree a -> ST s (Tree a)) -> Tree a -> ST s (Tree a)
forall a b. (a -> b) -> a -> b
$ Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced Array (Tree a)
arr
    else do
        PrimArray Int
sizes <- MutablePrimArray (PrimState (ST s)) Int -> ST s (PrimArray Int)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
sizes -- safe because the mutable @sizes@ isn't used afterwards
        Tree a -> ST s (Tree a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Tree a -> ST s (Tree a)) -> Tree a -> ST s (Tree a)
forall a b. (a -> b) -> a -> b
$ Array (Tree a) -> PrimArray Int -> Tree a
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced Array (Tree a)
arr PrimArray Int
sizes

-- Integer log base 2.
log2 :: Int -> Int
log2 :: Int -> Int
log2 Int
x = Int
bitSizeMinus1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros Int
x
  where
    bitSizeMinus1 :: Int
bitSizeMinus1 = Int -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Int
0 :: Int) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
{-# INLINE log2 #-}

instance Show1 Vector where
    liftShowsPrec :: (Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> Vector a -> ShowS
liftShowsPrec Int -> a -> ShowS
sp [a] -> ShowS
sl Int
p Vector a
v = (Int -> [a] -> ShowS) -> [Char] -> Int -> [a] -> ShowS
forall a. (Int -> a -> ShowS) -> [Char] -> Int -> a -> ShowS
showsUnaryWith ((Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> [a] -> ShowS
forall (f :: * -> *) a.
Show1 f =>
(Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> f a -> ShowS
liftShowsPrec Int -> a -> ShowS
sp [a] -> ShowS
sl) [Char]
"fromList" Int
p (Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v)

instance (Show a) => Show (Vector a) where
    showsPrec :: Int -> Vector a -> ShowS
showsPrec = Int -> Vector a -> ShowS
forall (f :: * -> *) a. (Show1 f, Show a) => Int -> f a -> ShowS
showsPrec1

instance Read1 Vector where
    liftReadPrec :: ReadPrec a -> ReadPrec [a] -> ReadPrec (Vector a)
liftReadPrec ReadPrec a
rp ReadPrec [a]
rl = ReadPrec (Vector a) -> ReadPrec (Vector a)
forall a. ReadPrec a -> ReadPrec a
readData (ReadPrec (Vector a) -> ReadPrec (Vector a))
-> ReadPrec (Vector a) -> ReadPrec (Vector a)
forall a b. (a -> b) -> a -> b
$ ReadPrec [a] -> [Char] -> ([a] -> Vector a) -> ReadPrec (Vector a)
forall a t. ReadPrec a -> [Char] -> (a -> t) -> ReadPrec t
readUnaryWith (ReadPrec a -> ReadPrec [a] -> ReadPrec [a]
forall (f :: * -> *) a.
Read1 f =>
ReadPrec a -> ReadPrec [a] -> ReadPrec (f a)
liftReadPrec ReadPrec a
rp ReadPrec [a]
rl) [Char]
"fromList" [a] -> Vector a
forall a. [a] -> Vector a
fromList
    liftReadListPrec :: ReadPrec a -> ReadPrec [a] -> ReadPrec [Vector a]
liftReadListPrec = ReadPrec a -> ReadPrec [a] -> ReadPrec [Vector a]
forall (f :: * -> *) a.
Read1 f =>
ReadPrec a -> ReadPrec [a] -> ReadPrec [f a]
liftReadListPrecDefault

instance (Read a) => Read (Vector a) where
    readPrec :: ReadPrec (Vector a)
readPrec = ReadPrec (Vector a)
forall (f :: * -> *) a. (Read1 f, Read a) => ReadPrec (f a)
readPrec1
    readListPrec :: ReadPrec [Vector a]
readListPrec = ReadPrec [Vector a]
forall a. Read a => ReadPrec [a]
readListPrecDefault

instance Eq1 Vector where
    liftEq :: (a -> b -> Bool) -> Vector a -> Vector b -> Bool
liftEq a -> b -> Bool
f Vector a
v1 Vector b
v2 = Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector b -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector b
v2 Bool -> Bool -> Bool
&& (a -> b -> Bool) -> [a] -> [b] -> Bool
forall (f :: * -> *) a b.
Eq1 f =>
(a -> b -> Bool) -> f a -> f b -> Bool
liftEq a -> b -> Bool
f (Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (Vector b -> [b]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector b
v2)

instance (Eq a) => Eq (Vector a) where
    == :: Vector a -> Vector a -> Bool
(==) = Vector a -> Vector a -> Bool
forall (f :: * -> *) a. (Eq1 f, Eq a) => f a -> f a -> Bool
eq1

instance Ord1 Vector where
    liftCompare :: (a -> b -> Ordering) -> Vector a -> Vector b -> Ordering
liftCompare a -> b -> Ordering
f Vector a
v1 Vector b
v2 = (a -> b -> Ordering) -> [a] -> [b] -> Ordering
forall (f :: * -> *) a b.
Ord1 f =>
(a -> b -> Ordering) -> f a -> f b -> Ordering
liftCompare a -> b -> Ordering
f (Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (Vector b -> [b]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector b
v2)

instance (Ord a) => Ord (Vector a) where
    compare :: Vector a -> Vector a -> Ordering
compare = Vector a -> Vector a -> Ordering
forall (f :: * -> *) a. (Ord1 f, Ord a) => f a -> f a -> Ordering
compare1

instance Semigroup (Vector a) where
    Vector a
v1 <> :: Vector a -> Vector a -> Vector a
<> Vector a
v2 = Vector a
v1 Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
>< Vector a
v2

instance Monoid (Vector a) where
    mempty :: Vector a
mempty = Vector a
forall a. Vector a
empty

instance Foldable Vector where
    foldr :: (a -> b -> b) -> b -> Vector a -> b
foldr a -> b -> b
f b
acc = Vector a -> b
go
      where
        go :: Vector a -> b
go Vector a
Empty = b
acc
        go (Root Int
_ Int
_ Tree a
tree) = Tree a -> b -> b
foldrTree Tree a
tree b
acc

        foldrTree :: Tree a -> b -> b
foldrTree (Balanced Array (Tree a)
arr) b
acc' = (Tree a -> b -> b) -> b -> Array (Tree a) -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Tree a -> b -> b
foldrTree b
acc' Array (Tree a)
arr
        foldrTree (Unbalanced Array (Tree a)
arr PrimArray Int
_) b
acc' = (Tree a -> b -> b) -> b -> Array (Tree a) -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Tree a -> b -> b
foldrTree b
acc' Array (Tree a)
arr
        foldrTree (Leaf Array a
arr) b
acc' = (a -> b -> b) -> b -> Array a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> b -> b
f b
acc' Array a
arr
    {-# INLINE foldr #-}

    foldl :: (b -> a -> b) -> b -> Vector a -> b
foldl b -> a -> b
f b
acc = Vector a -> b
go
      where
        go :: Vector a -> b
go Vector a
Empty = b
acc
        go (Root Int
_ Int
_ Tree a
tree) = b -> Tree a -> b
foldlTree b
acc Tree a
tree

        foldlTree :: b -> Tree a -> b
foldlTree b
acc' (Balanced Array (Tree a)
arr) = (b -> Tree a -> b) -> b -> Array (Tree a) -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl b -> Tree a -> b
foldlTree b
acc' Array (Tree a)
arr
        foldlTree b
acc' (Unbalanced Array (Tree a)
arr PrimArray Int
_) = (b -> Tree a -> b) -> b -> Array (Tree a) -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl b -> Tree a -> b
foldlTree b
acc' Array (Tree a)
arr
        foldlTree b
acc' (Leaf Array a
arr) = (b -> a -> b) -> b -> Array a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl b -> a -> b
f b
acc' Array a
arr
    {-# INLINE foldl #-}

    foldr' :: (a -> b -> b) -> b -> Vector a -> b
foldr' a -> b -> b
f b
acc = Vector a -> b
go
      where
        go :: Vector a -> b
go Vector a
Empty = b
acc
        go (Root Int
_ Int
_ Tree a
tree) = Tree a -> b -> b
foldrTree' Tree a
tree b
acc

        foldrTree' :: Tree a -> b -> b
foldrTree' (Balanced Array (Tree a)
arr) b
acc' = (Tree a -> b -> b) -> b -> Array (Tree a) -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr' Tree a -> b -> b
foldrTree' b
acc' Array (Tree a)
arr
        foldrTree' (Unbalanced Array (Tree a)
arr PrimArray Int
_) b
acc' = (Tree a -> b -> b) -> b -> Array (Tree a) -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr' Tree a -> b -> b
foldrTree' b
acc' Array (Tree a)
arr
        foldrTree' (Leaf Array a
arr) b
acc' = (a -> b -> b) -> b -> Array a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr' a -> b -> b
f b
acc' Array a
arr
    {-# INLINE foldr' #-}

    foldl' :: (b -> a -> b) -> b -> Vector a -> b
foldl' b -> a -> b
f b
acc = Vector a -> b
go
      where
        go :: Vector a -> b
go Vector a
Empty = b
acc
        go (Root Int
_ Int
_ Tree a
tree) = b -> Tree a -> b
foldlTree' b
acc Tree a
tree

        foldlTree' :: b -> Tree a -> b
foldlTree' b
acc' (Balanced Array (Tree a)
arr) = (b -> Tree a -> b) -> b -> Array (Tree a) -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> Tree a -> b
foldlTree' b
acc' Array (Tree a)
arr
        foldlTree' b
acc' (Unbalanced Array (Tree a)
arr PrimArray Int
_) = (b -> Tree a -> b) -> b -> Array (Tree a) -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> Tree a -> b
foldlTree' b
acc' Array (Tree a)
arr
        foldlTree' b
acc' (Leaf Array a
arr) = (b -> a -> b) -> b -> Array a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> a -> b
f b
acc' Array a
arr
    {-# INLINE foldl' #-}

    null :: Vector a -> Bool
null Vector a
Empty = Bool
True
    null Root{} = Bool
False
    {-# INLINE null #-}

    length :: Vector a -> Int
length Vector a
Empty = Int
0
    length (Root Int
s Int
_ Tree a
_) = Int
s
    {-# INLINE length #-}

instance FoldableWithIndex Int Vector where
    ifoldr :: (Int -> a -> b -> b) -> b -> Vector a -> b
ifoldr Int -> a -> b -> b
f b
z0 Vector a
v = (a -> (Int -> b) -> Int -> b) -> (Int -> b) -> Vector a -> Int -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
x Int -> b
g !Int
i -> Int -> a -> b -> b
f Int
i a
x (Int -> b
g (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))) (b -> Int -> b
forall a b. a -> b -> a
const b
z0) Vector a
v Int
0

    ifoldl :: (Int -> b -> a -> b) -> b -> Vector a -> b
ifoldl Int -> b -> a -> b
f b
z0 Vector a
v = ((Int -> b) -> a -> Int -> b) -> (Int -> b) -> Vector a -> Int -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Int -> b
g a
x !Int
i -> Int -> b -> a -> b
f Int
i (Int -> b
g (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) a
x) (b -> Int -> b
forall a b. a -> b -> a
const b
z0) Vector a
v (Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

instance Functor Vector where
    fmap :: (a -> b) -> Vector a -> Vector b
fmap = (a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
map
    a
x <$ :: a -> Vector b -> Vector a
<$ Vector b
v = [a] -> Vector a
forall a. [a] -> Vector a
fromList (Int -> a -> [a]
forall a. Int -> a -> [a]
replicate (Vector b -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector b
v) a
x)

instance FunctorWithIndex Int Vector where
    imap :: (Int -> a -> b) -> Vector a -> Vector b
imap Int -> a -> b
f Vector a
v = Identity (Vector b) -> Vector b
forall a. Identity a -> a
runIdentity (Identity (Vector b) -> Vector b)
-> Identity (Vector b) -> Vector b
forall a b. (a -> b) -> a -> b
$ Indexed Identity (Vector b) -> Int -> Identity (Vector b)
forall (f :: * -> *) a. Indexed f a -> Int -> f a
evalIndexed ((a -> Indexed Identity b)
-> Vector a -> Indexed Identity (Vector b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((Int -> WithIndex (Identity b)) -> Indexed Identity b
forall (f :: * -> *) a. (Int -> WithIndex (f a)) -> Indexed f a
Indexed ((Int -> WithIndex (Identity b)) -> Indexed Identity b)
-> (a -> Int -> WithIndex (Identity b)) -> a -> Indexed Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int -> WithIndex (Identity b)
f') Vector a
v) Int
0
      where
        f' :: a -> Int -> WithIndex (Identity b)
f' a
x Int
i = Int
i Int -> WithIndex (Identity b) -> WithIndex (Identity b)
`seq` Int -> Identity b -> WithIndex (Identity b)
forall a. Int -> a -> WithIndex a
WithIndex (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (b -> Identity b
forall a. a -> Identity a
Identity (Int -> a -> b
f Int
i a
x))

instance Traversable Vector where
    traverse :: (a -> f b) -> Vector a -> f (Vector b)
traverse a -> f b
_ Vector a
Empty = Vector b -> f (Vector b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector b
forall a. Vector a
Empty
    traverse a -> f b
f (Root Int
size Int
sh Tree a
tree) = Int -> Int -> Tree b -> Vector b
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (Tree b -> Vector b) -> f (Tree b) -> f (Vector b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree a -> f (Tree b)
traverseTree Tree a
tree
      where
        traverseTree :: Tree a -> f (Tree b)
traverseTree (Balanced Array (Tree a)
arr) = Array (Tree b) -> Tree b
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree b) -> Tree b) -> f (Array (Tree b)) -> f (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Tree a -> f (Tree b)) -> Array (Tree a) -> f (Array (Tree b))
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array a -> f (Array b)
A.traverse' Tree a -> f (Tree b)
traverseTree Array (Tree a)
arr
        traverseTree (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = Array (Tree b) -> PrimArray Int -> Tree b
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (Array (Tree b) -> PrimArray Int -> Tree b)
-> f (Array (Tree b)) -> f (PrimArray Int -> Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Tree a -> f (Tree b)) -> Array (Tree a) -> f (Array (Tree b))
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array a -> f (Array b)
A.traverse' Tree a -> f (Tree b)
traverseTree Array (Tree a)
arr f (PrimArray Int -> Tree b) -> f (PrimArray Int) -> f (Tree b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> PrimArray Int -> f (PrimArray Int)
forall (f :: * -> *) a. Applicative f => a -> f a
pure PrimArray Int
sizes
        traverseTree (Leaf Array a
arr) = Array b -> Tree b
forall a. Array a -> Tree a
Leaf (Array b -> Tree b) -> f (Array b) -> f (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Array a -> f (Array b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array a -> f (Array b)
A.traverse a -> f b
f Array a
arr

instance TraversableWithIndex Int Vector where
    itraverse :: (Int -> a -> f b) -> Vector a -> f (Vector b)
itraverse Int -> a -> f b
f Vector a
v = Indexed f (Vector b) -> Int -> f (Vector b)
forall (f :: * -> *) a. Indexed f a -> Int -> f a
evalIndexed ((a -> Indexed f b) -> Vector a -> Indexed f (Vector b)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((Int -> WithIndex (f b)) -> Indexed f b
forall (f :: * -> *) a. (Int -> WithIndex (f a)) -> Indexed f a
Indexed ((Int -> WithIndex (f b)) -> Indexed f b)
-> (a -> Int -> WithIndex (f b)) -> a -> Indexed f b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int -> WithIndex (f b)
f') Vector a
v) Int
0
      where
        f' :: a -> Int -> WithIndex (f b)
f' a
x Int
i = Int
i Int -> WithIndex (f b) -> WithIndex (f b)
`seq` Int -> f b -> WithIndex (f b)
forall a. Int -> a -> WithIndex a
WithIndex (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> a -> f b
f Int
i a
x)

instance Applicative Vector where
    pure :: a -> Vector a
pure = a -> Vector a
forall a. a -> Vector a
singleton
    Vector (a -> b)
fs <*> :: Vector (a -> b) -> Vector a -> Vector b
<*> Vector a
xs = (Vector b -> (a -> b) -> Vector b)
-> Vector b -> Vector (a -> b) -> Vector b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Vector b
acc a -> b
f -> Vector b
acc Vector b -> Vector b -> Vector b
forall a. Vector a -> Vector a -> Vector a
>< (a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
map a -> b
f Vector a
xs) Vector b
forall a. Vector a
empty Vector (a -> b)
fs
    liftA2 :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
liftA2 a -> b -> c
f Vector a
xs Vector b
ys = (Vector c -> a -> Vector c) -> Vector c -> Vector a -> Vector c
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Vector c
acc a
x -> Vector c
acc Vector c -> Vector c -> Vector c
forall a. Vector a -> Vector a -> Vector a
>< (b -> c) -> Vector b -> Vector c
forall a b. (a -> b) -> Vector a -> Vector b
map (a -> b -> c
f a
x) Vector b
ys) Vector c
forall a. Vector a
empty Vector a
xs
    Vector a
xs *> :: Vector a -> Vector b -> Vector b
*> Vector b
ys = (Vector b -> a -> Vector b) -> Vector b -> Vector a -> Vector b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Vector b
acc a
_ -> Vector b
acc Vector b -> Vector b -> Vector b
forall a. Vector a -> Vector a -> Vector a
>< Vector b
ys) Vector b
forall a. Vector a
empty Vector a
xs
    Vector a
xs <* :: Vector a -> Vector b -> Vector a
<* Vector b
ys = (Vector a -> a -> Vector a) -> Vector a -> Vector a -> Vector a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Vector a
acc a
x -> Vector a
acc Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
>< [a] -> Vector a
forall a. [a] -> Vector a
fromList (Int -> a -> [a]
forall a. Int -> a -> [a]
replicate (Vector b -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector b
ys) a
x)) Vector a
forall a. Vector a
empty Vector a
xs

instance Monad Vector where
    Vector a
xs >>= :: Vector a -> (a -> Vector b) -> Vector b
>>= a -> Vector b
f = (Vector b -> a -> Vector b) -> Vector b -> Vector a -> Vector b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Vector b
acc a
x -> Vector b
acc Vector b -> Vector b -> Vector b
forall a. Vector a -> Vector a -> Vector a
>< a -> Vector b
f a
x) Vector b
forall a. Vector a
empty Vector a
xs

instance Alternative Vector where
    empty :: Vector a
empty = Vector a
forall a. Vector a
empty
    <|> :: Vector a -> Vector a -> Vector a
(<|>) = Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
(><)

instance MonadPlus Vector

instance MonadFail Vector where
    fail :: [Char] -> Vector a
fail [Char]
_ = Vector a
forall a. Vector a
empty

instance MonadFix Vector where
    mfix :: (a -> Vector a) -> Vector a
mfix a -> Vector a
f = [a] -> Vector a
forall a. [a] -> Vector a
fromList ([a] -> Vector a) -> [a] -> Vector a
forall a b. (a -> b) -> a -> b
$ (Int -> a) -> [Int] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
i -> let x :: a
x = Int -> Vector a -> a
forall a. HasCallStack => Int -> Vector a -> a
index Int
i (a -> Vector a
f a
x) in a
x) [Int
0..Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (a -> Vector a
f a
forall a. a
err) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      where
        err :: a
err = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"mfix for Data.RRBVector.Vector applied to strict function"

instance MonadZip Vector where
    mzipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
mzipWith = (a -> b -> c) -> Vector a -> Vector b -> Vector c
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith
    mzip :: Vector a -> Vector b -> Vector (a, b)
mzip = Vector a -> Vector b -> Vector (a, b)
forall a b. Vector a -> Vector b -> Vector (a, b)
zip
    munzip :: Vector (a, b) -> (Vector a, Vector b)
munzip = Vector (a, b) -> (Vector a, Vector b)
forall a b. Vector (a, b) -> (Vector a, Vector b)
unzip

instance Exts.IsList (Vector a) where
    type Item (Vector a) = a
    fromList :: [Item (Vector a)] -> Vector a
fromList = [Item (Vector a)] -> Vector a
forall a. [a] -> Vector a
fromList
    toList :: Vector a -> [Item (Vector a)]
toList = Vector a -> [Item (Vector a)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList

instance (a ~ Char) => Exts.IsString (Vector a) where
    fromString :: [Char] -> Vector a
fromString = [Char] -> Vector a
forall a. [a] -> Vector a
fromList

instance (NFData a) => NFData (Vector a) where
    rnf :: Vector a -> ()
rnf = Vector a -> ()
forall (f :: * -> *) a. (NFData1 f, NFData a) => f a -> ()
rnf1

instance NFData1 Vector where
    liftRnf :: (a -> ()) -> Vector a -> ()
liftRnf a -> ()
f = (() -> a -> ()) -> () -> Vector a -> ()
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\()
_ a
x -> a -> ()
f a
x) ()

-- | \(O(1)\). The empty vector.
--
-- > empty = fromList []
empty :: Vector a
empty :: Vector a
empty = Vector a
forall a. Vector a
Empty

-- | \(O(1)\). A vector with a single element.
--
-- > singleton x = fromList [x]
singleton :: a -> Vector a
singleton :: a -> Vector a
singleton a
x = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root Int
1 Int
0 (Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ a -> Array a
forall a. a -> Array a
A.singleton a
x)

-- | \(O(n)\). Create a new vector from a list.
fromList :: [a] -> Vector a
fromList :: [a] -> Vector a
fromList [] = Vector a
forall a. Vector a
Empty
fromList [a
x] = a -> Vector a
forall a. a -> Vector a
singleton a
x
fromList [a]
ls = case (Array a -> Tree a) -> [a] -> [Tree a]
forall a a. (Array a -> a) -> [a] -> [a]
nodes Array a -> Tree a
forall a. Array a -> Tree a
Leaf [a]
ls of
    [Tree a
tree] -> Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize Int
0 Tree a
tree) Int
0 Tree a
tree -- tree is a single leaf
    [Tree a]
ls' -> Int -> [Tree a] -> Vector a
forall a. Int -> [Tree a] -> Vector a
iterateNodes Int
blockShift [Tree a]
ls'
  where
    nodes :: (Array a -> a) -> [a] -> [a]
nodes Array a -> a
f [a]
trees = (forall s. ST s [a]) -> [a]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [a]) -> [a]) -> (forall s. ST s [a]) -> [a]
forall a b. (a -> b) -> a -> b
$ do
        Buffer s a
buffer <- Int -> ST s (Buffer s a)
forall s a. Int -> ST s (Buffer s a)
Buffer.new Int
blockSize
        let loop :: [a] -> ST s [a]
loop [] = do
                Array a
result <- Buffer s a -> ST s (Array a)
forall s a. Buffer s a -> ST s (Array a)
Buffer.get Buffer s a
buffer
                [a] -> ST s [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Array a -> a
f Array a
result]
            loop (a
t : [a]
ts) = do
                Int
size <- Buffer s a -> ST s Int
forall s a. Buffer s a -> ST s Int
Buffer.size Buffer s a
buffer
                if Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blockSize then do
                    Array a
result <- Buffer s a -> ST s (Array a)
forall s a. Buffer s a -> ST s (Array a)
Buffer.get Buffer s a
buffer
                    Buffer s a -> a -> ST s ()
forall s a. Buffer s a -> a -> ST s ()
Buffer.push Buffer s a
buffer a
t
                    [a]
rest <- [a] -> ST s [a]
loop [a]
ts
                    [a] -> ST s [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array a -> a
f Array a
result a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
rest)
                else do
                    Buffer s a -> a -> ST s ()
forall s a. Buffer s a -> a -> ST s ()
Buffer.push Buffer s a
buffer a
t
                    [a] -> ST s [a]
loop [a]
ts
        [a] -> ST s [a]
loop [a]
trees
    {-# INLINE nodes #-}

    iterateNodes :: Int -> [Tree a] -> Vector a
iterateNodes Int
sh [Tree a]
trees = case (Array (Tree a) -> Tree a) -> [Tree a] -> [Tree a]
forall a a. (Array a -> a) -> [a] -> [a]
nodes Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced [Tree a]
trees of
        [Tree a
tree] -> Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize Int
sh Tree a
tree) Int
sh Tree a
tree
        [Tree a]
trees' -> Int -> [Tree a] -> Vector a
iterateNodes (Int -> Int
up Int
sh) [Tree a]
trees'

-- | \(O(\log n)\). The element at the index or 'Nothing' if the index is out of range.
lookup :: Int -> Vector a -> Maybe a
lookup :: Int -> Vector a -> Maybe a
lookup Int
_ Vector a
Empty = Maybe a
forall a. Maybe a
Nothing
lookup Int
i (Root Int
size Int
sh Tree a
tree)
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
size = Maybe a
forall a. Maybe a
Nothing  -- index out of range
    | Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tree a -> a
forall p. Int -> Int -> Tree p -> p
lookupTree Int
i Int
sh Tree a
tree
  where
    lookupTree :: Int -> Int -> Tree p -> p
lookupTree Int
i Int
sh (Balanced Array (Tree p)
arr) = Int -> Int -> Tree p -> p
lookupTree Int
i (Int -> Int
down Int
sh) (Array (Tree p) -> Int -> Tree p
forall a. Array a -> Int -> a
A.index Array (Tree p)
arr (Int -> Int -> Int
radixIndex Int
i Int
sh))
    lookupTree Int
i Int
sh (Unbalanced Array (Tree p)
arr PrimArray Int
sizes) =
        let (Int
idx, Int
subIdx) = PrimArray Int -> Int -> Int -> (Int, Int)
relaxedRadixIndex PrimArray Int
sizes Int
i Int
sh
        in Int -> Int -> Tree p -> p
lookupTree Int
subIdx (Int -> Int
down Int
sh) (Array (Tree p) -> Int -> Tree p
forall a. Array a -> Int -> a
A.index Array (Tree p)
arr Int
idx)
    lookupTree Int
i Int
_ (Leaf Array p
arr) = Array p -> Int -> p
forall a. Array a -> Int -> a
A.index Array p
arr (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
blockMask)

-- | \(O(\log n)\). The element at the index. Calls 'error' if the index is out of range.
index :: HasCallStack => Int -> Vector a -> a
index :: Int -> Vector a -> a
index Int
i = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"AMT.index: index out of range") (Maybe a -> a) -> (Vector a -> Maybe a) -> Vector a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Vector a -> Maybe a
forall a. Int -> Vector a -> Maybe a
lookup Int
i

-- | \(O(\log n)\). A flipped version of 'lookup'.
(!?) :: Vector a -> Int -> Maybe a
!? :: Vector a -> Int -> Maybe a
(!?) = (Int -> Vector a -> Maybe a) -> Vector a -> Int -> Maybe a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Vector a -> Maybe a
forall a. Int -> Vector a -> Maybe a
lookup

-- | \(O(\log n)\). A flipped version of 'index'.
(!) :: HasCallStack => Vector a -> Int -> a
(!) = (Int -> Vector a -> a) -> Vector a -> Int -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Vector a -> a
forall a. HasCallStack => Int -> Vector a -> a
index

-- | \(O(\log n)\). Update the element at the index with a new element.
-- If the index is out of range, the original vector is returned.
update :: Int -> a -> Vector a -> Vector a
update :: Int -> a -> Vector a -> Vector a
update Int
_ a
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
update Int
i a
x v :: Vector a
v@(Root Int
size Int
sh Tree a
tree)
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
size = Vector a
v  -- index out of range
    | Bool
otherwise = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (Int -> Int -> Tree a -> Tree a
adjustTree Int
i Int
sh Tree a
tree)
  where
    adjustTree :: Int -> Int -> Tree a -> Tree a
adjustTree Int
i Int
sh (Balanced Array (Tree a)
arr) = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (Int -> Int -> Int
radixIndex Int
i Int
sh) (Int -> Int -> Tree a -> Tree a
adjustTree Int
i (Int -> Int
down Int
sh)))
    adjustTree Int
i Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) =
        let (Int
idx, Int
subIdx) = PrimArray Int -> Int -> Int -> (Int, Int)
relaxedRadixIndex PrimArray Int
sizes Int
i Int
sh
        in Array (Tree a) -> PrimArray Int -> Tree a
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Int
idx (Int -> Int -> Tree a -> Tree a
adjustTree Int
subIdx (Int -> Int
down Int
sh))) PrimArray Int
sizes
    adjustTree Int
i Int
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Int -> a -> Array a
forall a. Array a -> Int -> a -> Array a
A.update Array a
arr (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
blockMask) a
x)

-- | \(O(\log n)\). Adjust the element at the index by applying the function to it.
-- If the index is out of range, the original vector is returned.
adjust :: Int -> (a -> a) -> Vector a -> Vector a
adjust :: Int -> (a -> a) -> Vector a -> Vector a
adjust Int
_ a -> a
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
adjust Int
i a -> a
f v :: Vector a
v@(Root Int
size Int
sh Tree a
tree)
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
size = Vector a
v  -- index out of range
    | Bool
otherwise = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (Int -> Int -> Tree a -> Tree a
adjustTree Int
i Int
sh Tree a
tree)
  where
    adjustTree :: Int -> Int -> Tree a -> Tree a
adjustTree Int
i Int
sh (Balanced Array (Tree a)
arr) = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (Int -> Int -> Int
radixIndex Int
i Int
sh) (Int -> Int -> Tree a -> Tree a
adjustTree Int
i (Int -> Int
down Int
sh)))
    adjustTree Int
i Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) =
        let (Int
idx, Int
subIdx) = PrimArray Int -> Int -> Int -> (Int, Int)
relaxedRadixIndex PrimArray Int
sizes Int
i Int
sh
        in Array (Tree a) -> PrimArray Int -> Tree a
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Int
idx (Int -> Int -> Tree a -> Tree a
adjustTree Int
subIdx (Int -> Int
down Int
sh))) PrimArray Int
sizes
    adjustTree Int
i Int
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Int -> (a -> a) -> Array a
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust Array a
arr (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
blockMask) a -> a
f)

-- | \(O(\log n)\). Like 'adjust', but the result of the function is forced.
adjust' :: Int -> (a -> a) -> Vector a -> Vector a
adjust' :: Int -> (a -> a) -> Vector a -> Vector a
adjust' Int
_ a -> a
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
adjust' Int
i a -> a
f v :: Vector a
v@(Root Int
size Int
sh Tree a
tree)
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
size = Vector a
v  -- index out of range
    | Bool
otherwise = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (Int -> Int -> Tree a -> Tree a
adjustTree Int
i Int
sh Tree a
tree)
  where
    adjustTree :: Int -> Int -> Tree a -> Tree a
adjustTree Int
i Int
sh (Balanced Array (Tree a)
arr) = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (Int -> Int -> Int
radixIndex Int
i Int
sh) (Int -> Int -> Tree a -> Tree a
adjustTree Int
i (Int -> Int
down Int
sh)))
    adjustTree Int
i Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) =
        let (Int
idx, Int
subIdx) = PrimArray Int -> Int -> Int -> (Int, Int)
relaxedRadixIndex PrimArray Int
sizes Int
i Int
sh
        in Array (Tree a) -> PrimArray Int -> Tree a
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Int
idx (Int -> Int -> Tree a -> Tree a
adjustTree Int
subIdx (Int -> Int
down Int
sh))) PrimArray Int
sizes
    adjustTree Int
i Int
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Int -> (a -> a) -> Array a
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array a
arr (Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
blockMask) a -> a
f)

-- | \(O(n)\). Apply the function to every element.
--
-- >>> map (+ 1) (fromList [1, 2, 3])
-- fromList [2,3,4]
map :: (a -> b) -> Vector a -> Vector b
map :: (a -> b) -> Vector a -> Vector b
map a -> b
_ Vector a
Empty = Vector b
forall a. Vector a
Empty
map a -> b
f (Root Int
size Int
sh Tree a
tree) = Int -> Int -> Tree b -> Vector b
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (Tree a -> Tree b
mapTree Tree a
tree)
  where
    mapTree :: Tree a -> Tree b
mapTree (Balanced Array (Tree a)
arr) = Array (Tree b) -> Tree b
forall a. Array (Tree a) -> Tree a
Balanced ((Tree a -> Tree b) -> Array (Tree a) -> Array (Tree b)
forall a b. (a -> b) -> Array a -> Array b
A.map' Tree a -> Tree b
mapTree Array (Tree a)
arr)
    mapTree (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = Array (Tree b) -> PrimArray Int -> Tree b
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced ((Tree a -> Tree b) -> Array (Tree a) -> Array (Tree b)
forall a b. (a -> b) -> Array a -> Array b
A.map' Tree a -> Tree b
mapTree Array (Tree a)
arr) PrimArray Int
sizes
    mapTree (Leaf Array a
arr) = Array b -> Tree b
forall a. Array a -> Tree a
Leaf ((a -> b) -> Array a -> Array b
forall a b. (a -> b) -> Array a -> Array b
A.map a -> b
f Array a
arr)

-- | \(O(n)\). Reverse the vector.
--
-- >>> reverse (fromList [1, 2, 3])
-- fromList [3,2,1]
reverse :: Vector a -> Vector a
reverse :: Vector a -> Vector a
reverse = [a] -> Vector a
forall a. [a] -> Vector a
fromList ([a] -> Vector a) -> (Vector a -> [a]) -> Vector a -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> a -> [a]) -> [a] -> Vector a -> [a]
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((a -> [a] -> [a]) -> [a] -> a -> [a]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) [] -- convert the vector to a reverse list and then rebuild

-- | \(O(\min(n_1, n_2))\). Take two vectors and return a vector of corresponding pairs.
-- If one input is longer, excess elements are discarded from the right end.
--
-- > zip = zipWith (,)
zip :: Vector a -> Vector b -> Vector (a, b)
zip :: Vector a -> Vector b -> Vector (a, b)
zip Vector a
v1 Vector b
v2 = [(a, b)] -> Vector (a, b)
forall a. [a] -> Vector a
fromList ([(a, b)] -> Vector (a, b)) -> [(a, b)] -> Vector (a, b)
forall a b. (a -> b) -> a -> b
$ [a] -> [b] -> [(a, b)]
forall a b. [a] -> [b] -> [(a, b)]
List.zip (Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (Vector b -> [b]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector b
v2)

-- | \(O(\min(n_1, n_2))\). 'zipWith' generalizes 'zip' by zipping with the function.
zipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith :: (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith a -> b -> c
f Vector a
v1 Vector b
v2 = [c] -> Vector c
forall a. [a] -> Vector a
fromList ([c] -> Vector c) -> [c] -> Vector c
forall a b. (a -> b) -> a -> b
$ (a -> b -> c) -> [a] -> [b] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith a -> b -> c
f (Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (Vector b -> [b]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector b
v2)

-- | \(O(n)\). Unzip a vector of pairs.
--
-- >>> unzip (fromList [(1, "a"), (2, "b"), (3, "c")])
-- (fromList [1,2,3],fromList ["a","b","c"])
unzip :: Vector (a, b) -> (Vector a, Vector b)
unzip :: Vector (a, b) -> (Vector a, Vector b)
unzip Vector (a, b)
v = (((a, b) -> a) -> Vector (a, b) -> Vector a
forall a b. (a -> b) -> Vector a -> Vector b
map (a, b) -> a
forall a b. (a, b) -> a
fst Vector (a, b)
v, ((a, b) -> b) -> Vector (a, b) -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
map (a, b) -> b
forall a b. (a, b) -> b
snd Vector (a, b)
v)

-- | \(O(\log n)\). The first element and the vector without the first element, or 'Nothing' if the vector is empty.
--
-- >>> viewl (fromList [1, 2, 3])
-- Just (1,fromList [2,3])
viewl :: Vector a -> Maybe (a, Vector a)
viewl :: Vector a -> Maybe (a, Vector a)
viewl Vector a
Empty = Maybe (a, Vector a)
forall a. Maybe a
Nothing
viewl v :: Vector a
v@(Root Int
_ Int
_ Tree a
tree) = let !tail :: Vector a
tail = Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
drop Int
1 Vector a
v in (a, Vector a) -> Maybe (a, Vector a)
forall a. a -> Maybe a
Just (Tree a -> a
forall p. Tree p -> p
headTree Tree a
tree, Vector a
tail)
  where
    headTree :: Tree p -> p
headTree (Balanced Array (Tree p)
arr) = Tree p -> p
headTree (Array (Tree p) -> Tree p
forall a. Array a -> a
A.head Array (Tree p)
arr)
    headTree (Unbalanced Array (Tree p)
arr PrimArray Int
_) = Tree p -> p
headTree (Array (Tree p) -> Tree p
forall a. Array a -> a
A.head Array (Tree p)
arr)
    headTree (Leaf Array p
arr) = Array p -> p
forall a. Array a -> a
A.head Array p
arr

-- | \(O(\log n)\). The vector without the last element and the last element, or 'Nothing' if the vector is empty.
--
-- >>> viewr (fromList [1, 2, 3])
-- Just (fromList [1,2],3)
viewr :: Vector a -> Maybe (Vector a, a)
viewr :: Vector a -> Maybe (Vector a, a)
viewr Vector a
Empty = Maybe (Vector a, a)
forall a. Maybe a
Nothing
viewr v :: Vector a
v@(Root Int
size Int
_ Tree a
tree) = let !init :: Vector a
init = Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
take (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Vector a
v in (Vector a, a) -> Maybe (Vector a, a)
forall a. a -> Maybe a
Just (Vector a
init, Tree a -> a
forall p. Tree p -> p
lastTree Tree a
tree)
  where
    lastTree :: Tree p -> p
lastTree (Balanced Array (Tree p)
arr) = Tree p -> p
lastTree (Array (Tree p) -> Tree p
forall a. Array a -> a
A.last Array (Tree p)
arr)
    lastTree (Unbalanced Array (Tree p)
arr PrimArray Int
_) = Tree p -> p
lastTree (Array (Tree p) -> Tree p
forall a. Array a -> a
A.last Array (Tree p)
arr)
    lastTree (Leaf Array p
arr) = Array p -> p
forall a. Array a -> a
A.last Array p
arr

-- | \(O(\log n)\). Split the vector at the given index.
--
-- > splitAt n v = (take n v, drop n v)
splitAt :: Int -> Vector a -> (Vector a, Vector a)
splitAt :: Int -> Vector a -> (Vector a, Vector a)
splitAt Int
n Vector a
v =
    let !left :: Vector a
left = Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
take Int
n Vector a
v
        !right :: Vector a
right = Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
drop Int
n Vector a
v
    in (Vector a
left, Vector a
right)

-- | \(O(\log n)\). Insert an element at the given index.
insertAt :: Int -> a -> Vector a -> Vector a
insertAt :: Int -> a -> Vector a -> Vector a
insertAt Int
i a
x Vector a
v = let (Vector a
left, Vector a
right) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
splitAt Int
i Vector a
v in (Vector a
left Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
|> a
x) Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
>< Vector a
right

-- | \(O(\log n)\). Delete the element at the given index.
deleteAt :: Int -> Vector a -> Vector a
deleteAt :: Int -> Vector a -> Vector a
deleteAt Int
i Vector a
v = let (Vector a
left, Vector a
right) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
splitAt (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Vector a
v in Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
take Int
i Vector a
left Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
>< Vector a
right

-- concatenation

-- | \(O(\log \max(n_1, n_2))\). Concatenates two vectors.
--
-- >>> fromList [1, 2, 3] >< fromList [4, 5]
-- fromList [1,2,3,4,5]
(><) :: Vector a -> Vector a -> Vector a
Vector a
Empty >< :: Vector a -> Vector a -> Vector a
>< Vector a
v = Vector a
v
Vector a
v >< Vector a
Empty = Vector a
v
Root Int
size1 Int
sh1 Tree a
tree1 >< Root Int
size2 Int
sh2 Tree a
tree2 =
    let maxShift :: Int
maxShift = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
sh1 Int
sh2
        newTree :: Tree a
newTree = Tree a -> Int -> Tree a -> Int -> Tree a
forall a. Tree a -> Int -> Tree a -> Int -> Tree a
mergeTrees Tree a
tree1 Int
sh1 Tree a
tree2 Int
sh2
    in case Tree a -> Maybe (Tree a)
forall a. Tree a -> Maybe (Tree a)
singleTree Tree a
newTree of
        Just Tree a
newTree -> Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
size2) Int
maxShift Tree a
newTree
        Maybe (Tree a)
Nothing -> Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
size2) (Int -> Int
up Int
maxShift) Tree a
newTree
  where
    mergeTrees :: Tree a -> Int -> Tree a -> Int -> Tree a
mergeTrees (Leaf Array a
arr1) Int
_ (Leaf Array a
arr2) Int
_ = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$
        if Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blockSize then Tree a -> Tree a -> Array (Tree a)
forall a. a -> a -> Array a
A.from2 (Array a -> Tree a
forall a. Array a -> Tree a
Leaf Array a
arr1) (Array a -> Tree a
forall a. Array a -> Tree a
Leaf Array a
arr2)
        else if Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
blockSize then Tree a -> Array (Tree a)
forall a. a -> Array a
A.singleton (Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a
arr1 Array a -> Array a -> Array a
forall a. Semigroup a => a -> a -> a
<> Array a
arr2))
        else
            let (Array a
left, Array a
right) = Array a -> Int -> (Array a, Array a)
forall a. Array a -> Int -> (Array a, Array a)
A.splitAt (Array a
arr1 Array a -> Array a -> Array a
forall a. Semigroup a => a -> a -> a
<> Array a
arr2) Int
blockSize
            in Tree a -> Tree a -> Array (Tree a)
forall a. a -> a -> Array a
A.from2 (Array a -> Tree a
forall a. Array a -> Tree a
Leaf Array a
left) (Array a -> Tree a
forall a. Array a -> Tree a
Leaf Array a
right)
    mergeTrees Tree a
tree1 Int
sh1 Tree a
tree2 Int
sh2 = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
sh1 Int
sh2 of
        Ordering
LT ->
            let right :: Array (Tree a)
right = Tree a -> Array (Tree a)
forall a. Tree a -> Array (Tree a)
treeToArray Tree a
tree2
                (Tree a
rightHead, Array (Tree a)
rightTail) = Array (Tree a) -> (Tree a, Array (Tree a))
forall a. Array a -> (a, Array a)
viewl Array (Tree a)
right
                merged :: Tree a
merged = Tree a -> Int -> Tree a -> Int -> Tree a
mergeTrees Tree a
tree1 Int
sh1 Tree a
rightHead (Int -> Int
down Int
sh2)
            in Int -> Array (Tree a) -> Array (Tree a) -> Array (Tree a) -> Tree a
forall a.
Int -> Array (Tree a) -> Array (Tree a) -> Array (Tree a) -> Tree a
mergeRebalance Int
sh2 Array (Tree a)
forall a. Array a
A.empty (Tree a -> Array (Tree a)
forall a. Tree a -> Array (Tree a)
treeToArray Tree a
merged) Array (Tree a)
rightTail
        Ordering
GT ->
            let left :: Array (Tree a)
left = Tree a -> Array (Tree a)
forall a. Tree a -> Array (Tree a)
treeToArray Tree a
tree1
                (Array (Tree a)
leftInit, Tree a
leftLast) = Array (Tree a) -> (Array (Tree a), Tree a)
forall b. Array b -> (Array b, b)
viewr Array (Tree a)
left
                merged :: Tree a
merged = Tree a -> Int -> Tree a -> Int -> Tree a
mergeTrees Tree a
leftLast (Int -> Int
down Int
sh1) Tree a
tree2 Int
sh2
            in Int -> Array (Tree a) -> Array (Tree a) -> Array (Tree a) -> Tree a
forall a.
Int -> Array (Tree a) -> Array (Tree a) -> Array (Tree a) -> Tree a
mergeRebalance Int
sh1 Array (Tree a)
leftInit (Tree a -> Array (Tree a)
forall a. Tree a -> Array (Tree a)
treeToArray Tree a
merged) Array (Tree a)
forall a. Array a
A.empty
        Ordering
EQ ->
            let left :: Array (Tree a)
left = Tree a -> Array (Tree a)
forall a. Tree a -> Array (Tree a)
treeToArray Tree a
tree1
                right :: Array (Tree a)
right = Tree a -> Array (Tree a)
forall a. Tree a -> Array (Tree a)
treeToArray Tree a
tree2
                (Array (Tree a)
leftInit, Tree a
leftLast) = Array (Tree a) -> (Array (Tree a), Tree a)
forall b. Array b -> (Array b, b)
viewr Array (Tree a)
left
                (Tree a
rightHead, Array (Tree a)
rightTail) = Array (Tree a) -> (Tree a, Array (Tree a))
forall a. Array a -> (a, Array a)
viewl Array (Tree a)
right
                merged :: Tree a
merged = Tree a -> Int -> Tree a -> Int -> Tree a
mergeTrees Tree a
leftLast (Int -> Int
down Int
sh1) Tree a
rightHead (Int -> Int
down Int
sh2)
            in Int -> Array (Tree a) -> Array (Tree a) -> Array (Tree a) -> Tree a
forall a.
Int -> Array (Tree a) -> Array (Tree a) -> Array (Tree a) -> Tree a
mergeRebalance Int
sh1 Array (Tree a)
leftInit (Tree a -> Array (Tree a)
forall a. Tree a -> Array (Tree a)
treeToArray Tree a
merged) Array (Tree a)
rightTail
      where
        viewl :: Array a -> (a, Array a)
viewl Array a
arr = (Array a -> a
forall a. Array a -> a
A.head Array a
arr, Array a -> Int -> Array a
forall a. Array a -> Int -> Array a
A.drop Array a
arr Int
1)
        viewr :: Array b -> (Array b, b)
viewr Array b
arr = (Array b -> Int -> Array b
forall a. Array a -> Int -> Array a
A.take Array b
arr (Array b -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array b
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1), Array b -> b
forall a. Array a -> a
A.last Array b
arr)

    -- the type annotations are necessary to compile
    mergeRebalance :: forall a. Int -> A.Array (Tree a) -> A.Array (Tree a) -> A.Array (Tree a) -> Tree a
    mergeRebalance :: Int -> Array (Tree a) -> Array (Tree a) -> Array (Tree a) -> Tree a
mergeRebalance Int
sh Array (Tree a)
left Array (Tree a)
center Array (Tree a)
right
        | Int
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blockShift = (Tree a -> Array a) -> (Array a -> Tree a) -> Tree a
forall t. (Tree a -> Array t) -> (Array t -> Tree a) -> Tree a
mergeRebalance' (\(Leaf Array a
arr) -> Array a
arr) Array a -> Tree a
forall a. Array a -> Tree a
Leaf
        | Bool
otherwise = (Tree a -> Array (Tree a)) -> (Array (Tree a) -> Tree a) -> Tree a
forall t. (Tree a -> Array t) -> (Array t -> Tree a) -> Tree a
mergeRebalance' Tree a -> Array (Tree a)
forall a. Tree a -> Array (Tree a)
treeToArray (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes (Int -> Int
down Int
sh))
      where
        mergeRebalance' :: (Tree a -> A.Array t) -> (A.Array t -> Tree a) -> Tree a
        mergeRebalance' :: (Tree a -> Array t) -> (Array t -> Tree a) -> Tree a
mergeRebalance' Tree a -> Array t
extract Array t -> Tree a
construct = (forall s. ST s (Tree a)) -> Tree a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Tree a)) -> Tree a)
-> (forall s. ST s (Tree a)) -> Tree a
forall a b. (a -> b) -> a -> b
$ do
            Buffer s (Tree a)
newRoot <- Int -> ST s (Buffer s (Tree a))
forall s a. Int -> ST s (Buffer s a)
Buffer.new Int
blockSize
            Buffer s (Tree a)
newSubtree <- Int -> ST s (Buffer s (Tree a))
forall s a. Int -> ST s (Buffer s a)
Buffer.new Int
blockSize
            Buffer s t
newNode <- Int -> ST s (Buffer s t)
forall s a. Int -> ST s (Buffer s a)
Buffer.new Int
blockSize
            [Tree a] -> (Tree a -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Array (Tree a) -> [Tree a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Array (Tree a)
left [Tree a] -> [Tree a] -> [Tree a]
forall a. [a] -> [a] -> [a]
++ Array (Tree a) -> [Tree a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Array (Tree a)
center [Tree a] -> [Tree a] -> [Tree a]
forall a. [a] -> [a] -> [a]
++ Array (Tree a) -> [Tree a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Array (Tree a)
right) ((Tree a -> ST s ()) -> ST s ()) -> (Tree a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Tree a
subtree ->
                Array t -> (t -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Tree a -> Array t
extract Tree a
subtree) ((t -> ST s ()) -> ST s ()) -> (t -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \t
x -> do
                    Int
lenNode <- Buffer s t -> ST s Int
forall s a. Buffer s a -> ST s Int
Buffer.size Buffer s t
newNode
                    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lenNode Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blockSize) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                        (Array t -> Tree a) -> Buffer s t -> Buffer s (Tree a) -> ST s ()
forall a a s. (Array a -> a) -> Buffer s a -> Buffer s a -> ST s ()
pushTo Array t -> Tree a
construct Buffer s t
newNode Buffer s (Tree a)
newSubtree
                        Int
lenSubtree <- Buffer s (Tree a) -> ST s Int
forall s a. Buffer s a -> ST s Int
Buffer.size Buffer s (Tree a)
newSubtree
                        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lenSubtree Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blockSize) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ (Array (Tree a) -> Tree a)
-> Buffer s (Tree a) -> Buffer s (Tree a) -> ST s ()
forall a a s. (Array a -> a) -> Buffer s a -> Buffer s a -> ST s ()
pushTo (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh) Buffer s (Tree a)
newSubtree Buffer s (Tree a)
newRoot
                    Buffer s t -> t -> ST s ()
forall s a. Buffer s a -> a -> ST s ()
Buffer.push Buffer s t
newNode t
x
            (Array t -> Tree a) -> Buffer s t -> Buffer s (Tree a) -> ST s ()
forall a a s. (Array a -> a) -> Buffer s a -> Buffer s a -> ST s ()
pushTo Array t -> Tree a
construct Buffer s t
newNode Buffer s (Tree a)
newSubtree
            (Array (Tree a) -> Tree a)
-> Buffer s (Tree a) -> Buffer s (Tree a) -> ST s ()
forall a a s. (Array a -> a) -> Buffer s a -> Buffer s a -> ST s ()
pushTo (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh) Buffer s (Tree a)
newSubtree Buffer s (Tree a)
newRoot
            Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes (Int -> Int
up Int
sh) (Array (Tree a) -> Tree a)
-> ST s (Array (Tree a)) -> ST s (Tree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Buffer s (Tree a) -> ST s (Array (Tree a))
forall s a. Buffer s a -> ST s (Array a)
Buffer.get Buffer s (Tree a)
newRoot
        {-# INLINE mergeRebalance' #-}

        pushTo :: (Array a -> a) -> Buffer s a -> Buffer s a -> ST s ()
pushTo Array a -> a
f Buffer s a
from Buffer s a
to = do
            Array a
result <- Buffer s a -> ST s (Array a)
forall s a. Buffer s a -> ST s (Array a)
Buffer.get Buffer s a
from
            Buffer s a -> a -> ST s ()
forall s a. Buffer s a -> a -> ST s ()
Buffer.push Buffer s a
to (Array a -> a
f Array a
result)
        {-# INLINE pushTo #-}

    singleTree :: Tree a -> Maybe (Tree a)
singleTree (Balanced Array (Tree a)
arr)
        | Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Tree a -> Maybe (Tree a)
forall a. a -> Maybe a
Just (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
arr)
    singleTree (Unbalanced Array (Tree a)
arr PrimArray Int
_)
        | Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Tree a -> Maybe (Tree a)
forall a. a -> Maybe a
Just (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
arr)
    singleTree Tree a
_ = Maybe (Tree a)
forall a. Maybe a
Nothing

-- | \(O(\log n)\). Add an element to the left end of the vector.
--
-- >>> 1 <| fromList [2, 3, 4]
-- fromList [1,2,3,4]
(<|) :: a -> Vector a -> Vector a
a
x <| :: a -> Vector a -> Vector a
<| Vector a
Empty = a -> Vector a
forall a. a -> Vector a
singleton a
x
a
x <| Root Int
size Int
sh Tree a
tree
    | Int
insertShift Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
sh = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
insertShift (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
insertShift (Tree a -> Tree a -> Array (Tree a)
forall a. a -> a -> Array a
A.from2 (a -> Int -> Tree a
forall a. a -> Int -> Tree a
newBranch a
x Int
sh) Tree a
tree))
    | Bool
otherwise = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
sh (Int -> Tree a -> Tree a
consTree Int
sh Tree a
tree)
  where
    consTree :: Int -> Tree a -> Tree a
consTree Int
sh (Balanced Array (Tree a)
arr)
        | Int
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
insertShift = Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.cons Array (Tree a)
arr (a -> Int -> Tree a
forall a. a -> Int -> Tree a
newBranch a
x (Int -> Int
down Int
sh)))
        | Bool
otherwise = Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Int
0 (Int -> Tree a -> Tree a
consTree (Int -> Int
down Int
sh)))
    consTree Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
_)
        | Int
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
insertShift = Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.cons Array (Tree a)
arr (a -> Int -> Tree a
forall a. a -> Int -> Tree a
newBranch a
x (Int -> Int
down Int
sh)))
        | Bool
otherwise = Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Int
0 (Int -> Tree a -> Tree a
consTree (Int -> Int
down Int
sh)))
    consTree Int
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ Array a -> a -> Array a
forall a. Array a -> a -> Array a
A.cons Array a
arr a
x

    insertShift :: Int
insertShift = Int -> Int -> Int -> Tree a -> Int
forall a. Int -> Int -> Int -> Tree a -> Int
computeShift Int
size Int
sh (Int -> Int
up Int
sh) Tree a
tree

    -- compute the shift at which the new branch needs to be inserted (0 means there is space in the leaf)
    -- the index is computed for efficient calculation of the shift in a balanced subtree
    computeShift :: Int -> Int -> Int -> Tree a -> Int
computeShift Int
i Int
sh Int
min (Balanced Array (Tree a)
_) =
        let newShift :: Int
newShift = (Int -> Int
log2 Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
blockShift) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
blockShift
        in if Int
newShift Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
sh then Int
min else Int
newShift
    computeShift Int
_ Int
sh Int
min (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) =
        let i' :: Int
i' = PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes Int
0 -- the size of the first subtree
            newMin :: Int
newMin = if Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
sh else Int
min
        in Int -> Int -> Int -> Tree a -> Int
computeShift Int
i' (Int -> Int
down Int
sh) Int
newMin (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
arr)
    computeShift Int
_ Int
_ Int
min (Leaf Array a
arr) = if Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
0 else Int
min

-- | \(O(\log n)\). Add an element to the right end of the vector.
--
-- >>> fromList [1, 2, 3] |> 4
-- fromList [1,2,3,4]
(|>) :: Vector a -> a -> Vector a
Vector a
Empty |> :: Vector a -> a -> Vector a
|> a
x = a -> Vector a
forall a. a -> Vector a
singleton a
x
Root Int
size Int
sh Tree a
tree |> a
x
    | Int
insertShift Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
sh = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
insertShift (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
insertShift (Tree a -> Tree a -> Array (Tree a)
forall a. a -> a -> Array a
A.from2 Tree a
tree (a -> Int -> Tree a
forall a. a -> Int -> Tree a
newBranch a
x Int
sh)))
    | Bool
otherwise = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
sh (Int -> Tree a -> Tree a
snocTree Int
sh Tree a
tree)
  where
    snocTree :: Int -> Tree a -> Tree a
snocTree Int
sh (Balanced Array (Tree a)
arr)
        | Int
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
insertShift = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.snoc Array (Tree a)
arr (a -> Int -> Tree a
forall a. a -> Int -> Tree a
newBranch a
x (Int -> Int
down Int
sh)) -- the current subtree is fully balanced
        | Bool
otherwise = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> Tree a -> Tree a
snocTree (Int -> Int
down Int
sh))
    snocTree Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
sizes)
        | Int
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
insertShift = Array (Tree a) -> PrimArray Int -> Tree a
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.snoc Array (Tree a)
arr (a -> Int -> Tree a
forall a. a -> Int -> Tree a
newBranch a
x (Int -> Int
down Int
sh))) PrimArray Int
newSizesSnoc
        | Bool
otherwise = Array (Tree a) -> PrimArray Int -> Tree a
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int -> Tree a -> Tree a
snocTree (Int -> Int
down Int
sh))) PrimArray Int
newSizesAdjust
      where
        -- snoc the last size + 1
        newSizesSnoc :: PrimArray Int
newSizesSnoc = (forall s. ST s (PrimArray Int)) -> PrimArray Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (PrimArray Int)) -> PrimArray Int)
-> (forall s. ST s (PrimArray Int)) -> PrimArray Int
forall a b. (a -> b) -> a -> b
$ do
            let lenSizes :: Int
lenSizes = PrimArray Int -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
sizes
            MutablePrimArray s Int
newArr <- Int -> ST s (MutablePrimArray (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (Int
lenSizes Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            MutablePrimArray (PrimState (ST s)) Int
-> Int -> PrimArray Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
newArr Int
0 PrimArray Int
sizes Int
0 Int
lenSizes
            let lastSize :: Int
lastSize = PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (Int
lenSizes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            MutablePrimArray (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
newArr Int
lenSizes (Int
lastSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            MutablePrimArray (PrimState (ST s)) Int -> ST s (PrimArray Int)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
newArr
        -- adjust the last size with (+ 1)
        newSizesAdjust :: PrimArray Int
newSizesAdjust = (forall s. ST s (PrimArray Int)) -> PrimArray Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (PrimArray Int)) -> PrimArray Int)
-> (forall s. ST s (PrimArray Int)) -> PrimArray Int
forall a b. (a -> b) -> a -> b
$ do
            let lenSizes :: Int
lenSizes = PrimArray Int -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
sizes
            MutablePrimArray s Int
newArr <- Int -> ST s (MutablePrimArray (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
lenSizes
            MutablePrimArray (PrimState (ST s)) Int
-> Int -> PrimArray Int -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
newArr Int
0 PrimArray Int
sizes Int
0 Int
lenSizes
            let lastSize :: Int
lastSize = PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (Int
lenSizes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            MutablePrimArray (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
newArr (Int
lenSizes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (Int
lastSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            MutablePrimArray (PrimState (ST s)) Int -> ST s (PrimArray Int)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
newArr
    snocTree Int
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ Array a -> a -> Array a
forall a. Array a -> a -> Array a
A.snoc Array a
arr a
x

    insertShift :: Int
insertShift = Int -> Int -> Int -> Tree a -> Int
forall a. Int -> Int -> Int -> Tree a -> Int
computeShift Int
size Int
sh (Int -> Int
up Int
sh) Tree a
tree

    -- compute the shift at which the new branch needs to be inserted (0 means there is space in the leaf)
    -- the index is computed for efficient calculation of the shift in a balanced subtree
    computeShift :: Int -> Int -> Int -> Tree a -> Int
computeShift Int
i Int
sh Int
min (Balanced Array (Tree a)
_) =
        let newShift :: Int
newShift = (Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
blockShift) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
blockShift
        in if Int
newShift Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
sh then Int
min else Int
newShift
    computeShift Int
_ Int
sh Int
min (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) =
        let i' :: Int
i' = PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (PrimArray Int -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
sizes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (PrimArray Int -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Int
sizes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) -- sizes has at least 2 elements, otherwise the node would be balanced
            newMin :: Int
newMin = if Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
sh else Int
min
        in Int -> Int -> Int -> Tree a -> Int
computeShift Int
i' (Int -> Int
down Int
sh) Int
newMin (Array (Tree a) -> Tree a
forall a. Array a -> a
A.last Array (Tree a)
arr)
    computeShift Int
_ Int
_ Int
min (Leaf Array a
arr) = if Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
0 else Int
min

-- create a new tree with shift @sh@
newBranch :: a -> Int -> Tree a
newBranch :: a -> Int -> Tree a
newBranch a
x = Int -> Tree a
go
  where
    go :: Int -> Tree a
go Int
0 = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ a -> Array a
forall a. a -> Array a
A.singleton a
x
    go Int
sh = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Tree a) -> Array (Tree a) -> Tree a
forall a b. (a -> b) -> a -> b
$ Tree a -> Array (Tree a)
forall a. a -> Array a
A.singleton (Int -> Tree a
go (Int -> Int
down Int
sh))
{-# INLINE newBranch #-}

-- splitting

-- | \(O(\log n)\). The first @i@ elements of the vector.
-- If @i@ is negative, the empty vector is returned. If the vector contains less than @i@ elements, the whole vector is returned.
take :: Int -> Vector a -> Vector a
take :: Int -> Vector a -> Vector a
take Int
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
take Int
n v :: Vector a
v@(Root Int
size Int
sh Tree a
tree)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Vector a
forall a. Vector a
empty
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
size = Vector a
v
    | Bool
otherwise = Vector a -> Vector a
forall a. Vector a -> Vector a
normalize (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root Int
n Int
sh (Int -> Int -> Tree a -> Tree a
forall a. Int -> Int -> Tree a -> Tree a
takeTree (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
sh Tree a
tree)
  where
    -- the initial @i@ is @n - 1@ -- the index of the last element in the new tree
    takeTree :: Int -> Int -> Tree a -> Tree a
takeTree Int
i Int
sh (Balanced Array (Tree a)
arr) =
        let idx :: Int
idx = Int -> Int -> Int
radixIndex Int
i Int
sh
            newArr :: Array (Tree a)
newArr = Array (Tree a) -> Int -> Array (Tree a)
forall a. Array a -> Int -> Array a
A.take Array (Tree a)
arr (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        in Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Int
idx (Int -> Int -> Tree a -> Tree a
takeTree Int
i (Int -> Int
down Int
sh)))
    takeTree Int
i Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) =
        let (Int
idx, Int
subIdx) = PrimArray Int -> Int -> Int -> (Int, Int)
relaxedRadixIndex PrimArray Int
sizes Int
i Int
sh
            newArr :: Array (Tree a)
newArr = Array (Tree a) -> Int -> Array (Tree a)
forall a. Array a -> Int -> Array a
A.take Array (Tree a)
arr (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        in Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Int
idx (Int -> Int -> Tree a -> Tree a
takeTree Int
subIdx (Int -> Int
down Int
sh)))
    takeTree Int
i Int
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Int -> Array a
forall a. Array a -> Int -> Array a
A.take Array a
arr ((Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
blockMask) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))

-- | \(O(\log n)\). The vector without the first @i@ elements
-- If @i@ is negative, the whole vector is returned. If the vector contains less than @i@ elements, the empty vector is returned.
drop :: Int -> Vector a -> Vector a
drop :: Int -> Vector a -> Vector a
drop Int
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
drop Int
n v :: Vector a
v@(Root Int
size Int
sh Tree a
tree)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Vector a
v
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
size = Vector a
forall a. Vector a
empty
    | Bool
otherwise = Vector a -> Vector a
forall a. Vector a -> Vector a
normalize (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) Int
sh (Int -> Int -> Tree a -> Tree a
forall a. Int -> Int -> Tree a -> Tree a
dropTree Int
n Int
sh Tree a
tree)
  where
    dropTree :: Int -> Int -> Tree a -> Tree a
dropTree Int
n Int
sh (Balanced Array (Tree a)
arr) =
        let idx :: Int
idx = Int -> Int -> Int
radixIndex Int
n Int
sh
            newArr :: Array (Tree a)
newArr = Array (Tree a) -> Int -> Array (Tree a)
forall a. Array a -> Int -> Array a
A.drop Array (Tree a)
arr Int
idx
        in Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Int
0 (Int -> Int -> Tree a -> Tree a
dropTree Int
n (Int -> Int
down Int
sh)))
    dropTree Int
n Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) =
        let (Int
idx, Int
subIdx) = PrimArray Int -> Int -> Int -> (Int, Int)
relaxedRadixIndex PrimArray Int
sizes Int
n Int
sh
            newArr :: Array (Tree a)
newArr = Array (Tree a) -> Int -> Array (Tree a)
forall a. Array a -> Int -> Array a
A.drop Array (Tree a)
arr Int
idx
        in Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Int
0 (Int -> Int -> Tree a -> Tree a
dropTree Int
subIdx (Int -> Int
down Int
sh)))
    dropTree Int
n Int
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Int -> Array a
forall a. Array a -> Int -> Array a
A.drop Array a
arr (Int
n Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
blockMask))

normalize :: Vector a -> Vector a
normalize :: Vector a -> Vector a
normalize (Root Int
size Int
sh (Balanced Array (Tree a)
arr))
    | Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Vector a -> Vector a
forall a. Vector a -> Vector a
normalize (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size (Int -> Int
down Int
sh) (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
arr)
normalize (Root Int
size Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
_))
    | Array (Tree a) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = Vector a -> Vector a
forall a. Vector a -> Vector a
normalize (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size (Int -> Int
down Int
sh) (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
arr)
normalize Vector a
v = Vector a
v