{-# 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, replicate
    -- ** 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 qualified Data.List as List
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
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 |>

type Shift = Int

-- Invariant: Children of a Balanced node are always balanced.
-- A Leaf node is considered balanced.
-- Nodes are always non-empty.
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 :: Shift
blockShift = Shift
4

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

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

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

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

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

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

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

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

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

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

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

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

instance (Show a) => Show (Vector a) where
    showsPrec :: Shift -> Vector a -> ShowS
showsPrec = Shift -> Vector a -> ShowS
forall (f :: * -> *) a. (Show1 f, Show a) => Shift -> 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 -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Vector a
v1 Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Vector b -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
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 -> Vector a -> Vector a
(<>) = Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
(><)
    stimes :: b -> Vector a -> Vector a
stimes = b -> Vector a -> Vector a
forall b a. (Integral b, Monoid a) => b -> a -> a
stimesMonoid

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 Shift
_ Shift
_ 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 Shift
_) 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 Shift
_ Shift
_ 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 Shift
_) = (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 Shift
_ Shift
_ 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 Shift
_) 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 Shift
_ Shift
_ 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 Shift
_) = (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

    length :: Vector a -> Shift
length Vector a
Empty = Shift
0
    length (Root Shift
s Shift
_ Tree a
_) = Shift
s

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

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

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 = Shift -> a -> Vector a
forall a. Shift -> a -> Vector a
replicate (Vector b -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Vector b
v) a
x

instance FunctorWithIndex Int Vector where
    imap :: (Shift -> a -> b) -> Vector a -> Vector b
imap Shift -> 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) -> Shift -> Identity (Vector b)
forall (f :: * -> *) a. Indexed f a -> Shift -> 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 ((Shift -> WithIndex (Identity b)) -> Indexed Identity b
forall (f :: * -> *) a. (Shift -> WithIndex (f a)) -> Indexed f a
Indexed ((Shift -> WithIndex (Identity b)) -> Indexed Identity b)
-> (a -> Shift -> WithIndex (Identity b))
-> a
-> Indexed Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Shift -> WithIndex (Identity b)
f') Vector a
v) Shift
0
      where
        f' :: a -> Shift -> WithIndex (Identity b)
f' a
x !Shift
i = Shift -> Identity b -> WithIndex (Identity b)
forall a. Shift -> a -> WithIndex a
WithIndex (Shift
i Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1) (b -> Identity b
forall a. a -> Identity a
Identity (Shift -> a -> b
f Shift
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 Shift
size Shift
sh Tree a
tree) = Shift -> Shift -> Tree b -> Vector b
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
size Shift
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 Shift
sizes) = (Array (Tree b) -> PrimArray Shift -> Tree b)
-> PrimArray Shift -> Array (Tree b) -> Tree b
forall a b c. (a -> b -> c) -> b -> a -> c
flip Array (Tree b) -> PrimArray Shift -> Tree b
forall a. Array (Tree a) -> PrimArray Shift -> Tree a
Unbalanced PrimArray Shift
sizes (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 (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
    {-# INLINE traverse #-}

instance TraversableWithIndex Int Vector where
    itraverse :: (Shift -> a -> f b) -> Vector a -> f (Vector b)
itraverse Shift -> a -> f b
f Vector a
v = Indexed f (Vector b) -> Shift -> f (Vector b)
forall (f :: * -> *) a. Indexed f a -> Shift -> 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 ((Shift -> WithIndex (f b)) -> Indexed f b
forall (f :: * -> *) a. (Shift -> WithIndex (f a)) -> Indexed f a
Indexed ((Shift -> WithIndex (f b)) -> Indexed f b)
-> (a -> Shift -> WithIndex (f b)) -> a -> Indexed f b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Shift -> WithIndex (f b)
f') Vector a
v) Shift
0
      where
        f' :: a -> Shift -> WithIndex (f b)
f' a
x !Shift
i = Shift -> f b -> WithIndex (f b)
forall a. Shift -> a -> WithIndex a
WithIndex (Shift
i Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1) (Shift -> a -> f b
f Shift
i a
x)
    {-# INLINE itraverse #-}

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
>< Shift -> a -> Vector a
forall a. Shift -> a -> Vector a
replicate (Vector b -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
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
$ (Shift -> a) -> [Shift] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Shift
i -> let x :: a
x = Shift -> Vector a -> a
forall a. HasCallStack => Shift -> Vector a -> a
index Shift
i (a -> Vector a
f a
x) in a
x) [Shift
0..Vector a -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length (a -> Vector a
f a
forall a. a
err) Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
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 = Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
1 Shift
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] -> Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root (Shift -> Tree a -> Shift
forall a. Shift -> Tree a -> Shift
treeSize Shift
0 Tree a
tree) Shift
0 Tree a
tree -- tree is a single leaf
    [Tree a]
ls' -> Shift -> [Tree a] -> Vector a
forall a. Shift -> [Tree a] -> Vector a
iterateNodes Shift
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 <- Shift -> ST s (Buffer s a)
forall s a. Shift -> ST s (Buffer s a)
Buffer.new Shift
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
                let !x :: a
x = Array a -> a
f Array a
result
                [a] -> ST s [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [a
x]
            loop (a
t : [a]
ts) = do
                Shift
size <- Buffer s a -> ST s Shift
forall s a. Buffer s a -> ST s Shift
Buffer.size Buffer s a
buffer
                if Shift
size Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Shift
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
                    let !x :: a
x = Array a -> a
f Array a
result
                    [a] -> ST s [a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
x 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 :: Shift -> [Tree a] -> Vector a
iterateNodes Shift
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] -> Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root (Shift -> Tree a -> Shift
forall a. Shift -> Tree a -> Shift
treeSize Shift
sh Tree a
tree) Shift
sh Tree a
tree
        [Tree a]
trees' -> Shift -> [Tree a] -> Vector a
iterateNodes (Shift -> Shift
up Shift
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 :: Shift -> a -> Vector a
replicate Shift
n a
x
    | Shift
n Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
<= Shift
0 = Vector a
forall a. Vector a
Empty
    | Shift
n Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
<= Shift
blockSize = Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
n Shift
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
$ Shift -> a -> Array a
forall a. Shift -> a -> Array a
A.replicate Shift
n a
x)
    | Bool
otherwise = Shift -> Tree a -> Tree a -> Vector a
iterateNodes Shift
blockShift (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
$ Shift -> a -> Array a
forall a. Shift -> a -> Array a
A.replicate Shift
blockSize a
x) (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
$ Shift -> a -> Array a
forall a. Shift -> a -> Array a
A.replicate (Shift
lastIdx Shift -> Shift -> Shift
forall a. Bits a => a -> a -> a
.&. Shift
blockMask Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1) a
x)
  where
    lastIdx :: Shift
lastIdx = Shift
n Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
1

    -- @full@ is a full subtree, @rest@ is the last subtree
    iterateNodes :: Shift -> Tree a -> Tree a -> Vector a
iterateNodes !Shift
sh !Tree a
full !Tree a
rest =
        let subtreesM1 :: Shift
subtreesM1 = Shift
lastIdx Shift -> Shift -> Shift
forall a. Bits a => a -> Shift -> a
`unsafeShiftR` Shift
sh -- the number of subtrees minus 1
            full' :: Tree a
full' = 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
$ Shift -> Tree a -> Array (Tree a)
forall a. Shift -> a -> Array a
A.replicate Shift
blockSize Tree a
full
            rest' :: Tree a
rest' = 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
$ Shift -> Tree a -> Tree a -> Array (Tree a)
forall a. Shift -> a -> a -> Array a
A.replicateSnoc (Shift
subtreesM1 Shift -> Shift -> Shift
forall a. Bits a => a -> a -> a
.&. Shift
blockMask) Tree a
full Tree a
rest
        in if Shift
subtreesM1 Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
< Shift
blockSize then Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
n Shift
sh Tree a
rest' else Shift -> Tree a -> Tree a -> Vector a
iterateNodes (Shift -> Shift
up Shift
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 :: Shift -> Vector a -> Maybe a
lookup Shift
_ Vector a
Empty = Maybe a
forall a. Maybe a
Nothing
lookup Shift
i (Root Shift
size Shift
sh Tree a
tree)
    | Shift
i Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
< Shift
0 Bool -> Bool -> Bool
|| Shift
i Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
>= Shift
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
$ Shift -> Shift -> Tree a -> a
forall p. Shift -> Shift -> Tree p -> p
lookupTree Shift
i Shift
sh Tree a
tree
  where
    lookupTree :: Shift -> Shift -> Tree p -> p
lookupTree Shift
i Shift
sh (Balanced Array (Tree p)
arr) = Shift -> Shift -> Tree p -> p
lookupTree Shift
i (Shift -> Shift
down Shift
sh) (Array (Tree p) -> Shift -> Tree p
forall a. Array a -> Shift -> a
A.index Array (Tree p)
arr (Shift -> Shift -> Shift
radixIndex Shift
i Shift
sh))
    lookupTree Shift
i Shift
sh (Unbalanced Array (Tree p)
arr PrimArray Shift
sizes) =
        let (Shift
idx, Shift
subIdx) = PrimArray Shift -> Shift -> Shift -> (Shift, Shift)
relaxedRadixIndex PrimArray Shift
sizes Shift
i Shift
sh
        in Shift -> Shift -> Tree p -> p
lookupTree Shift
subIdx (Shift -> Shift
down Shift
sh) (Array (Tree p) -> Shift -> Tree p
forall a. Array a -> Shift -> a
A.index Array (Tree p)
arr Shift
idx)
    lookupTree Shift
i Shift
_ (Leaf Array p
arr) = Array p -> Shift -> p
forall a. Array a -> Shift -> a
A.index Array p
arr (Shift
i Shift -> Shift -> Shift
forall a. Bits a => a -> a -> a
.&. Shift
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 :: Shift -> Vector a -> a
index Shift
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
. Shift -> Vector a -> Maybe a
forall a. Shift -> Vector a -> Maybe a
lookup Shift
i
{-# INLINE index #-}

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

-- | \(O(\log n)\). A flipped version of 'index'.
(!) :: HasCallStack => Vector a -> Int -> a
(!) = (Shift -> Vector a -> a) -> Vector a -> Shift -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip Shift -> Vector a -> a
forall a. HasCallStack => Shift -> 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 :: Shift -> a -> Vector a -> Vector a
update Shift
_ a
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
update Shift
i a
x v :: Vector a
v@(Root Shift
size Shift
sh Tree a
tree)
    | Shift
i Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
< Shift
0 Bool -> Bool -> Bool
|| Shift
i Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
>= Shift
size = Vector a
v  -- index out of range
    | Bool
otherwise = Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
size Shift
sh (Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
i Shift
sh Tree a
tree)
  where
    adjustTree :: Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
i Shift
sh (Balanced Array (Tree a)
arr) = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (Shift -> Shift -> Shift
radixIndex Shift
i Shift
sh) (Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
i (Shift -> Shift
down Shift
sh)))
    adjustTree Shift
i Shift
sh (Unbalanced Array (Tree a)
arr PrimArray Shift
sizes) =
        let (Shift
idx, Shift
subIdx) = PrimArray Shift -> Shift -> Shift -> (Shift, Shift)
relaxedRadixIndex PrimArray Shift
sizes Shift
i Shift
sh
        in Array (Tree a) -> PrimArray Shift -> Tree a
forall a. Array (Tree a) -> PrimArray Shift -> Tree a
Unbalanced (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Shift
idx (Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
subIdx (Shift -> Shift
down Shift
sh))) PrimArray Shift
sizes
    adjustTree Shift
i Shift
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Shift -> a -> Array a
forall a. Array a -> Shift -> a -> Array a
A.update Array a
arr (Shift
i Shift -> Shift -> Shift
forall a. Bits a => a -> a -> a
.&. Shift
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 :: Shift -> (a -> a) -> Vector a -> Vector a
adjust Shift
_ a -> a
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
adjust Shift
i a -> a
f v :: Vector a
v@(Root Shift
size Shift
sh Tree a
tree)
    | Shift
i Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
< Shift
0 Bool -> Bool -> Bool
|| Shift
i Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
>= Shift
size = Vector a
v  -- index out of range
    | Bool
otherwise = Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
size Shift
sh (Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
i Shift
sh Tree a
tree)
  where
    adjustTree :: Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
i Shift
sh (Balanced Array (Tree a)
arr) = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (Shift -> Shift -> Shift
radixIndex Shift
i Shift
sh) (Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
i (Shift -> Shift
down Shift
sh)))
    adjustTree Shift
i Shift
sh (Unbalanced Array (Tree a)
arr PrimArray Shift
sizes) =
        let (Shift
idx, Shift
subIdx) = PrimArray Shift -> Shift -> Shift -> (Shift, Shift)
relaxedRadixIndex PrimArray Shift
sizes Shift
i Shift
sh
        in Array (Tree a) -> PrimArray Shift -> Tree a
forall a. Array (Tree a) -> PrimArray Shift -> Tree a
Unbalanced (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Shift
idx (Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
subIdx (Shift -> Shift
down Shift
sh))) PrimArray Shift
sizes
    adjustTree Shift
i Shift
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Shift -> (a -> a) -> Array a
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust Array a
arr (Shift
i Shift -> Shift -> Shift
forall a. Bits a => a -> a -> a
.&. Shift
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' :: Shift -> (a -> a) -> Vector a -> Vector a
adjust' Shift
_ a -> a
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
adjust' Shift
i a -> a
f v :: Vector a
v@(Root Shift
size Shift
sh Tree a
tree)
    | Shift
i Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
< Shift
0 Bool -> Bool -> Bool
|| Shift
i Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
>= Shift
size = Vector a
v  -- index out of range
    | Bool
otherwise = Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
size Shift
sh (Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
i Shift
sh Tree a
tree)
  where
    adjustTree :: Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
i Shift
sh (Balanced Array (Tree a)
arr) = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (Shift -> Shift -> Shift
radixIndex Shift
i Shift
sh) (Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
i (Shift -> Shift
down Shift
sh)))
    adjustTree Shift
i Shift
sh (Unbalanced Array (Tree a)
arr PrimArray Shift
sizes) =
        let (Shift
idx, Shift
subIdx) = PrimArray Shift -> Shift -> Shift -> (Shift, Shift)
relaxedRadixIndex PrimArray Shift
sizes Shift
i Shift
sh
        in Array (Tree a) -> PrimArray Shift -> Tree a
forall a. Array (Tree a) -> PrimArray Shift -> Tree a
Unbalanced (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Shift
idx (Shift -> Shift -> Tree a -> Tree a
adjustTree Shift
subIdx (Shift -> Shift
down Shift
sh))) PrimArray Shift
sizes
    adjustTree Shift
i Shift
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Shift -> (a -> a) -> Array a
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array a
arr (Shift
i Shift -> Shift -> Shift
forall a. Bits a => a -> a -> a
.&. Shift
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 Shift
size Shift
sh Tree a
tree) = Shift -> Shift -> Tree b -> Vector b
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
size Shift
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 Shift
sizes) = Array (Tree b) -> PrimArray Shift -> Tree b
forall a. Array (Tree a) -> PrimArray Shift -> 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 Shift
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 = (a -> b -> (a, b)) -> Vector a -> Vector b -> Vector (a, b)
forall a b c. (a -> b -> c) -> Vector a -> Vector b -> Vector c
zipWith (,)

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

        viewrArr :: Array b -> (Array b, b)
viewrArr Array b
arr = (Array b -> Shift -> Array b
forall a. Array a -> Shift -> Array a
A.take Array b
arr (Array b -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Array b
arr Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
1), Array b -> b
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 :: Shift
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
mergeRebalance !Shift
sh !Array (Tree a)
left !Array (Tree a)
center !Array (Tree a)
right
        | Shift
sh Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Shift
blockShift = (Tree a -> Array a) -> (Array a -> Tree a) -> Array (Tree a)
forall t.
(Tree a -> Array t) -> (Array t -> Tree a) -> Array (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) -> Array (Tree a)
forall t.
(Tree a -> Array t) -> (Array t -> Tree a) -> Array (Tree a)
mergeRebalance' Tree a -> Array (Tree a)
forall a. Tree a -> Array (Tree a)
treeToArray (Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes (Shift -> Shift
down Shift
sh))
      where
        mergeRebalance' :: (Tree a -> A.Array t) -> (A.Array t -> Tree a) -> A.Array (Tree a)
        mergeRebalance' :: (Tree a -> Array t) -> (Array t -> Tree a) -> Array (Tree a)
mergeRebalance' Tree a -> Array t
extract Array t -> Tree a
construct = (forall s. ST s (Array (Tree a))) -> Array (Tree a)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array (Tree a))) -> Array (Tree a))
-> (forall s. ST s (Array (Tree a))) -> Array (Tree a)
forall a b. (a -> b) -> a -> b
$ do
            Buffer s (Tree a)
newRoot <- Shift -> ST s (Buffer s (Tree a))
forall s a. Shift -> ST s (Buffer s a)
Buffer.new Shift
blockSize
            Buffer s (Tree a)
newSubtree <- Shift -> ST s (Buffer s (Tree a))
forall s a. Shift -> ST s (Buffer s a)
Buffer.new Shift
blockSize
            Buffer s t
newNode <- Shift -> ST s (Buffer s t)
forall s a. Shift -> ST s (Buffer s a)
Buffer.new Shift
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
                    Shift
lenNode <- Buffer s t -> ST s Shift
forall s a. Buffer s a -> ST s Shift
Buffer.size Buffer s t
newNode
                    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Shift
lenNode Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Shift
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
                        Shift
lenSubtree <- Buffer s (Tree a) -> ST s Shift
forall s a. Buffer s a -> ST s Shift
Buffer.size Buffer s (Tree a)
newSubtree
                        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Shift
lenSubtree Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Shift
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 (Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
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 (Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
sh) Buffer s (Tree a)
newSubtree Buffer s (Tree a)
newRoot
            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 (a -> ST s ()) -> a -> ST s ()
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 <| :: a -> Vector a -> Vector a
<| Vector a
Empty = a -> Vector a
forall a. a -> Vector a
singleton a
x
a
x <| Root Shift
size Shift
sh Tree a
tree
    | Shift
insertShift Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
> Shift
sh = Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root (Shift
size Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1) Shift
insertShift (Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
insertShift (let !new :: Tree a
new = a -> Shift -> Tree a
forall a. a -> Shift -> Tree a
newBranch a
x Shift
sh in Tree a -> Tree a -> Array (Tree a)
forall a. a -> a -> Array a
A.from2 Tree a
new Tree a
tree))
    | Bool
otherwise = Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root (Shift
size Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1) Shift
sh (Shift -> Tree a -> Tree a
consTree Shift
sh Tree a
tree)
  where
    consTree :: Shift -> Tree a -> Tree a
consTree Shift
sh (Balanced Array (Tree a)
arr)
        | Shift
sh Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Shift
insertShift = Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
sh (Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.cons Array (Tree a)
arr (Tree a -> Array (Tree a)) -> Tree a -> Array (Tree a)
forall a b. (a -> b) -> a -> b
$! a -> Shift -> Tree a
forall a. a -> Shift -> Tree a
newBranch a
x (Shift -> Shift
down Shift
sh))
        | Bool
otherwise = Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
sh (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Shift
0 (Shift -> Tree a -> Tree a
consTree (Shift -> Shift
down Shift
sh)))
    consTree Shift
sh (Unbalanced Array (Tree a)
arr PrimArray Shift
_)
        | Shift
sh Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Shift
insertShift = Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
sh (Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.cons Array (Tree a)
arr (Tree a -> Array (Tree a)) -> Tree a -> Array (Tree a)
forall a b. (a -> b) -> a -> b
$! a -> Shift -> Tree a
forall a. a -> Shift -> Tree a
newBranch a
x (Shift -> Shift
down Shift
sh))
        | Bool
otherwise = Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
sh (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Shift
0 (Shift -> Tree a -> Tree a
consTree (Shift -> Shift
down Shift
sh)))
    consTree Shift
_ (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 :: Shift
insertShift = Shift -> Shift -> Shift -> Tree a -> Shift
forall a. Shift -> Shift -> Shift -> Tree a -> Shift
computeShift Shift
size Shift
sh (Shift -> Shift
up Shift
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 :: Shift -> Shift -> Shift -> Tree a -> Shift
computeShift !Shift
sz !Shift
sh !Shift
min (Balanced Array (Tree a)
_) =
        let newShift :: Shift
newShift = (Shift -> Shift
log2 Shift
sz Shift -> Shift -> Shift
forall a. Integral a => a -> a -> a
`div` Shift
blockShift) Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
* Shift
blockShift
        in if Shift
newShift Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
> Shift
sh then Shift
min else Shift
newShift
    computeShift Shift
_ Shift
sh Shift
min (Unbalanced Array (Tree a)
arr PrimArray Shift
sizes) =
        let sz' :: Shift
sz' = PrimArray Shift -> Shift -> Shift
forall a. Prim a => PrimArray a -> Shift -> a
indexPrimArray PrimArray Shift
sizes Shift
0 -- the size of the first subtree
            newMin :: Shift
newMin = if Array (Tree a) -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Array (Tree a)
arr Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
< Shift
blockSize then Shift
sh else Shift
min
        in Shift -> Shift -> Shift -> Tree a -> Shift
computeShift Shift
sz' (Shift -> Shift
down Shift
sh) Shift
newMin (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
arr)
    computeShift Shift
_ Shift
_ Shift
min (Leaf Array a
arr) = if Array a -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Array a
arr Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
< Shift
blockSize then Shift
0 else Shift
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 Shift
size Shift
sh Tree a
tree |> a
x
    | Shift
insertShift Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
> Shift
sh = Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root (Shift
size Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1) Shift
insertShift (Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
insertShift (Tree a -> Tree a -> Array (Tree a)
forall a. a -> a -> Array a
A.from2 Tree a
tree (Tree a -> Array (Tree a)) -> Tree a -> Array (Tree a)
forall a b. (a -> b) -> a -> b
$! a -> Shift -> Tree a
forall a. a -> Shift -> Tree a
newBranch a
x Shift
sh))
    | Bool
otherwise = Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root (Shift
size Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1) Shift
sh (Shift -> Tree a -> Tree a
snocTree Shift
sh Tree a
tree)
  where
    snocTree :: Shift -> Tree a -> Tree a
snocTree Shift
sh (Balanced Array (Tree a)
arr)
        | Shift
sh Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Shift
insertShift = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.snoc Array (Tree a)
arr (Tree a -> Array (Tree a)) -> Tree a -> Array (Tree a)
forall a b. (a -> b) -> a -> b
$! a -> Shift -> Tree a
forall a. a -> Shift -> Tree a
newBranch a
x (Shift -> Shift
down Shift
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) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (Array (Tree a) -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Array (Tree a)
arr Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
1) (Shift -> Tree a -> Tree a
snocTree (Shift -> Shift
down Shift
sh))
    snocTree Shift
sh (Unbalanced Array (Tree a)
arr PrimArray Shift
sizes)
        | Shift
sh Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Shift
insertShift = Array (Tree a) -> PrimArray Shift -> Tree a
forall a. Array (Tree a) -> PrimArray Shift -> Tree a
Unbalanced (Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.snoc Array (Tree a)
arr (Tree a -> Array (Tree a)) -> Tree a -> Array (Tree a)
forall a b. (a -> b) -> a -> b
$! a -> Shift -> Tree a
forall a. a -> Shift -> Tree a
newBranch a
x (Shift -> Shift
down Shift
sh)) PrimArray Shift
newSizesSnoc
        | Bool
otherwise = Array (Tree a) -> PrimArray Shift -> Tree a
forall a. Array (Tree a) -> PrimArray Shift -> Tree a
Unbalanced (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr (Array (Tree a) -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Array (Tree a)
arr Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
1) (Shift -> Tree a -> Tree a
snocTree (Shift -> Shift
down Shift
sh))) PrimArray Shift
newSizesAdjust
      where
        -- snoc the last size + 1
        newSizesSnoc :: PrimArray Shift
newSizesSnoc = (forall s. ST s (PrimArray Shift)) -> PrimArray Shift
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (PrimArray Shift)) -> PrimArray Shift)
-> (forall s. ST s (PrimArray Shift)) -> PrimArray Shift
forall a b. (a -> b) -> a -> b
$ do
            let lenSizes :: Shift
lenSizes = PrimArray Shift -> Shift
forall a. Prim a => PrimArray a -> Shift
sizeofPrimArray PrimArray Shift
sizes
            MutablePrimArray s Shift
newArr <- Shift -> ST s (MutablePrimArray (PrimState (ST s)) Shift)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Shift -> m (MutablePrimArray (PrimState m) a)
newPrimArray (Shift
lenSizes Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1)
            MutablePrimArray (PrimState (ST s)) Shift
-> Shift -> PrimArray Shift -> Shift -> Shift -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Shift -> PrimArray a -> Shift -> Shift -> m ()
copyPrimArray MutablePrimArray s Shift
MutablePrimArray (PrimState (ST s)) Shift
newArr Shift
0 PrimArray Shift
sizes Shift
0 Shift
lenSizes
            let lastSize :: Shift
lastSize = PrimArray Shift -> Shift -> Shift
forall a. Prim a => PrimArray a -> Shift -> a
indexPrimArray PrimArray Shift
sizes (Shift
lenSizes Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
1)
            MutablePrimArray (PrimState (ST s)) Shift
-> Shift -> Shift -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Shift -> a -> m ()
writePrimArray MutablePrimArray s Shift
MutablePrimArray (PrimState (ST s)) Shift
newArr Shift
lenSizes (Shift
lastSize Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1)
            MutablePrimArray (PrimState (ST s)) Shift -> ST s (PrimArray Shift)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s Shift
MutablePrimArray (PrimState (ST s)) Shift
newArr
        -- adjust the last size with (+ 1)
        newSizesAdjust :: PrimArray Shift
newSizesAdjust = (forall s. ST s (PrimArray Shift)) -> PrimArray Shift
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (PrimArray Shift)) -> PrimArray Shift)
-> (forall s. ST s (PrimArray Shift)) -> PrimArray Shift
forall a b. (a -> b) -> a -> b
$ do
            let lenSizes :: Shift
lenSizes = PrimArray Shift -> Shift
forall a. Prim a => PrimArray a -> Shift
sizeofPrimArray PrimArray Shift
sizes
            MutablePrimArray s Shift
newArr <- Shift -> ST s (MutablePrimArray (PrimState (ST s)) Shift)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Shift -> m (MutablePrimArray (PrimState m) a)
newPrimArray Shift
lenSizes
            MutablePrimArray (PrimState (ST s)) Shift
-> Shift -> PrimArray Shift -> Shift -> Shift -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Shift -> PrimArray a -> Shift -> Shift -> m ()
copyPrimArray MutablePrimArray s Shift
MutablePrimArray (PrimState (ST s)) Shift
newArr Shift
0 PrimArray Shift
sizes Shift
0 Shift
lenSizes
            let lastSize :: Shift
lastSize = PrimArray Shift -> Shift -> Shift
forall a. Prim a => PrimArray a -> Shift -> a
indexPrimArray PrimArray Shift
sizes (Shift
lenSizes Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
1)
            MutablePrimArray (PrimState (ST s)) Shift
-> Shift -> Shift -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Shift -> a -> m ()
writePrimArray MutablePrimArray s Shift
MutablePrimArray (PrimState (ST s)) Shift
newArr (Shift
lenSizes Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
1) (Shift
lastSize Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1)
            MutablePrimArray (PrimState (ST s)) Shift -> ST s (PrimArray Shift)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s Shift
MutablePrimArray (PrimState (ST s)) Shift
newArr
    snocTree Shift
_ (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 :: Shift
insertShift = Shift -> Shift -> Shift -> Tree a -> Shift
forall a. Shift -> Shift -> Shift -> Tree a -> Shift
computeShift Shift
size Shift
sh (Shift -> Shift
up Shift
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 :: Shift -> Shift -> Shift -> Tree a -> Shift
computeShift !Shift
sz !Shift
sh !Shift
min (Balanced Array (Tree a)
_) =
        let newShift :: Shift
newShift = (Shift -> Shift
forall b. FiniteBits b => b -> Shift
countTrailingZeros Shift
sz Shift -> Shift -> Shift
forall a. Integral a => a -> a -> a
`div` Shift
blockShift) Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
* Shift
blockShift
        in if Shift
newShift Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
> Shift
sh then Shift
min else Shift
newShift
    computeShift Shift
_ Shift
sh Shift
min (Unbalanced Array (Tree a)
arr PrimArray Shift
sizes) =
        let lastIdx :: Shift
lastIdx = PrimArray Shift -> Shift
forall a. Prim a => PrimArray a -> Shift
sizeofPrimArray PrimArray Shift
sizes Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
1
            sz' :: Shift
sz' = PrimArray Shift -> Shift -> Shift
forall a. Prim a => PrimArray a -> Shift -> a
indexPrimArray PrimArray Shift
sizes Shift
lastIdx Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- PrimArray Shift -> Shift -> Shift
forall a. Prim a => PrimArray a -> Shift -> a
indexPrimArray PrimArray Shift
sizes (Shift
lastIdx Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
1) -- the size of the last subtree
            newMin :: Shift
newMin = if Array (Tree a) -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Array (Tree a)
arr Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
< Shift
blockSize then Shift
sh else Shift
min
        in Shift -> Shift -> Shift -> Tree a -> Shift
computeShift Shift
sz' (Shift -> Shift
down Shift
sh) Shift
newMin (Array (Tree a) -> Tree a
forall a. Array a -> a
A.last Array (Tree a)
arr)
    computeShift Shift
_ Shift
_ Shift
min (Leaf Array a
arr) = if Array a -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Array a
arr Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
< Shift
blockSize then Shift
0 else Shift
min

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

-- 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 :: Shift -> Vector a -> Vector a
take Shift
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
take Shift
n v :: Vector a
v@(Root Shift
size Shift
sh Tree a
tree)
    | Shift
n Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
<= Shift
0 = Vector a
forall a. Vector a
empty
    | Shift
n Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
>= Shift
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
$ Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
n Shift
sh (Shift -> Shift -> Tree a -> Tree a
forall a. Shift -> Shift -> Tree a -> Tree a
takeTree (Shift
n Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
1) Shift
sh Tree a
tree)
  where
    -- the initial @i@ is @n - 1@ -- the index of the last element in the new tree
    takeTree :: Shift -> Shift -> Tree a -> Tree a
takeTree Shift
i Shift
sh (Balanced Array (Tree a)
arr) =
        let idx :: Shift
idx = Shift -> Shift -> Shift
radixIndex Shift
i Shift
sh
            newArr :: Array (Tree a)
newArr = Array (Tree a) -> Shift -> Array (Tree a)
forall a. Array a -> Shift -> Array a
A.take Array (Tree a)
arr (Shift
idx Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1)
        in Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Shift
idx (Shift -> Shift -> Tree a -> Tree a
takeTree Shift
i (Shift -> Shift
down Shift
sh)))
    takeTree Shift
i Shift
sh (Unbalanced Array (Tree a)
arr PrimArray Shift
sizes) =
        let (Shift
idx, Shift
subIdx) = PrimArray Shift -> Shift -> Shift -> (Shift, Shift)
relaxedRadixIndex PrimArray Shift
sizes Shift
i Shift
sh
            newArr :: Array (Tree a)
newArr = Array (Tree a) -> Shift -> Array (Tree a)
forall a. Array a -> Shift -> Array a
A.take Array (Tree a)
arr (Shift
idx Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
1)
        in Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
sh (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Shift
idx (Shift -> Shift -> Tree a -> Tree a
takeTree Shift
subIdx (Shift -> Shift
down Shift
sh)))
    takeTree Shift
i Shift
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Shift -> Array a
forall a. Array a -> Shift -> Array a
A.take Array a
arr ((Shift
i Shift -> Shift -> Shift
forall a. Bits a => a -> a -> a
.&. Shift
blockMask) Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
+ Shift
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 :: Shift -> Vector a -> Vector a
drop Shift
_ Vector a
Empty = Vector a
forall a. Vector a
Empty
drop Shift
n v :: Vector a
v@(Root Shift
size Shift
sh Tree a
tree)
    | Shift
n Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
<= Shift
0 = Vector a
v
    | Shift
n Shift -> Shift -> Bool
forall a. Ord a => a -> a -> Bool
>= Shift
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
$ Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root (Shift
size Shift -> Shift -> Shift
forall a. Num a => a -> a -> a
- Shift
n) Shift
sh (Shift -> Shift -> Tree a -> Tree a
forall a. Shift -> Shift -> Tree a -> Tree a
dropTree Shift
n Shift
sh Tree a
tree)
  where
    dropTree :: Shift -> Shift -> Tree a -> Tree a
dropTree Shift
n Shift
sh (Balanced Array (Tree a)
arr) =
        let idx :: Shift
idx = Shift -> Shift -> Shift
radixIndex Shift
n Shift
sh
            newArr :: Array (Tree a)
newArr = Array (Tree a) -> Shift -> Array (Tree a)
forall a. Array a -> Shift -> Array a
A.drop Array (Tree a)
arr Shift
idx
        in Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
sh (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Shift
0 (Shift -> Shift -> Tree a -> Tree a
dropTree Shift
n (Shift -> Shift
down Shift
sh)))
    dropTree Shift
n Shift
sh (Unbalanced Array (Tree a)
arr PrimArray Shift
sizes) =
        let (Shift
idx, Shift
subIdx) = PrimArray Shift -> Shift -> Shift -> (Shift, Shift)
relaxedRadixIndex PrimArray Shift
sizes Shift
n Shift
sh
            newArr :: Array (Tree a)
newArr = Array (Tree a) -> Shift -> Array (Tree a)
forall a. Array a -> Shift -> Array a
A.drop Array (Tree a)
arr Shift
idx
        in Shift -> Array (Tree a) -> Tree a
forall a. Shift -> Array (Tree a) -> Tree a
computeSizes Shift
sh (Array (Tree a) -> Shift -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Shift -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Shift
0 (Shift -> Shift -> Tree a -> Tree a
dropTree Shift
subIdx (Shift -> Shift
down Shift
sh)))
    dropTree Shift
n Shift
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Shift -> Array a
forall a. Array a -> Shift -> Array a
A.drop Array a
arr (Shift
n Shift -> Shift -> Shift
forall a. Bits a => a -> a -> a
.&. Shift
blockMask))

normalize :: Vector a -> Vector a
normalize :: Vector a -> Vector a
normalize (Root Shift
size Shift
sh (Balanced Array (Tree a)
arr))
    | Array (Tree a) -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Array (Tree a)
arr Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Shift
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
$ Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
size (Shift -> Shift
down Shift
sh) (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
arr)
normalize (Root Shift
size Shift
sh (Unbalanced Array (Tree a)
arr PrimArray Shift
_))
    | Array (Tree a) -> Shift
forall (t :: * -> *) a. Foldable t => t a -> Shift
length Array (Tree a)
arr Shift -> Shift -> Bool
forall a. Eq a => a -> a -> Bool
== Shift
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
$ Shift -> Shift -> Tree a -> Vector a
forall a. Shift -> Shift -> Tree a -> Vector a
Root Shift
size (Shift -> Shift
down Shift
sh) (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
arr)
normalize Vector a
v = Vector a
v