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

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

import Control.Applicative (Alternative, liftA2)
import qualified Control.Applicative
import Control.DeepSeq
import Control.Monad (when, MonadPlus)
import Control.Monad.ST (runST)
import qualified Control.Monad.Fail as Fail
import Control.Monad.Fix (MonadFix(..))
import Control.Monad.Zip (MonadZip(..))

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

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

import Data.Primitive.PrimArray hiding (sizeofPrimArray) -- use @length@ of the @A.Array@ instead
import qualified Data.RRBVector.Internal.Array as A
import qualified Data.RRBVector.Internal.Buffer as Buffer

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

type Shift = Int

-- Invariants:
-- Children of a Balanced node are always balanced (a Leaf node is considered balanced).
-- Nodes are always non-empty.
-- The two arrays in an Unbalanced node always have the same size.
data Tree a
    = Balanced {-# UNPACK #-} !(A.Array (Tree a))
    | Unbalanced {-# UNPACK #-} !(A.Array (Tree a)) !(PrimArray Int)
    | Leaf {-# UNPACK #-} !(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
        !Shift -- shift (blockShift * height)
        !(Tree a)

-- The number of bits used per level.
blockShift :: Shift
blockShift :: Int
blockShift = Int
4

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

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

up :: Shift -> Shift
up :: Int -> Int
up Int
sh = Int
sh forall a. Num a => a -> a -> a
+ Int
blockShift

down :: Shift -> Shift
down :: Int -> Int
down Int
sh = Int
sh forall a. Num a => a -> a -> a
- Int
blockShift

radixIndex :: Int -> Shift -> Int
radixIndex :: Int -> Int -> Int
radixIndex Int
i Int
sh = Int
i forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
sh forall a. Bits a => a -> a -> a
.&. Int
blockMask

relaxedRadixIndex :: PrimArray Int -> Int -> Shift -> (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 forall a. Eq a => a -> a -> Bool
== Int
0 then Int
i else Int
i forall a. Num a => a -> a -> a
- forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (Int
idx 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 = 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 forall a. Ord a => a -> a -> Bool
< Int
current then Int
idx else Int -> Int
loop (Int
idx forall a. Num a => a -> a -> a
+ Int
1)

treeToArray :: Tree a -> A.Array (Tree a)
treeToArray :: forall a. 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
_) = forall a. HasCallStack => [Char] -> a
error [Char]
"treeToArray: leaf"

treeBalanced :: Tree a -> Bool
treeBalanced :: forall a. 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 :: Shift -> Tree a -> Int
treeSize :: forall a. Int -> Tree a -> Int
treeSize = 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 forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr
    go Int
acc Int
_ (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = Int
acc forall a. Num a => a -> a -> a
+ forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr forall a. Num a => a -> a -> a
- Int
1)
    go Int
acc Int
sh (Balanced Array (Tree a)
arr) =
        let i :: Int
i = forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr forall a. Num a => a -> a -> a
- Int
1
        in Int -> Int -> Tree a -> Int
go (Int
acc forall a. Num a => a -> a -> a
+ Int
i forall a. Num a => a -> a -> a
* (Int
1 forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
sh)) (Int -> Int
down Int
sh) (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 :: Shift -> A.Array (Tree a) -> Tree a
computeSizes :: forall a. Int -> Array (Tree a) -> Tree a
computeSizes !Int
sh Array (Tree a)
arr
    | Bool
isBalanced = forall a. Array (Tree a) -> Tree a
Balanced Array (Tree a)
arr
    | Bool
otherwise = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
        MutablePrimArray s Int
sizes <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr)
        let loop :: Int -> Int -> ST s (Tree a)
loop Int
acc Int
i
                | Int
i forall a. Ord a => a -> a -> Bool
< Int
len =
                    let size :: Int
size = forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh) (forall a. Array a -> Int -> a
A.index Array (Tree a)
arr Int
i)
                        acc' :: Int
acc' = Int
acc forall a. Num a => a -> a -> a
+ Int
size
                    in forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s Int
sizes Int
i Int
acc' forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> Int -> ST s (Tree a)
loop Int
acc' (Int
i forall a. Num a => a -> a -> a
+ Int
1)
                | Bool
otherwise = do
                    PrimArray Int
sizes <- forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s Int
sizes -- safe because the mutable @sizes@ isn't used afterwards
                    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced Array (Tree a)
arr PrimArray Int
sizes
        Int -> Int -> ST s (Tree a)
loop Int
0 Int
0
  where
    maxSize :: Int
maxSize = Int
1 forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
sh -- the maximum size of a subtree

    len :: Int
len = forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr

    lenM1 :: Int
lenM1 = Int
len forall a. Num a => a -> a -> a
- Int
1

    isBalanced :: Bool
isBalanced = Int -> Bool
go Int
0
      where
        go :: Int -> Bool
go Int
i
            | Int
i forall a. Ord a => a -> a -> Bool
< Int
lenM1 = forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh) Tree a
subtree forall a. Eq a => a -> a -> Bool
== Int
maxSize Bool -> Bool -> Bool
&& Int -> Bool
go (Int
i forall a. Num a => a -> a -> a
+ Int
1)
            | Bool
otherwise = forall a. Tree a -> Bool
treeBalanced Tree a
subtree
          where
            subtree :: Tree a
subtree = forall a. Array a -> Int -> a
A.index Array (Tree a)
arr Int
i

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

instance Show1 Vector where
    liftShowsPrec :: forall a.
(Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> Vector a -> ShowS
liftShowsPrec Int -> a -> ShowS
sp [a] -> ShowS
sl Int
p Vector a
v = forall a. (Int -> a -> ShowS) -> [Char] -> Int -> a -> ShowS
showsUnaryWith (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 (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 = forall (f :: * -> *) a. (Show1 f, Show a) => Int -> f a -> ShowS
showsPrec1

instance Read1 Vector where
    liftReadPrec :: forall a. ReadPrec a -> ReadPrec [a] -> ReadPrec (Vector a)
liftReadPrec ReadPrec a
rp ReadPrec [a]
rl = forall a. ReadPrec a -> ReadPrec a
readData forall a b. (a -> b) -> a -> b
$ forall a t. ReadPrec a -> [Char] -> (a -> t) -> ReadPrec t
readUnaryWith (forall (f :: * -> *) a.
Read1 f =>
ReadPrec a -> ReadPrec [a] -> ReadPrec (f a)
liftReadPrec ReadPrec a
rp ReadPrec [a]
rl) [Char]
"fromList" forall a. [a] -> Vector a
fromList
    liftReadListPrec :: forall a. ReadPrec a -> ReadPrec [a] -> ReadPrec [Vector a]
liftReadListPrec = 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 = forall (f :: * -> *) a. (Read1 f, Read a) => ReadPrec (f a)
readPrec1
    readListPrec :: ReadPrec [Vector a]
readListPrec = forall a. Read a => ReadPrec [a]
readListPrecDefault

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

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

instance Ord1 Vector where
    liftCompare :: forall a b.
(a -> b -> Ordering) -> Vector a -> Vector b -> Ordering
liftCompare a -> b -> Ordering
f Vector a
v1 Vector b
v2 = forall (f :: * -> *) a b.
Ord1 f =>
(a -> b -> Ordering) -> f a -> f b -> Ordering
liftCompare a -> b -> Ordering
f (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (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 = forall (f :: * -> *) a. (Ord1 f, Ord a) => f a -> f a -> Ordering
compare1

instance Semigroup (Vector a) where
    <> :: Vector a -> Vector a -> Vector a
(<>) = forall a. Vector a -> Vector a -> Vector a
(><)
    sconcat :: NonEmpty (Vector a) -> Vector a
sconcat (Vector a
v :| [Vector a]
vs) = Vector a
v forall a. Vector a -> Vector a -> Vector a
>< forall a. Monoid a => [a] -> a
mconcat [Vector a]
vs
    stimes :: forall b. Integral b => b -> Vector a -> Vector a
stimes = forall b a. (Integral b, Monoid a) => b -> a -> a
stimesMonoid

instance Monoid (Vector a) where
    mempty :: Vector a
mempty = forall a. Vector a
empty
    mconcat :: [Vector a] -> Vector a
mconcat = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Vector a -> Vector a -> Vector a
(><) forall a. Vector a
empty

instance Foldable Vector where
    foldr :: forall a b. (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' = 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' = 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' = 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 :: forall b a. (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) = 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
_) = 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) = 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' :: forall a b. (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' = 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' = 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' = 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' :: forall b a. (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) = 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
_) = 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) = 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 :: forall a. Vector a -> Bool
null Vector a
Empty = Bool
True
    null Root{} = Bool
False

    length :: forall a. Vector a -> Int
length Vector a
Empty = Int
0
    length (Root Int
s Int
_ Tree a
_) = Int
s

instance FoldableWithIndex Int Vector where
    ifoldMap :: forall m a. Monoid m => (Int -> a -> m) -> Vector a -> m
ifoldMap Int -> a -> m
f = forall i (f :: * -> *) a b.
FoldableWithIndex i f =>
(i -> a -> b -> b) -> b -> f a -> b
ifoldr (\Int
i a
x m
acc -> Int -> a -> m
f Int
i a
x forall a. Semigroup a => a -> a -> a
<> m
acc) forall a. Monoid a => a
mempty
    {-# INLINE ifoldMap #-}

    ifoldr :: forall a b. (Int -> a -> b -> b) -> b -> Vector a -> b
ifoldr Int -> 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
sh Tree a
tree) = Int -> Int -> Tree a -> b -> b
ifoldrTree Int
0 Int
sh Tree a
tree b
acc

        ifoldrTree :: Int -> Int -> Tree a -> b -> b
ifoldrTree !Int
i !Int
sh (Balanced Array (Tree a)
arr) b
acc' = forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
A.ifoldrStep Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> Tree a -> b -> b
ifoldrTree (Int -> Int
down Int
sh)) b
acc' Array (Tree a)
arr
        ifoldrTree Int
i Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
_) b
acc' = forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
A.ifoldrStep Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> Tree a -> b -> b
ifoldrTree (Int -> Int
down Int
sh)) b
acc' Array (Tree a)
arr
        ifoldrTree Int
i Int
_ (Leaf Array a
arr) b
acc' = forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
A.ifoldrStep Int
i (\a
_ -> Int
1) Int -> a -> b -> b
f b
acc' Array a
arr
    {-# INLINE ifoldr #-}

    ifoldl :: forall b a. (Int -> b -> a -> b) -> b -> Vector a -> b
ifoldl Int -> b -> a -> b
f b
acc = Vector a -> b
go
      where
        go :: Vector a -> b
go Vector a
Empty = b
acc
        go (Root Int
size Int
sh Tree a
tree) = Int -> Int -> b -> Tree a -> b
ifoldlTree (Int
size forall a. Num a => a -> a -> a
- Int
1) Int
sh b
acc Tree a
tree

        ifoldlTree :: Int -> Int -> b -> Tree a -> b
ifoldlTree !Int
i !Int
sh b
acc' (Balanced Array (Tree a)
arr) = forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
A.ifoldlStep Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> b -> Tree a -> b
ifoldlTree (Int -> Int
down Int
sh)) b
acc' Array (Tree a)
arr
        ifoldlTree Int
i Int
sh b
acc' (Unbalanced Array (Tree a)
arr PrimArray Int
_) = forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
A.ifoldlStep Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> b -> Tree a -> b
ifoldlTree (Int -> Int
down Int
sh)) b
acc' Array (Tree a)
arr
        ifoldlTree Int
i Int
_ b
acc' (Leaf Array a
arr) = forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
A.ifoldlStep Int
i (\a
_ -> Int
1) Int -> b -> a -> b
f b
acc' Array a
arr
    {-# INLINE ifoldl #-}

    ifoldr' :: forall a b. (Int -> a -> b -> b) -> b -> Vector a -> b
ifoldr' Int -> a -> b -> b
f b
acc = Vector a -> b
go
      where
        go :: Vector a -> b
go Vector a
Empty = b
acc
        go (Root Int
size Int
sh Tree a
tree) = Int -> Int -> Tree a -> b -> b
ifoldrTree' (Int
size forall a. Num a => a -> a -> a
- Int
1) Int
sh Tree a
tree b
acc

        ifoldrTree' :: Int -> Int -> Tree a -> b -> b
ifoldrTree' !Int
i !Int
sh (Balanced Array (Tree a)
arr) b
acc' = forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
A.ifoldrStep' Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> Tree a -> b -> b
ifoldrTree' (Int -> Int
down Int
sh)) b
acc' Array (Tree a)
arr
        ifoldrTree' Int
i Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
_) b
acc' = forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
A.ifoldrStep' Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> Tree a -> b -> b
ifoldrTree' (Int -> Int
down Int
sh)) b
acc' Array (Tree a)
arr
        ifoldrTree' Int
i Int
_ (Leaf Array a
arr) b
acc' = forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
A.ifoldrStep' Int
i (\a
_ -> Int
1) Int -> a -> b -> b
f b
acc' Array a
arr
    {-# INLINE ifoldr' #-}

    ifoldl' :: forall b a. (Int -> b -> a -> b) -> b -> Vector a -> b
ifoldl' Int -> 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
sh Tree a
tree) = Int -> Int -> b -> Tree a -> b
ifoldlTree' Int
0 Int
sh b
acc Tree a
tree

        ifoldlTree' :: Int -> Int -> b -> Tree a -> b
ifoldlTree' !Int
i !Int
sh b
acc' (Balanced Array (Tree a)
arr) = forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
A.ifoldlStep' Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> b -> Tree a -> b
ifoldlTree' (Int -> Int
down Int
sh)) b
acc' Array (Tree a)
arr
        ifoldlTree' Int
i Int
sh b
acc' (Unbalanced Array (Tree a)
arr PrimArray Int
_) = forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
A.ifoldlStep' Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> Int -> b -> Tree a -> b
ifoldlTree' (Int -> Int
down Int
sh)) b
acc' Array (Tree a)
arr
        ifoldlTree' Int
i Int
_ b
acc' (Leaf Array a
arr) = forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
A.ifoldlStep' Int
i (\a
_ -> Int
1) Int -> b -> a -> b
f b
acc' Array a
arr
    {-# INLINE ifoldl' #-}

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

instance FunctorWithIndex Int Vector where
    imap :: forall a b. (Int -> a -> b) -> Vector a -> Vector b
imap Int -> a -> b
_ Vector a
Empty = forall a. Vector a
Empty
    imap Int -> a -> b
f (Root Int
size Int
sh Tree a
tree) = forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (Int -> Int -> Tree a -> Tree b
imapTree Int
0 Int
sh Tree a
tree)
      where
        imapTree :: Int -> Int -> Tree a -> Tree b
imapTree !Int
i !Int
sh (Balanced Array (Tree a)
arr) = forall a. Array (Tree a) -> Tree a
Balanced (forall a b.
Int -> (a -> Int) -> (Int -> a -> b) -> Array a -> Array b
A.imapStep' Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (\Int
i -> Int -> Int -> Tree a -> Tree b
imapTree Int
i (Int -> Int
down Int
sh)) Array (Tree a)
arr)
        imapTree Int
i Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (forall a b.
Int -> (a -> Int) -> (Int -> a -> b) -> Array a -> Array b
A.imapStep' Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (\Int
i -> Int -> Int -> Tree a -> Tree b
imapTree Int
i (Int -> Int
down Int
sh)) Array (Tree a)
arr) PrimArray Int
sizes
        imapTree Int
i Int
_ (Leaf Array a
arr) = forall a. Array a -> Tree a
Leaf (forall a b.
Int -> (a -> Int) -> (Int -> a -> b) -> Array a -> Array b
A.imapStep Int
i (\a
_ -> Int
1) Int -> a -> b
f Array a
arr)

instance Traversable Vector where
    traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Vector a -> f (Vector b)
traverse a -> f b
_ Vector a
Empty = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Vector a
Empty
    traverse a -> f b
f (Root Int
size Int
sh Tree a
tree) = forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh 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) = forall a. Array (Tree a) -> Tree a
Balanced forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f 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) = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced PrimArray Int
sizes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f 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 (Leaf Array a
arr) = forall a. Array a -> Tree a
Leaf forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array a -> f (Array b)
A.traverse a -> f b
f Array a
arr
    {-# INLINE traverse #-}

instance TraversableWithIndex Int Vector where
    itraverse :: forall (f :: * -> *) a b.
Applicative f =>
(Int -> a -> f b) -> Vector a -> f (Vector b)
itraverse Int -> a -> f b
_ Vector a
Empty = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Vector a
Empty
    itraverse Int -> a -> f b
f (Root Int
size Int
sh Tree a
tree) = forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> Tree a -> f (Tree b)
itraverseTree Int
0 Int
sh Tree a
tree
      where
        itraverseTree :: Int -> Int -> Tree a -> f (Tree b)
itraverseTree !Int
i !Int
sh (Balanced Array (Tree a)
arr) = forall a. Array (Tree a) -> Tree a
Balanced forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a b.
Applicative f =>
Int -> (a -> Int) -> (Int -> a -> f b) -> Array a -> f (Array b)
A.itraverseStep' Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (\Int
i -> Int -> Int -> Tree a -> f (Tree b)
itraverseTree Int
i (Int -> Int
down Int
sh)) Array (Tree a)
arr
        itraverseTree Int
i Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced PrimArray Int
sizes forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a b.
Applicative f =>
Int -> (a -> Int) -> (Int -> a -> f b) -> Array a -> f (Array b)
A.itraverseStep' Int
i (forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) (\Int
i -> Int -> Int -> Tree a -> f (Tree b)
itraverseTree Int
i (Int -> Int
down Int
sh)) Array (Tree a)
arr
        itraverseTree Int
i Int
_ (Leaf Array a
arr) = forall a. Array a -> Tree a
Leaf forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a b.
Applicative f =>
Int -> (a -> Int) -> (Int -> a -> f b) -> Array a -> f (Array b)
A.itraverseStep Int
i (\a
_ -> Int
1) Int -> a -> f b
f Array a
arr
    {-# INLINE itraverse #-}

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

instance Monad Vector where
    Vector a
xs >>= :: forall a b. Vector a -> (a -> Vector b) -> Vector b
>>= a -> Vector b
f = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Vector b
acc a
x -> Vector b
acc forall a. Vector a -> Vector a -> Vector a
>< a -> Vector b
f a
x) forall a. Vector a
empty Vector a
xs
    >> :: forall a b. Vector a -> Vector b -> Vector b
(>>) = forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)
#if !(MIN_VERSION_base(4,13,0))
    fail = Fail.fail
#endif

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

instance MonadPlus Vector

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

instance MonadFix Vector where
    mfix :: forall a. (a -> Vector a) -> Vector a
mfix a -> Vector a
f = case a -> Vector a
f forall {a}. a
err of
        Vector a
Empty -> forall a. Vector a
Empty
        Root Int
size Int
_ Tree a
_ -> forall a. [a] -> Vector a
fromList forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
i -> let x :: a
x = a -> Vector a
f a
x forall a. HasCallStack => Vector a -> Int -> a
! Int
i in a
x) [Int
0..Int
size forall a. Num a => a -> a -> a
- Int
1]
      where
        err :: a
err = forall a. HasCallStack => [Char] -> a
error [Char]
"mfix for Data.RRBVector.Vector applied to strict function"

instance MonadZip Vector where
    mzipWith :: forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
mzipWith = forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith
    mzip :: forall a b. Vector a -> Vector b -> Vector (a, b)
mzip = forall a b. Vector a -> Vector b -> Vector (a, b)
zip
    munzip :: forall a b. Vector (a, b) -> (Vector a, Vector b)
munzip = 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 = forall a. [a] -> Vector a
fromList
    toList :: Vector a -> [Item (Vector a)]
toList = forall (t :: * -> *) a. Foldable t => t a -> [a]
toList

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

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

instance NFData1 Vector where
    liftRnf :: forall a. (a -> ()) -> Vector a -> ()
liftRnf a -> ()
f = 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 :: forall a. Vector a
empty = forall a. Vector a
Empty

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

-- | \(O(n)\). Create a new vector from a list.
fromList :: [a] -> Vector a
fromList :: forall a. [a] -> Vector a
fromList [] = forall a. Vector a
Empty
fromList [a
x] = forall a. a -> Vector a
singleton a
x
fromList [a]
ls = case forall {a} {a}. (Array a -> a) -> [a] -> [a]
nodes forall a. Array a -> Tree a
Leaf [a]
ls of
    [Tree a
tree] -> forall a. Int -> Int -> Tree a -> Vector a
Root (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' -> 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 a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
        Buffer s a
buffer <- forall s a. Int -> ST s (Buffer s a)
Buffer.new Int
blockSize
        let loop :: [a] -> ST s [a]
loop [] = do
                Array a
result <- forall s a. Buffer s a -> ST s (Array a)
Buffer.get Buffer s a
buffer
                forall (f :: * -> *) a. Applicative f => a -> f a
pure [Array a -> a
f Array a
result]
            loop (a
t : [a]
ts) = do
                Int
size <- forall s a. Buffer s a -> ST s Int
Buffer.size Buffer s a
buffer
                if Int
size forall a. Eq a => a -> a -> Bool
== Int
blockSize then do
                    Array a
result <- forall s a. Buffer s a -> ST s (Array a)
Buffer.get Buffer s a
buffer
                    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
                    forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array a -> a
f Array a
result forall a. a -> [a] -> [a]
: [a]
rest)
                else do
                    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 forall {a} {a}. (Array a -> a) -> [a] -> [a]
nodes forall a. Array (Tree a) -> Tree a
Balanced [Tree a]
trees of
        [Tree a
tree] -> forall a. Int -> Int -> Tree a -> Vector a
Root (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)\). @replicate n x@ creates a vector of length @n@ with every element set to @x@.
--
-- >>> replicate 5 42
-- fromList [42,42,42,42,42]
--
-- @since 0.1.1.0
replicate :: Int -> a -> Vector a
replicate :: forall a. Int -> a -> Vector a
replicate Int
n a
x
    | Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = forall a. Vector a
Empty
    | Int
n forall a. Ord a => a -> a -> Bool
<= Int
blockSize = forall a. Int -> Int -> Tree a -> Vector a
Root Int
n Int
0 (forall a. Array a -> Tree a
Leaf forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> Array a
A.replicate Int
n a
x)
    | Bool
otherwise = Int -> Tree a -> Tree a -> Vector a
iterateNodes Int
blockShift (forall a. Array a -> Tree a
Leaf forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> Array a
A.replicate Int
blockSize a
x) (forall a. Array a -> Tree a
Leaf forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> Array a
A.replicate (Int
lastIdx forall a. Bits a => a -> a -> a
.&. Int
blockMask forall a. Num a => a -> a -> a
+ Int
1) a
x)
  where
    lastIdx :: Int
lastIdx = Int
n forall a. Num a => a -> a -> a
- Int
1

    -- @full@ is a full subtree, @rest@ is the last subtree
    iterateNodes :: Int -> Tree a -> Tree a -> Vector a
iterateNodes !Int
sh !Tree a
full !Tree a
rest =
        let subtreesM1 :: Int
subtreesM1 = Int
lastIdx forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
sh -- the number of subtrees minus 1
            full' :: Tree a
full' = forall a. Array (Tree a) -> Tree a
Balanced forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> Array a
A.replicate Int
blockSize Tree a
full
            rest' :: Tree a
rest' = forall a. Array (Tree a) -> Tree a
Balanced forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> a -> Array a
A.replicateSnoc (Int
subtreesM1 forall a. Bits a => a -> a -> a
.&. Int
blockMask) Tree a
full Tree a
rest
        in if Int
subtreesM1 forall a. Ord a => a -> a -> Bool
< Int
blockSize then forall a. Int -> Int -> Tree a -> Vector a
Root Int
n Int
sh Tree a
rest' else Int -> Tree a -> Tree a -> Vector a
iterateNodes (Int -> Int
up Int
sh) Tree a
full' Tree a
rest'

-- | \(O(\log n)\). The element at the index or 'Nothing' if the index is out of range.
lookup :: Int -> Vector a -> Maybe a
lookup :: forall a. Int -> Vector a -> Maybe a
lookup !Int
_ Vector a
Empty = forall a. Maybe a
Nothing
lookup Int
i (Root Int
size Int
sh Tree a
tree)
    | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i forall a. Ord a => a -> a -> Bool
>= Int
size = forall a. Maybe a
Nothing  -- index out of range
    | Bool
otherwise = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall {t}. Int -> Int -> Tree t -> t
lookupTree Int
i Int
sh Tree a
tree
  where
    lookupTree :: Int -> Int -> Tree t -> t
lookupTree Int
i Int
sh (Balanced Array (Tree t)
arr) = Int -> Int -> Tree t -> t
lookupTree Int
i (Int -> Int
down Int
sh) (forall a. Array a -> Int -> a
A.index Array (Tree t)
arr (Int -> Int -> Int
radixIndex Int
i Int
sh))
    lookupTree Int
i Int
sh (Unbalanced Array (Tree t)
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 t -> t
lookupTree Int
subIdx (Int -> Int
down Int
sh) (forall a. Array a -> Int -> a
A.index Array (Tree t)
arr Int
idx)
    lookupTree Int
i Int
_ (Leaf Array t
arr) = forall a. Array a -> Int -> a
A.index Array t
arr (Int
i forall a. Bits a => a -> a -> a
.&. Int
blockMask)
{-# INLINE lookup #-}

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

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

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

-- | \(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 :: forall a. Int -> a -> Vector a -> Vector a
update !Int
_ a
_ Vector a
Empty = forall a. Vector a
Empty
update Int
i a
x v :: Vector a
v@(Root Int
size Int
sh Tree a
tree)
    | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i forall a. Ord a => a -> a -> Bool
>= Int
size = Vector a
v  -- index out of range
    | Bool
otherwise = forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (Int -> Int -> Tree a -> Tree a
updateTree Int
i Int
sh Tree a
tree)
  where
    updateTree :: Int -> Int -> Tree a -> Tree a
updateTree Int
i Int
sh (Balanced Array (Tree a)
arr) = forall a. Array (Tree a) -> Tree a
Balanced (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
updateTree Int
i (Int -> Int
down Int
sh)))
    updateTree 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 forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Int
idx (Int -> Int -> Tree a -> Tree a
updateTree Int
subIdx (Int -> Int
down Int
sh))) PrimArray Int
sizes
    updateTree Int
i Int
_ (Leaf Array a
arr) = forall a. Array a -> Tree a
Leaf (forall a. Array a -> Int -> a -> Array a
A.update Array a
arr (Int
i 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 :: forall a. Int -> (a -> a) -> Vector a -> Vector a
adjust !Int
_ a -> a
_ Vector a
Empty = 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 forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i forall a. Ord a => a -> a -> Bool
>= Int
size = Vector a
v  -- index out of range
    | Bool
otherwise = 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) = forall a. Array (Tree a) -> Tree a
Balanced (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 forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (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) = forall a. Array a -> Tree a
Leaf (forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust Array a
arr (Int
i 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' :: forall a. Int -> (a -> a) -> Vector a -> Vector a
adjust' !Int
_ a -> a
_ Vector a
Empty = 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 forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i forall a. Ord a => a -> a -> Bool
>= Int
size = Vector a
v  -- index out of range
    | Bool
otherwise = 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) = forall a. Array (Tree a) -> Tree a
Balanced (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 forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (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) = forall a. Array a -> Tree a
Leaf (forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array a
arr (Int
i 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 :: forall a b. (a -> b) -> Vector a -> Vector b
map a -> b
_ Vector a
Empty = forall a. Vector a
Empty
map a -> b
f (Root Int
size Int
sh Tree a
tree) = 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) = forall a. Array (Tree a) -> Tree a
Balanced (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) = forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (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) = forall a. Array a -> Tree a
Leaf (forall a b. (a -> b) -> Array a -> Array b
A.map a -> b
f Array a
arr)

-- | \(O(n)\). Like 'map', but the results of the function are forced.
--
-- @since 0.2.0.0
map' :: (a -> b) -> Vector a -> Vector b
map' :: forall a b. (a -> b) -> Vector a -> Vector b
map' a -> b
_ Vector a
Empty = forall a. Vector a
Empty
map' a -> b
f (Root Int
size Int
sh Tree a
tree) = 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) = forall a. Array (Tree a) -> Tree a
Balanced (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) = forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (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) = forall a. Array a -> Tree a
Leaf (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 :: forall a. Vector a -> Vector a
reverse Vector a
v
    | forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v forall a. Ord a => a -> a -> Bool
<= Int
1 = Vector a
v
    | Bool
otherwise = forall a. [a] -> Vector a
fromList (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) [] Vector a
v) -- 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 :: Vector a -> Vector b -> Vector (a, b)
zip :: forall a b. Vector a -> Vector b -> Vector (a, b)
zip Vector a
v1 Vector b
v2 = forall a. [a] -> Vector a
fromList forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
List.zip (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (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 :: forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith a -> b -> c
f Vector a
v1 Vector b
v2 = forall a. [a] -> Vector a
fromList forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
List.zipWith a -> b -> c
f (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector b
v2)

-- TODO: unzip = unzipWith id
-- | \(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 :: forall a b. Vector (a, b) -> (Vector a, Vector b)
unzip Vector (a, b)
Empty = (forall a. Vector a
Empty, forall a. Vector a
Empty)
unzip (Root Int
size Int
sh Tree (a, b)
tree) = case forall {a} {a}. Tree (a, a) -> (Tree a, Tree a)
unzipTree Tree (a, b)
tree of
    (!Tree a
left, !Tree b
right) -> (forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh Tree a
left, forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh Tree b
right)
  where
    unzipTree :: Tree (a, a) -> (Tree a, Tree a)
unzipTree (Balanced Array (Tree (a, a))
arr) = case forall a b c. (a -> (b, c)) -> Array a -> (Array b, Array c)
A.unzipWith Tree (a, a) -> (Tree a, Tree a)
unzipTree Array (Tree (a, a))
arr of
        (!Array (Tree a)
left, !Array (Tree a)
right) -> (forall a. Array (Tree a) -> Tree a
Balanced Array (Tree a)
left, forall a. Array (Tree a) -> Tree a
Balanced Array (Tree a)
right)
    unzipTree (Unbalanced Array (Tree (a, a))
arr PrimArray Int
sizes) = case forall a b c. (a -> (b, c)) -> Array a -> (Array b, Array c)
A.unzipWith Tree (a, a) -> (Tree a, Tree a)
unzipTree Array (Tree (a, a))
arr of
        (!Array (Tree a)
left, !Array (Tree a)
right) -> (forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced Array (Tree a)
left PrimArray Int
sizes, forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced Array (Tree a)
right PrimArray Int
sizes)
    unzipTree (Leaf Array (a, a)
arr) = case forall a b c. (a -> (b, c)) -> Array a -> (Array b, Array c)
A.unzipWith forall a. a -> a
id Array (a, a)
arr of
        (!Array a
left, !Array a
right) -> (forall a. Array a -> Tree a
Leaf Array a
left, forall a. Array a -> Tree a
Leaf Array a
right)

-- | \(O(n)\). Unzip a vector with a function.
--
-- > unzipWith f = unzip . map f
--
-- @since 0.2.0.0
unzipWith :: (a -> (b, c)) -> Vector a -> (Vector b, Vector c)
unzipWith :: forall a b c. (a -> (b, c)) -> Vector a -> (Vector b, Vector c)
unzipWith a -> (b, c)
_ Vector a
Empty = (forall a. Vector a
Empty, forall a. Vector a
Empty)
unzipWith a -> (b, c)
f (Root Int
size Int
sh Tree a
tree) = case Tree a -> (Tree b, Tree c)
unzipTree Tree a
tree of
    (!Tree b
left, !Tree c
right) -> (forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh Tree b
left, forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh Tree c
right)
  where
    unzipTree :: Tree a -> (Tree b, Tree c)
unzipTree (Balanced Array (Tree a)
arr) = case forall a b c. (a -> (b, c)) -> Array a -> (Array b, Array c)
A.unzipWith Tree a -> (Tree b, Tree c)
unzipTree Array (Tree a)
arr of
        (!Array (Tree b)
left, !Array (Tree c)
right) -> (forall a. Array (Tree a) -> Tree a
Balanced Array (Tree b)
left, forall a. Array (Tree a) -> Tree a
Balanced Array (Tree c)
right)
    unzipTree (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = case forall a b c. (a -> (b, c)) -> Array a -> (Array b, Array c)
A.unzipWith Tree a -> (Tree b, Tree c)
unzipTree Array (Tree a)
arr of
        (!Array (Tree b)
left, !Array (Tree c)
right) -> (forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced Array (Tree b)
left PrimArray Int
sizes, forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced Array (Tree c)
right PrimArray Int
sizes)
    unzipTree (Leaf Array a
arr) = case forall a b c. (a -> (b, c)) -> Array a -> (Array b, Array c)
A.unzipWith a -> (b, c)
f Array a
arr of
        (!Array b
left, !Array c
right) -> (forall a. Array a -> Tree a
Leaf Array b
left, forall a. Array a -> Tree a
Leaf Array c
right)

-- | \(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 :: forall a. Vector a -> Maybe (a, Vector a)
viewl Vector a
Empty = forall a. Maybe a
Nothing
viewl v :: Vector a
v@(Root Int
_ Int
_ Tree a
tree) = let !tail :: Vector a
tail = forall a. Int -> Vector a -> Vector a
drop Int
1 Vector a
v in forall a. a -> Maybe a
Just (forall {t}. Tree t -> t
headTree Tree a
tree, Vector a
tail)
  where
    headTree :: Tree t -> t
headTree (Balanced Array (Tree t)
arr) = Tree t -> t
headTree (forall a. Array a -> a
A.head Array (Tree t)
arr)
    headTree (Unbalanced Array (Tree t)
arr PrimArray Int
_) = Tree t -> t
headTree (forall a. Array a -> a
A.head Array (Tree t)
arr)
    headTree (Leaf Array t
arr) = forall a. Array a -> a
A.head Array t
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 :: forall a. Vector a -> Maybe (Vector a, a)
viewr Vector a
Empty = forall a. Maybe a
Nothing
viewr v :: Vector a
v@(Root Int
size Int
_ Tree a
tree) = let !init :: Vector a
init = forall a. Int -> Vector a -> Vector a
take (Int
size forall a. Num a => a -> a -> a
- Int
1) Vector a
v in forall a. a -> Maybe a
Just (Vector a
init, forall {t}. Tree t -> t
lastTree Tree a
tree)
  where
    lastTree :: Tree t -> t
lastTree (Balanced Array (Tree t)
arr) = Tree t -> t
lastTree (forall a. Array a -> a
A.last Array (Tree t)
arr)
    lastTree (Unbalanced Array (Tree t)
arr PrimArray Int
_) = Tree t -> t
lastTree (forall a. Array a -> a
A.last Array (Tree t)
arr)
    lastTree (Leaf Array t
arr) = forall a. Array a -> a
A.last Array t
arr

-- | \(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 :: forall a. Int -> Vector a -> Vector a
take !Int
_ Vector a
Empty = forall a. Vector a
Empty
take Int
n v :: Vector a
v@(Root Int
size Int
sh Tree a
tree)
    | Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = forall a. Vector a
empty
    | Int
n forall a. Ord a => a -> a -> Bool
>= Int
size = Vector a
v
    | Bool
otherwise = forall a. Vector a -> Vector a
normalize forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> Tree a -> Vector a
Root Int
n Int
sh (forall a. Int -> Int -> Tree a -> Tree a
takeTree (Int
n forall a. Num a => a -> a -> a
- Int
1) Int
sh Tree a
tree)

-- | \(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 :: forall a. Int -> Vector a -> Vector a
drop !Int
_ Vector a
Empty = forall a. Vector a
Empty
drop Int
n v :: Vector a
v@(Root Int
size Int
sh Tree a
tree)
    | Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = Vector a
v
    | Int
n forall a. Ord a => a -> a -> Bool
>= Int
size = forall a. Vector a
empty
    | Bool
otherwise = forall a. Vector a -> Vector a
normalize forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size forall a. Num a => a -> a -> a
- Int
n) Int
sh (forall a. Int -> Int -> Tree a -> Tree a
dropTree Int
n Int
sh Tree a
tree)

-- | \(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 :: forall a. Int -> Vector a -> (Vector a, Vector a)
splitAt !Int
_ Vector a
Empty = (forall a. Vector a
Empty, forall a. Vector a
Empty)
splitAt Int
n v :: Vector a
v@(Root Int
size Int
sh Tree a
tree)
    | Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = (forall a. Vector a
empty, Vector a
v)
    | Int
n forall a. Ord a => a -> a -> Bool
>= Int
size = (Vector a
v, forall a. Vector a
empty)
    | Bool
otherwise =
        let !left :: Vector a
left = forall a. Vector a -> Vector a
normalize forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> Tree a -> Vector a
Root Int
n Int
sh (forall a. Int -> Int -> Tree a -> Tree a
takeTree (Int
n forall a. Num a => a -> a -> a
- Int
1) Int
sh Tree a
tree)
            !right :: Vector a
right = forall a. Vector a -> Vector a
normalize forall a b. (a -> b) -> a -> b
$ forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size forall a. Num a => a -> a -> a
- Int
n) Int
sh (forall a. Int -> Int -> Tree a -> Tree a
dropTree Int
n Int
sh Tree a
tree)
        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 :: forall a. Int -> a -> Vector a -> Vector a
insertAt Int
i a
x Vector a
v = let (Vector a
left, Vector a
right) = forall a. Int -> Vector a -> (Vector a, Vector a)
splitAt Int
i Vector a
v in (Vector a
left forall a. Vector a -> a -> Vector a
|> a
x) 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 :: forall a. Int -> Vector a -> Vector a
deleteAt Int
i Vector a
v = let (Vector a
left, Vector a
right) = forall a. Int -> Vector a -> (Vector a, Vector a)
splitAt (Int
i forall a. Num a => a -> a -> a
+ Int
1) Vector a
v in forall a. Int -> Vector a -> Vector a
take Int
i Vector a
left 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 >< :: forall a. 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 = forall a. Ord a => a -> a -> a
max Int
sh1 Int
sh2
        upMaxShift :: Int
upMaxShift = Int -> Int
up Int
maxShift
        newArr :: Array (Tree a)
newArr = forall {a}. Tree a -> Int -> Tree a -> Int -> Array (Tree a)
mergeTrees Tree a
tree1 Int
sh1 Tree a
tree2 Int
sh2
    in if forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
newArr forall a. Eq a => a -> a -> Bool
== Int
1
        then forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size1 forall a. Num a => a -> a -> a
+ Int
size2) Int
maxShift (forall a. Array a -> a
A.head Array (Tree a)
newArr)
        else forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size1 forall a. Num a => a -> a -> a
+ Int
size2) Int
upMaxShift (forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
upMaxShift Array (Tree a)
newArr)
  where
    mergeTrees :: Tree a -> Int -> Tree a -> Int -> Array (Tree a)
mergeTrees tree1 :: Tree a
tree1@(Leaf Array a
arr1) !Int
_ tree2 :: Tree a
tree2@(Leaf Array a
arr2) !Int
_
        | forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr1 forall a. Eq a => a -> a -> Bool
== Int
blockSize = forall a. a -> a -> Array a
A.from2 Tree a
tree1 Tree a
tree2
        | forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr1 forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr2 forall a. Ord a => a -> a -> Bool
<= Int
blockSize = forall a. a -> Array a
A.singleton forall a b. (a -> b) -> a -> b
$! forall a. Array a -> Tree a
Leaf (Array a
arr1 forall a. Array a -> Array a -> Array a
A.++ Array a
arr2)
        | Bool
otherwise =
            let (Array a
left, Array a
right) = forall a. Array a -> Int -> (Array a, Array a)
A.splitAt (Array a
arr1 forall a. Array a -> Array a -> Array a
A.++ Array a
arr2) Int
blockSize -- 'A.splitAt' doesn't copy anything
                !leftTree :: Tree a
leftTree = forall a. Array a -> Tree a
Leaf Array a
left
                !rightTree :: Tree a
rightTree = forall a. Array a -> Tree a
Leaf Array a
right
            in forall a. a -> a -> Array a
A.from2 Tree a
leftTree Tree a
rightTree
    mergeTrees Tree a
tree1 Int
sh1 Tree a
tree2 Int
sh2 = case forall a. Ord a => a -> a -> Ordering
compare Int
sh1 Int
sh2 of
        Ordering
LT ->
            let !right :: Array (Tree a)
right = forall a. Tree a -> Array (Tree a)
treeToArray Tree a
tree2
                (Tree a
rightHead, Array (Tree a)
rightTail) = forall {a}. Array a -> (a, Array a)
viewlArr Array (Tree a)
right
                merged :: Array (Tree a)
merged = Tree a -> Int -> Tree a -> Int -> Array (Tree a)
mergeTrees Tree a
tree1 Int
sh1 Tree a
rightHead (Int -> Int
down Int
sh2)
            in forall a.
Int
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
mergeRebalance Int
sh2 forall a. Array a
A.empty Array (Tree a)
merged Array (Tree a)
rightTail
        Ordering
GT ->
            let !left :: Array (Tree a)
left = forall a. Tree a -> Array (Tree a)
treeToArray Tree a
tree1
                (Array (Tree a)
leftInit, Tree a
leftLast) = forall {b}. Array b -> (Array b, b)
viewrArr Array (Tree a)
left
                merged :: Array (Tree a)
merged = Tree a -> Int -> Tree a -> Int -> Array (Tree a)
mergeTrees Tree a
leftLast (Int -> Int
down Int
sh1) Tree a
tree2 Int
sh2
            in forall a.
Int
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
mergeRebalance Int
sh1 Array (Tree a)
leftInit Array (Tree a)
merged forall a. Array a
A.empty
        Ordering
EQ ->
            let !left :: Array (Tree a)
left = forall a. Tree a -> Array (Tree a)
treeToArray Tree a
tree1
                !right :: Array (Tree a)
right = forall a. Tree a -> Array (Tree a)
treeToArray Tree a
tree2
                (Array (Tree a)
leftInit, Tree a
leftLast) = forall {b}. Array b -> (Array b, b)
viewrArr Array (Tree a)
left
                (Tree a
rightHead, Array (Tree a)
rightTail) = forall {a}. Array a -> (a, Array a)
viewlArr Array (Tree a)
right
                merged :: Array (Tree a)
merged = Tree a -> Int -> Tree a -> Int -> Array (Tree a)
mergeTrees Tree a
leftLast (Int -> Int
down Int
sh1) Tree a
rightHead (Int -> Int
down Int
sh2)
            in forall a.
Int
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
mergeRebalance Int
sh1 Array (Tree a)
leftInit Array (Tree a)
merged Array (Tree a)
rightTail
      where
        viewlArr :: Array a -> (a, Array a)
viewlArr Array a
arr = (forall a. Array a -> a
A.head Array a
arr, forall a. Array a -> Int -> Array a
A.drop Array a
arr Int
1)

        viewrArr :: Array b -> (Array b, b)
viewrArr Array b
arr = (forall a. Array a -> Int -> Array a
A.take Array b
arr (forall (t :: * -> *) a. Foldable t => t a -> Int
length Array b
arr forall a. Num a => a -> a -> a
- Int
1), forall a. Array a -> a
A.last Array b
arr)

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

-- | \(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 <| :: forall a. a -> Vector a -> Vector a
<| Vector a
Empty = forall a. a -> Vector a
singleton a
x
a
x <| Root Int
size Int
sh Tree a
tree
    | Int
insertShift forall a. Ord a => a -> a -> Bool
> Int
sh = forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size forall a. Num a => a -> a -> a
+ Int
1) Int
insertShift (forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
insertShift (let !new :: Tree a
new = forall a. a -> Int -> Tree a
newBranch a
x Int
sh in forall a. a -> a -> Array a
A.from2 Tree a
new Tree a
tree))
    | Bool
otherwise = forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size 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 forall a. Eq a => a -> a -> Bool
== Int
insertShift = forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (forall a. Array a -> a -> Array a
A.cons Array (Tree a)
arr forall a b. (a -> b) -> a -> b
$! forall a. a -> Int -> Tree a
newBranch a
x (Int -> Int
down Int
sh))
        | Bool
otherwise = forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (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 forall a. Eq a => a -> a -> Bool
== Int
insertShift = forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (forall a. Array a -> a -> Array a
A.cons Array (Tree a)
arr forall a b. (a -> b) -> a -> b
$! forall a. a -> Int -> Tree a
newBranch a
x (Int -> Int
down Int
sh))
        | Bool
otherwise = forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (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) = forall a. Array a -> Tree a
Leaf forall a b. (a -> b) -> a -> b
$ forall a. Array a -> a -> Array a
A.cons Array a
arr a
x

    insertShift :: Int
insertShift = 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 size is computed for efficient calculation of the shift in a balanced subtree
    computeShift :: Int -> Int -> Int -> Tree a -> Int
computeShift !Int
sz !Int
sh !Int
min (Balanced Array (Tree a)
_) =
        let newShift :: Int
newShift = (Int -> Int
log2 Int
sz forall a. Integral a => a -> a -> a
`div` Int
blockShift) forall a. Num a => a -> a -> a
* Int
blockShift
        in if Int
newShift 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 sz' :: Int
sz' = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes Int
0 -- the size of the first subtree
            newMin :: Int
newMin = if forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
sh else Int
min
        in Int -> Int -> Int -> Tree a -> Int
computeShift Int
sz' (Int -> Int
down Int
sh) Int
newMin (forall a. Array a -> a
A.head Array (Tree a)
arr)
    computeShift Int
_ Int
_ Int
min (Leaf Array a
arr) = if forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr 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 |> :: forall a. Vector a -> a -> Vector a
|> a
x = forall a. a -> Vector a
singleton a
x
Root Int
size Int
sh Tree a
tree |> a
x
    | Int
insertShift forall a. Ord a => a -> a -> Bool
> Int
sh = forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size forall a. Num a => a -> a -> a
+ Int
1) Int
insertShift (forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
insertShift (forall a. a -> a -> Array a
A.from2 Tree a
tree forall a b. (a -> b) -> a -> b
$! forall a. a -> Int -> Tree a
newBranch a
x Int
sh))
    | Bool
otherwise = forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size 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 forall a. Eq a => a -> a -> Bool
== Int
insertShift = forall a. Array (Tree a) -> Tree a
Balanced (forall a. Array a -> a -> Array a
A.snoc Array (Tree a)
arr forall a b. (a -> b) -> a -> b
$! forall a. a -> Int -> Tree a
newBranch a
x (Int -> Int
down Int
sh)) -- the current subtree is fully balanced
        | Bool
otherwise = forall a. Array (Tree a) -> Tree a
Balanced forall a b. (a -> b) -> a -> b
$ forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr 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 forall a. Eq a => a -> a -> Bool
== Int
insertShift = forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (forall a. Array a -> a -> Array a
A.snoc Array (Tree a)
arr forall a b. (a -> b) -> a -> b
$! forall a. a -> Int -> Tree a
newBranch a
x (Int -> Int
down Int
sh)) PrimArray Int
newSizesSnoc
        | Bool
otherwise = forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr forall a. Num a => a -> a -> a
- Int
1) (Int -> Tree a -> Tree a
snocTree (Int -> Int
down Int
sh))) PrimArray Int
newSizesAdjust
      where
        len :: Int
len = forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr
        -- snoc the last size + 1
        newSizesSnoc :: PrimArray Int
newSizesSnoc = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
            MutablePrimArray s Int
newArr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray (Int
len forall a. Num a => a -> a -> a
+ Int
1)
            forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s Int
newArr Int
0 PrimArray Int
sizes Int
0 Int
len
            let lastSize :: Int
lastSize = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (Int
len forall a. Num a => a -> a -> a
- Int
1)
            forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s Int
newArr Int
len (Int
lastSize forall a. Num a => a -> a -> a
+ Int
1)
            forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s Int
newArr
        -- adjust the last size with (+ 1)
        newSizesAdjust :: PrimArray Int
newSizesAdjust = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
            MutablePrimArray s Int
newArr <- forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
            forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s Int
newArr Int
0 PrimArray Int
sizes Int
0 Int
len
            let lastSize :: Int
lastSize = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (Int
len forall a. Num a => a -> a -> a
- Int
1)
            forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s Int
newArr (Int
len forall a. Num a => a -> a -> a
- Int
1) (Int
lastSize forall a. Num a => a -> a -> a
+ Int
1)
            forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s Int
newArr
    snocTree Int
_ (Leaf Array a
arr) = forall a. Array a -> Tree a
Leaf forall a b. (a -> b) -> a -> b
$ forall a. Array a -> a -> Array a
A.snoc Array a
arr a
x

    insertShift :: Int
insertShift = 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 size is computed for efficient calculation of the shift in a balanced subtree
    computeShift :: Int -> Int -> Int -> Tree a -> Int
computeShift !Int
sz !Int
sh !Int
min (Balanced Array (Tree a)
_) =
        let newShift :: Int
newShift = (forall b. FiniteBits b => b -> Int
countTrailingZeros Int
sz forall a. Integral a => a -> a -> a
`div` Int
blockShift) forall a. Num a => a -> a -> a
* Int
blockShift
        in if Int
newShift 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 lastIdx :: Int
lastIdx = forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr forall a. Num a => a -> a -> a
- Int
1
            sz' :: Int
sz' = forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes Int
lastIdx forall a. Num a => a -> a -> a
- forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (Int
lastIdx forall a. Num a => a -> a -> a
- Int
1) -- the size of the last subtree
            newMin :: Int
newMin = if forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
sh else Int
min
        in Int -> Int -> Int -> Tree a -> Int
computeShift Int
sz' (Int -> Int
down Int
sh) Int
newMin (forall a. Array a -> a
A.last Array (Tree a)
arr)
    computeShift Int
_ Int
_ Int
min (Leaf Array a
arr) = if forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
0 else Int
min

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

-- splitting

-- the initial @i@ is @n - 1@ -- the index of the last element in the new tree
takeTree :: Int -> Shift -> Tree a -> Tree a
takeTree :: forall a. 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 = forall a. Array a -> Int -> Array a
A.take Array (Tree a)
arr (Int
idx forall a. Num a => a -> a -> a
+ Int
1)
    in forall a. Array (Tree a) -> Tree a
Balanced (forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Int
idx (forall a. 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 = forall a. Array a -> Int -> Array a
A.take Array (Tree a)
arr (Int
idx forall a. Num a => a -> a -> a
+ Int
1)
    in forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Int
idx (forall a. Int -> Int -> Tree a -> Tree a
takeTree Int
subIdx (Int -> Int
down Int
sh)))
takeTree Int
i Int
_ (Leaf Array a
arr) = forall a. Array a -> Tree a
Leaf (forall a. Array a -> Int -> Array a
A.take Array a
arr ((Int
i forall a. Bits a => a -> a -> a
.&. Int
blockMask) forall a. Num a => a -> a -> a
+ Int
1))

dropTree :: Int -> Shift -> Tree a -> Tree a
dropTree :: forall a. 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 = forall a. Array a -> Int -> Array a
A.drop Array (Tree a)
arr Int
idx
    in forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Int
0 (forall a. 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 = forall a. Array a -> Int -> Array a
A.drop Array (Tree a)
arr Int
idx
    in forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Int
0 (forall a. Int -> Int -> Tree a -> Tree a
dropTree Int
subIdx (Int -> Int
down Int
sh)))
dropTree Int
n Int
_ (Leaf Array a
arr) = forall a. Array a -> Tree a
Leaf (forall a. Array a -> Int -> Array a
A.drop Array a
arr (Int
n forall a. Bits a => a -> a -> a
.&. Int
blockMask))

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