{-# 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
    , findIndexL, findIndexR, findIndicesL, findIndicesR
    -- * Transformations
    , map, map', reverse
    -- * Zipping and unzipping
    , zip, zipWith, unzip, unzipWith
    ) where

#if !(MIN_VERSION_base(4,18,0))
import Control.Applicative (Alternative, liftA2)
#else
import Control.Applicative (Alternative)
#endif
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 Int -> Int -> Int
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

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

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

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

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
_) = [Char] -> Array (Tree 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 = Int -> Int -> Tree a -> Int
forall {a}. Int -> Int -> Tree a -> Int
go Int
0
  where
    go :: Int -> Int -> Tree a -> Int
go !Int
acc !Int
_ (Leaf Array a
arr) = Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Array a -> Int
forall a. Array a -> Int
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (Array (Tree a) -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    go Int
acc Int
sh (Balanced Array (Tree a)
arr) =
        let i :: Int
i = Array (Tree a) -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
        in Int -> Int -> Tree a -> Int
go (Int
acc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
sh)) (Int -> Int
down Int
sh) (Array (Tree a) -> Int -> Tree a
forall a. Array a -> Int -> a
A.index Array (Tree a)
arr Int
i)
{-# INLINE treeSize #-}

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

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

    lenM1 :: Int
lenM1 = Int
len Int -> Int -> Int
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
lenM1 = Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh) Tree a
subtree Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
maxSize Bool -> Bool -> Bool
&& Int -> Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            | Bool
otherwise = Tree a -> Bool
forall a. Tree a -> Bool
treeBalanced Tree a
subtree
          where
            subtree :: Tree a
subtree = Array (Tree a) -> Int -> Tree a
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros Int
x
  where
    bitSizeMinus1 :: Int
bitSizeMinus1 = Int -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Int
0 :: Int) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

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 = (Int -> [a] -> ShowS) -> [Char] -> Int -> [a] -> ShowS
forall a. (Int -> a -> ShowS) -> [Char] -> Int -> a -> ShowS
showsUnaryWith ((Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> [a] -> ShowS
forall a.
(Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> [a] -> ShowS
forall (f :: * -> *) a.
Show1 f =>
(Int -> a -> ShowS) -> ([a] -> ShowS) -> Int -> f a -> ShowS
liftShowsPrec Int -> a -> ShowS
sp [a] -> ShowS
sl) [Char]
"fromList" Int
p (Vector a -> [a]
forall a. Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v)

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

instance Read1 Vector where
    liftReadPrec :: forall a. 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 a. 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 :: forall a. 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 :: 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 = Vector a -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector b -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector b
v2 Bool -> Bool -> Bool
&& (a -> b -> Bool) -> [a] -> [b] -> Bool
forall a b. (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 a. Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (Vector b -> [b]
forall a. Vector a -> [a]
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 :: forall a b.
(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 a b. (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 a. Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (Vector b -> [b]
forall a. Vector a -> [a]
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
(><)
    sconcat :: NonEmpty (Vector a) -> Vector a
sconcat (Vector a
v :| [Vector a]
vs) = Vector a
v Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector 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 = 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
    mconcat :: [Vector a] -> Vector a
mconcat = (Vector a -> Vector a -> Vector a)
-> Vector a -> [Vector a] -> Vector a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Vector a -> Vector a -> Vector a
forall a. Vector 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' = (Tree a -> b -> b) -> b -> Array (Tree a) -> b
forall a b. (a -> b -> b) -> b -> Array a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Tree a -> b -> b
foldrTree b
acc' Array (Tree a)
arr
        foldrTree (Unbalanced Array (Tree a)
arr PrimArray Int
_) b
acc' = (Tree a -> b -> b) -> b -> Array (Tree a) -> b
forall a b. (a -> b -> b) -> b -> Array 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 a b. (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 :: 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) = (b -> Tree a -> b) -> b -> Array (Tree a) -> b
forall b a. (b -> a -> b) -> b -> Array a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl b -> Tree a -> b
foldlTree b
acc' Array (Tree a)
arr
        foldlTree b
acc' (Unbalanced Array (Tree a)
arr PrimArray Int
_) = (b -> Tree a -> b) -> b -> Array (Tree a) -> b
forall b a. (b -> a -> b) -> b -> Array 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 b a. (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' :: 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' = (Tree a -> b -> b) -> b -> Array (Tree a) -> b
forall a b. (a -> b -> b) -> b -> Array a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr' Tree a -> b -> b
foldrTree' b
acc' Array (Tree a)
arr
        foldrTree' (Unbalanced Array (Tree a)
arr PrimArray Int
_) b
acc' = (Tree a -> b -> b) -> b -> Array (Tree a) -> b
forall a b. (a -> b -> b) -> b -> Array 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 a b. (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' :: 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) = (b -> Tree a -> b) -> b -> Array (Tree a) -> b
forall b a. (b -> a -> b) -> b -> Array a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' b -> Tree a -> b
foldlTree' b
acc' Array (Tree a)
arr
        foldlTree' b
acc' (Unbalanced Array (Tree a)
arr PrimArray Int
_) = (b -> Tree a -> b) -> b -> Array (Tree a) -> b
forall b a. (b -> a -> b) -> b -> Array 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 b a. (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 :: 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 = (Int -> a -> m -> m) -> m -> Vector a -> m
forall a b. (Int -> a -> b -> b) -> b -> Vector a -> b
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 m -> m -> m
forall a. Semigroup a => a -> a -> a
<> m
acc) m
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' = Int
-> (Tree a -> Int)
-> (Int -> Tree a -> b -> b)
-> b
-> Array (Tree a)
-> b
forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
A.ifoldrStep Int
i (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) ((Int -> Int -> Tree a -> b -> b) -> Int -> Int -> Tree a -> b -> b
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' = Int
-> (Tree a -> Int)
-> (Int -> Tree a -> b -> b)
-> b
-> Array (Tree a)
-> b
forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
A.ifoldrStep Int
i (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) ((Int -> Int -> Tree a -> b -> b) -> Int -> Int -> Tree a -> b -> b
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' = Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
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 Int -> Int -> Int
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) = Int
-> (Tree a -> Int)
-> (Int -> b -> Tree a -> b)
-> b
-> Array (Tree a)
-> b
forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
A.ifoldlStep Int
i (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) ((Int -> Int -> b -> Tree a -> b) -> Int -> Int -> b -> Tree a -> b
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
_) = Int
-> (Tree a -> Int)
-> (Int -> b -> Tree a -> b)
-> b
-> Array (Tree a)
-> b
forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
A.ifoldlStep Int
i (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) ((Int -> Int -> b -> Tree a -> b) -> Int -> Int -> b -> Tree a -> b
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) = Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
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 Int -> Int -> Int
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' = Int
-> (Tree a -> Int)
-> (Int -> Tree a -> b -> b)
-> b
-> Array (Tree a)
-> b
forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
A.ifoldrStep' Int
i (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) ((Int -> Int -> Tree a -> b -> b) -> Int -> Int -> Tree a -> b -> b
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' = Int
-> (Tree a -> Int)
-> (Int -> Tree a -> b -> b)
-> b
-> Array (Tree a)
-> b
forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
A.ifoldrStep' Int
i (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) ((Int -> Int -> Tree a -> b -> b) -> Int -> Int -> Tree a -> b -> b
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' = Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
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) = Int
-> (Tree a -> Int)
-> (Int -> b -> Tree a -> b)
-> b
-> Array (Tree a)
-> b
forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
A.ifoldlStep' Int
i (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) ((Int -> Int -> b -> Tree a -> b) -> Int -> Int -> b -> Tree a -> b
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
_) = Int
-> (Tree a -> Int)
-> (Int -> b -> Tree a -> b)
-> b
-> Array (Tree a)
-> b
forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
A.ifoldlStep' Int
i (Int -> Tree a -> Int
forall a. Int -> Tree a -> Int
treeSize (Int -> Int
down Int
sh)) ((Int -> Int -> b -> Tree a -> b) -> Int -> Int -> b -> Tree a -> b
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) = Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
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 = (a -> b) -> Vector a -> Vector b
forall a b. (a -> b) -> Vector a -> Vector b
map
    a
x <$ :: forall a b. a -> Vector b -> Vector a
<$ Vector b
v = Int -> a -> Vector a
forall a. Int -> a -> Vector a
replicate (Vector b -> Int
forall a. Vector a -> Int
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 = Vector b
forall a. Vector a
Empty
    imap Int -> a -> b
f (Root Int
size Int
sh Tree a
tree) = Int -> Int -> Tree b -> Vector b
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (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) = Array (Tree b) -> Tree b
forall a. Array (Tree a) -> Tree a
Balanced (Int
-> (Tree a -> Int)
-> (Int -> Tree a -> Tree b)
-> Array (Tree a)
-> Array (Tree b)
forall a b.
Int -> (a -> Int) -> (Int -> a -> b) -> Array a -> Array b
A.imapStep' Int
i (Int -> Tree a -> Int
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) = Array (Tree b) -> PrimArray Int -> Tree b
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced (Int
-> (Tree a -> Int)
-> (Int -> Tree a -> Tree b)
-> Array (Tree a)
-> Array (Tree b)
forall a b.
Int -> (a -> Int) -> (Int -> a -> b) -> Array a -> Array b
A.imapStep' Int
i (Int -> Tree a -> Int
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) = Array b -> Tree b
forall a. Array a -> Tree a
Leaf (Int -> (a -> Int) -> (Int -> a -> b) -> Array a -> Array b
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 = Vector b -> f (Vector b)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector b
forall a. Vector a
Empty
    traverse a -> f b
f (Root Int
size Int
sh Tree a
tree) = Int -> Int -> Tree b -> Vector b
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (Tree b -> Vector b) -> f (Tree b) -> f (Vector b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree a -> f (Tree b)
traverseTree Tree a
tree
      where
        traverseTree :: Tree a -> f (Tree b)
traverseTree (Balanced Array (Tree a)
arr) = Array (Tree b) -> Tree b
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree b) -> Tree b) -> f (Array (Tree b)) -> f (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Tree a -> f (Tree b)) -> Array (Tree a) -> f (Array (Tree b))
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array a -> f (Array b)
A.traverse' Tree a -> f (Tree b)
traverseTree Array (Tree a)
arr
        traverseTree (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = (Array (Tree b) -> PrimArray Int -> Tree b)
-> PrimArray Int -> Array (Tree b) -> Tree b
forall a b c. (a -> b -> c) -> b -> a -> c
flip Array (Tree b) -> PrimArray Int -> Tree b
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced PrimArray Int
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 :: forall (f :: * -> *) a b.
Applicative f =>
(Int -> a -> f b) -> Vector a -> f (Vector b)
itraverse Int -> a -> f b
_ Vector a
Empty = Vector b -> f (Vector b)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector b
forall a. Vector a
Empty
    itraverse Int -> a -> f b
f (Root Int
size Int
sh Tree a
tree) = Int -> Int -> Tree b -> Vector b
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (Tree b -> Vector b) -> f (Tree b) -> f (Vector b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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) = 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
<$> Int
-> (Tree a -> Int)
-> (Int -> Tree a -> f (Tree b))
-> Array (Tree a)
-> f (Array (Tree b))
forall (f :: * -> *) a b.
Applicative f =>
Int -> (a -> Int) -> (Int -> a -> f b) -> Array a -> f (Array b)
A.itraverseStep' Int
i (Int -> Tree a -> Int
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) = (Array (Tree b) -> PrimArray Int -> Tree b)
-> PrimArray Int -> Array (Tree b) -> Tree b
forall a b c. (a -> b -> c) -> b -> a -> c
flip Array (Tree b) -> PrimArray Int -> Tree b
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced PrimArray Int
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
<$> Int
-> (Tree a -> Int)
-> (Int -> Tree a -> f (Tree b))
-> Array (Tree a)
-> f (Array (Tree b))
forall (f :: * -> *) a b.
Applicative f =>
Int -> (a -> Int) -> (Int -> a -> f b) -> Array a -> f (Array b)
A.itraverseStep' Int
i (Int -> Tree a -> Int
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) = 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
<$> Int -> (a -> Int) -> (Int -> a -> f b) -> Array a -> f (Array 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 = a -> Vector a
forall a. a -> Vector a
singleton
    Vector (a -> b)
fs <*> :: forall a b. Vector (a -> b) -> Vector a -> Vector b
<*> Vector a
xs = (Vector b -> (a -> b) -> Vector b)
-> Vector b -> Vector (a -> b) -> Vector b
forall b a. (b -> a -> b) -> b -> Vector a -> 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 :: forall a b c. (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 b a. (b -> a -> b) -> b -> Vector a -> b
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 *> :: forall a b. Vector a -> Vector b -> Vector b
*> Vector b
ys = Int -> Vector b -> Vector b
forall b. Integral b => b -> Vector b -> Vector b
forall a b. (Semigroup a, Integral b) => b -> a -> a
stimes (Vector a -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
xs) Vector b
ys
    Vector a
xs <* :: forall a b. Vector a -> Vector b -> Vector a
<* Vector b
ys = (Vector a -> a -> Vector a) -> Vector a -> Vector a -> Vector a
forall b a. (b -> a -> b) -> b -> Vector a -> b
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
>< Int -> a -> Vector a
forall a. Int -> a -> Vector a
replicate (Vector b -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector b
ys) a
x) Vector a
forall a. Vector a
empty Vector a
xs

instance Monad Vector where
    Vector a
xs >>= :: forall a b. Vector a -> (a -> Vector b) -> Vector b
>>= a -> Vector b
f = (Vector b -> a -> Vector b) -> Vector b -> Vector a -> Vector b
forall b a. (b -> a -> b) -> b -> Vector a -> 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
    >> :: forall a b. Vector a -> Vector b -> Vector b
(>>) = Vector a -> Vector b -> Vector b
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 = Vector a
forall a. Vector a
empty
    <|> :: forall a. Vector a -> Vector a -> Vector 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]
_ = Vector a
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 a
forall {a}. a
err of
        Vector a
Empty -> Vector a
forall a. Vector a
Empty
        Root Int
size Int
_ Tree a
_ -> [a] -> Vector a
forall a. [a] -> Vector a
fromList ([a] -> Vector a) -> [a] -> Vector a
forall a b. (a -> b) -> a -> b
$ (Int -> a) -> [Int] -> [a]
forall 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 Vector a -> Int -> a
forall a. HasCallStack => Vector a -> Int -> a
! Int
i in a
x) [Int
0..Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      where
        err :: a
err = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"mfix for Data.RRBVector.Vector applied to strict function"

instance MonadZip Vector where
    mzipWith :: forall a b c. (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 :: forall a b. 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 :: forall a b. 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 = [a] -> Vector a
[Item (Vector a)] -> Vector a
forall a. [a] -> Vector a
fromList
    toList :: Vector a -> [Item (Vector a)]
toList = Vector a -> [a]
Vector a -> [Item (Vector a)]
forall a. Vector a -> [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
[Char] -> Vector Char
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 :: forall a. (a -> ()) -> Vector a -> ()
liftRnf a -> ()
f = (() -> a -> ()) -> () -> Vector a -> ()
forall b a. (b -> a -> b) -> b -> Vector a -> b
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 = Vector a
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 = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root Int
1 Int
0 (Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ a -> Array a
forall a. a -> Array a
A.singleton a
x)

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

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

-- | \(O(\log n)\). @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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Vector a
forall a. Vector a
Empty
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
blockSize = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root Int
n Int
0 (Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ Int -> a -> Array a
forall a. Int -> a -> Array a
A.replicate Int
n a
x)
    | Bool
otherwise = Int -> Tree a -> Tree a -> Vector a
iterateNodes Int
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
$ Int -> a -> Array a
forall a. Int -> a -> Array a
A.replicate Int
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
$ Int -> a -> Array a
forall a. Int -> a -> Array a
A.replicate (Int
lastIdx Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
blockMask Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
x)
  where
    lastIdx :: Int
lastIdx = Int
n Int -> Int -> Int
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 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
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
$ Int -> Tree a -> Array (Tree a)
forall a. Int -> a -> Array a
A.replicate Int
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
$ Int -> Tree a -> Tree a -> Array (Tree a)
forall a. Int -> a -> a -> Array a
A.replicateSnoc (Int
subtreesM1 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
blockMask) Tree a
full Tree a
rest
        in if Int
subtreesM1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int -> Int -> Tree a -> Vector a
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 = Maybe a
forall a. Maybe a
Nothing
lookup Int
i (Root Int
size Int
sh Tree a
tree)
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
size = Maybe a
forall a. Maybe a
Nothing  -- index out of range
    | Bool
otherwise = a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> a -> Maybe a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tree a -> a
forall {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) (Array (Tree t) -> Int -> Tree t
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) (Array (Tree t) -> Int -> Tree t
forall a. Array a -> Int -> a
A.index Array (Tree t)
arr Int
idx)
    lookupTree Int
i Int
_ (Leaf Array t
arr) = Array t -> Int -> t
forall a. Array a -> Int -> a
A.index Array t
arr (Int
i Int -> Int -> Int
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 = a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"index out of range") (Maybe a -> a) -> (Vector a -> Maybe a) -> Vector a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Vector a -> Maybe a
forall a. Int -> Vector a -> Maybe a
lookup Int
i
{-# INLINE index #-}

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

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

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

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

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

-- | \(O(n)\). 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 = Vector b
forall a. Vector a
Empty
map' a -> b
f (Root Int
size Int
sh Tree a
tree) = Int -> Int -> Tree b -> Vector b
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh (Tree a -> Tree b
mapTree Tree a
tree)
  where
    mapTree :: Tree a -> Tree b
mapTree (Balanced Array (Tree a)
arr) = Array (Tree b) -> Tree b
forall a. Array (Tree a) -> Tree a
Balanced ((Tree a -> Tree b) -> Array (Tree a) -> Array (Tree b)
forall a b. (a -> b) -> Array a -> Array b
A.map' Tree a -> Tree b
mapTree Array (Tree a)
arr)
    mapTree (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = Array (Tree b) -> PrimArray Int -> Tree b
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced ((Tree a -> Tree b) -> Array (Tree a) -> Array (Tree b)
forall a b. (a -> b) -> Array a -> Array b
A.map' Tree a -> Tree b
mapTree Array (Tree a)
arr) PrimArray Int
sizes
    mapTree (Leaf Array a
arr) = Array b -> Tree b
forall a. Array a -> Tree a
Leaf ((a -> b) -> Array a -> Array b
forall a b. (a -> b) -> Array a -> Array b
A.map' a -> b
f Array a
arr)

-- | \(O(n)\). Reverse the vector.
--
-- >>> reverse (fromList [1, 2, 3])
-- fromList [3,2,1]
reverse :: Vector a -> Vector a
reverse :: forall a. Vector a -> Vector a
reverse Vector a
v
    | Vector a -> Int
forall a. Vector a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Vector a
v Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = Vector a
v
    | Bool
otherwise = [a] -> Vector a
forall a. [a] -> Vector a
fromList (([a] -> a -> [a]) -> [a] -> Vector a -> [a]
forall b a. (b -> a -> b) -> b -> Vector a -> b
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 (:)) [] 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 = [(a, b)] -> Vector (a, b)
forall a. [a] -> Vector a
fromList ([(a, b)] -> Vector (a, b)) -> [(a, b)] -> Vector (a, b)
forall a b. (a -> b) -> a -> b
$ [a] -> [b] -> [(a, b)]
forall a b. [a] -> [b] -> [(a, b)]
List.zip (Vector a -> [a]
forall a. Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (Vector b -> [b]
forall a. Vector a -> [a]
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 = [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 a. Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Vector a
v1) (Vector b -> [b]
forall a. Vector a -> [a]
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 :: forall a b. Vector (a, b) -> (Vector a, Vector b)
unzip = (((a, b) -> (a, b)) -> Vector (a, b) -> (Vector a, Vector b))
-> ((a, b) -> (a, b)) -> Vector (a, b) -> (Vector a, Vector b)
forall a. a -> a
Exts.inline ((a, b) -> (a, b)) -> Vector (a, b) -> (Vector a, Vector b)
forall a b c. (a -> (b, c)) -> Vector a -> (Vector b, Vector c)
unzipWith (a, b) -> (a, b)
forall a. a -> a
id

-- | \(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 = (Vector b
forall a. Vector a
Empty, Vector c
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) -> (Int -> Int -> Tree b -> Vector b
forall a. Int -> Int -> Tree a -> Vector a
Root Int
size Int
sh Tree b
left, Int -> Int -> Tree c -> Vector c
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 (Tree a -> (Tree b, Tree c))
-> Array (Tree a) -> (Array (Tree b), Array (Tree c))
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) -> (Array (Tree b) -> Tree b
forall a. Array (Tree a) -> Tree a
Balanced Array (Tree b)
left, Array (Tree c) -> Tree c
forall a. Array (Tree a) -> Tree a
Balanced Array (Tree c)
right)
    unzipTree (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) = case (Tree a -> (Tree b, Tree c))
-> Array (Tree a) -> (Array (Tree b), Array (Tree c))
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) -> (Array (Tree b) -> PrimArray Int -> Tree b
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced Array (Tree b)
left PrimArray Int
sizes, Array (Tree c) -> PrimArray Int -> Tree c
forall a. Array (Tree a) -> PrimArray Int -> Tree a
Unbalanced Array (Tree c)
right PrimArray Int
sizes)
    unzipTree (Leaf Array a
arr) = case (a -> (b, c)) -> Array a -> (Array b, Array c)
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) -> (Array b -> Tree b
forall a. Array a -> Tree a
Leaf Array b
left, Array c -> Tree c
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 = Maybe (a, Vector a)
forall a. Maybe a
Nothing
viewl v :: Vector a
v@(Root Int
_ Int
_ Tree a
tree) = let !tail :: Vector a
tail = Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
drop Int
1 Vector a
v in (a, Vector a) -> Maybe (a, Vector a)
forall a. a -> Maybe a
Just (Tree a -> a
forall {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 (Array (Tree t) -> Tree t
forall a. Array a -> a
A.head Array (Tree t)
arr)
    headTree (Unbalanced Array (Tree t)
arr PrimArray Int
_) = Tree t -> t
headTree (Array (Tree t) -> Tree t
forall a. Array a -> a
A.head Array (Tree t)
arr)
    headTree (Leaf Array t
arr) = Array t -> t
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 = Maybe (Vector a, a)
forall a. Maybe a
Nothing
viewr v :: Vector a
v@(Root Int
size Int
_ Tree a
tree) = let !init :: Vector a
init = Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
take (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Vector a
v in (Vector a, a) -> Maybe (Vector a, a)
forall a. a -> Maybe a
Just (Vector a
init, Tree a -> a
forall {t}. Tree t -> t
lastTree Tree a
tree)
  where
    lastTree :: Tree t -> t
lastTree (Balanced Array (Tree t)
arr) = Tree t -> t
lastTree (Array (Tree t) -> Tree t
forall a. Array a -> a
A.last Array (Tree t)
arr)
    lastTree (Unbalanced Array (Tree t)
arr PrimArray Int
_) = Tree t -> t
lastTree (Array (Tree t) -> Tree t
forall a. Array a -> a
A.last Array (Tree t)
arr)
    lastTree (Leaf Array t
arr) = Array t -> t
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 = Vector a
forall a. Vector a
Empty
take Int
n v :: Vector a
v@(Root Int
size Int
sh Tree a
tree)
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = Vector a
forall a. Vector a
empty
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
size = Vector a
v
    | Bool
otherwise = Vector a -> Vector a
forall a. Vector a -> Vector a
normalize (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root Int
n Int
sh (Int -> Int -> Tree a -> Tree a
forall a. Int -> Int -> Tree a -> Tree a
takeTree (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
sh Tree a
tree)

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

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

-- | \(O(\log n)\). Insert an element at the given index, shifting the rest of the vector over.
-- If the index is negative, add the element to the left end of the vector.
-- If the index is bigger than or equal to the length of the vector, add the element to the right end of the vector.
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) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
splitAt Int
i Vector a
v in (Vector a
left Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
|> a
x) Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
>< Vector a
right

-- | \(O(\log n)\). Delete the element at the given index.
-- If the index is out of range, return the original vector.
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) = Int -> Vector a -> (Vector a, Vector a)
forall a. Int -> Vector a -> (Vector a, Vector a)
splitAt (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Vector a
v in Int -> Vector a -> Vector a
forall a. Int -> Vector a -> Vector a
take Int
i Vector a
left Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
>< Vector a
right

-- | \(O(n)\). Find the first index from the left that satisfies the predicate.
--
-- @since 0.2.1.0
findIndexL :: (a -> Bool) -> Vector a -> Maybe Int
findIndexL :: forall a. (a -> Bool) -> Vector a -> Maybe Int
findIndexL a -> Bool
f = (Int -> a -> Maybe Int -> Maybe Int)
-> Maybe Int -> Vector a -> Maybe Int
forall a b. (Int -> a -> b -> b) -> b -> Vector a -> b
forall i (f :: * -> *) a b.
FoldableWithIndex i f =>
(i -> a -> b -> b) -> b -> f a -> b
ifoldr (\Int
i a
x Maybe Int
acc -> if a -> Bool
f a
x then Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i else Maybe Int
acc) Maybe Int
forall a. Maybe a
Nothing
{-# INLINE findIndexL #-}

-- | \(O(n)\). Find the first index from the right that satisfies the predicate.
--
-- @since 0.2.1.0
findIndexR :: (a -> Bool) -> Vector a -> Maybe Int
findIndexR :: forall a. (a -> Bool) -> Vector a -> Maybe Int
findIndexR a -> Bool
f = (Int -> Maybe Int -> a -> Maybe Int)
-> Maybe Int -> Vector a -> Maybe Int
forall b a. (Int -> b -> a -> b) -> b -> Vector a -> b
forall i (f :: * -> *) b a.
FoldableWithIndex i f =>
(i -> b -> a -> b) -> b -> f a -> b
ifoldl (\Int
i Maybe Int
acc a
x -> if a -> Bool
f a
x then Int -> Maybe Int
forall a. a -> Maybe a
Just Int
i else Maybe Int
acc) Maybe Int
forall a. Maybe a
Nothing
{-# INLINE findIndexR #-}

-- | \(O(n)\). Find the indices that satisfy the predicate, starting from the left.
--
-- @since 0.2.1.0
findIndicesL :: (a -> Bool) -> Vector a -> [Int]
findIndicesL :: forall a. (a -> Bool) -> Vector a -> [Int]
findIndicesL a -> Bool
f = (Int -> a -> [Int] -> [Int]) -> [Int] -> Vector a -> [Int]
forall a b. (Int -> a -> b -> b) -> b -> Vector a -> b
forall i (f :: * -> *) a b.
FoldableWithIndex i f =>
(i -> a -> b -> b) -> b -> f a -> b
ifoldr (\Int
i a
x [Int]
acc -> if a -> Bool
f a
x then Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
acc else [Int]
acc) []
{-# INLINE findIndicesL #-}

-- | \(O(n)\). Find the indices that satisfy the predicate, starting from the right.
--
-- @since 0.2.1.0
findIndicesR :: (a -> Bool) -> Vector a -> [Int]
findIndicesR :: forall a. (a -> Bool) -> Vector a -> [Int]
findIndicesR a -> Bool
f = (Int -> [Int] -> a -> [Int]) -> [Int] -> Vector a -> [Int]
forall b a. (Int -> b -> a -> b) -> b -> Vector a -> b
forall i (f :: * -> *) b a.
FoldableWithIndex i f =>
(i -> b -> a -> b) -> b -> f a -> b
ifoldl (\Int
i [Int]
acc a
x -> if a -> Bool
f a
x then Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
acc else [Int]
acc) []
{-# INLINE findIndicesR #-}

-- 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 upMaxShift :: Int
upMaxShift = Int -> Int
up (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
sh1 Int
sh2)
        newArr :: Array (Tree a)
newArr = Tree a -> Int -> Tree a -> Int -> Array (Tree a)
forall {a}. Tree a -> Int -> Tree a -> Int -> Array (Tree a)
mergeTrees Tree a
tree1 Int
sh1 Tree a
tree2 Int
sh2
    in Vector a -> Vector a
forall a. Vector a -> Vector a
normalize (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
size2) Int
upMaxShift (Int -> Array (Tree a) -> Tree a
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
_
        | Array a -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blockSize = Tree a -> Tree a -> Array (Tree a)
forall a. a -> a -> Array a
A.from2 Tree a
tree1 Tree a
tree2
        | Array a -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Array a -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
blockSize = 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. Array a -> Array a -> Array a
A.++ Array a
arr2)
        | Bool
otherwise =
            let (Array a
left, Array a
right) = Array a -> Int -> (Array a, Array a)
forall a. Array a -> Int -> (Array a, Array a)
A.splitAt (Array a
arr1 Array a -> Array a -> Array a
forall a. Array a -> Array a -> Array a
A.++ Array a
arr2) Int
blockSize -- 'A.splitAt' doesn't copy anything
                !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 Int
sh1 Tree a
tree2 Int
sh2 = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
sh1 Int
sh2 of
        Ordering
LT ->
            let !right :: Array (Tree a)
right = Tree a -> Array (Tree a)
forall a. Tree a -> Array (Tree a)
treeToArray Tree a
tree2
                (Tree a
rightHead, Array (Tree a)
rightTail) = Array (Tree a) -> (Tree a, Array (Tree a))
forall {a}. Array a -> (a, Array a)
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 Int
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
forall a.
Int
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
mergeRebalance Int
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 -> Int -> Tree a -> Int -> Array (Tree a)
mergeTrees Tree a
leftLast (Int -> Int
down Int
sh1) Tree a
tree2 Int
sh2
            in Int
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
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)
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 -> 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 Int
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
-> Array (Tree a)
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 = (Array a -> a
forall a. Array a -> a
A.head Array a
arr, Array a -> Int -> Array a
forall a. Array a -> Int -> Array a
A.drop Array a
arr Int
1)

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

    -- the type 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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
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 (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes (Int -> Int
down Int
sh))
      where
        mergeRebalance' :: (Tree a -> A.Array t) -> (A.Array t -> Tree a) -> 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 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 <- Int -> ST s (Buffer s (Tree a))
forall s a. Int -> ST s (Buffer s a)
Buffer.new Int
blockSize
            Buffer s (Tree a)
newSubtree <- Int -> ST s (Buffer s (Tree a))
forall s a. Int -> ST s (Buffer s a)
Buffer.new Int
blockSize
            Buffer s t
newNode <- Int -> ST s (Buffer s t)
forall s a. Int -> ST s (Buffer s a)
Buffer.new Int
blockSize
            [Tree a] -> (Tree a -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Array (Tree a) -> [Tree a]
forall a. Array a -> [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 a. Array a -> [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 a. Array a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Array (Tree a)
right) ((Tree a -> ST s ()) -> ST s ()) -> (Tree a -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Tree a
subtree ->
                Array t -> (t -> ST s ()) -> ST s ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
t a -> (a -> f b) -> f ()
for_ (Tree a -> Array t
extract Tree a
subtree) ((t -> ST s ()) -> ST s ()) -> (t -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \t
x -> do
                    Int
lenNode <- Buffer s t -> ST s Int
forall s a. Buffer s a -> ST s Int
Buffer.size Buffer s t
newNode
                    Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lenNode Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blockSize) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
                        (Array t -> Tree a) -> Buffer s t -> Buffer s (Tree a) -> ST s ()
forall {a} {a} {s}.
(Array a -> a) -> Buffer s a -> Buffer s a -> ST s ()
pushTo Array t -> Tree a
construct Buffer s t
newNode Buffer s (Tree a)
newSubtree
                        Int
lenSubtree <- Buffer s (Tree a) -> ST s Int
forall s a. Buffer s a -> ST s Int
Buffer.size Buffer s (Tree a)
newSubtree
                        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
lenSubtree Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
blockSize) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ (Array (Tree a) -> Tree a)
-> Buffer s (Tree a) -> Buffer s (Tree a) -> ST s ()
forall {a} {a} {s}.
(Array a -> a) -> Buffer s a -> Buffer s a -> ST s ()
pushTo (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh) Buffer s (Tree a)
newSubtree Buffer s (Tree a)
newRoot
                    Buffer s t -> t -> ST s ()
forall s a. Buffer s a -> a -> ST s ()
Buffer.push Buffer s t
newNode t
x
            (Array t -> Tree a) -> Buffer s t -> Buffer s (Tree a) -> ST s ()
forall {a} {a} {s}.
(Array a -> a) -> Buffer s a -> Buffer s a -> ST s ()
pushTo Array t -> Tree a
construct Buffer s t
newNode Buffer s (Tree a)
newSubtree
            (Array (Tree a) -> Tree a)
-> Buffer s (Tree a) -> Buffer s (Tree a) -> ST s ()
forall {a} {a} {s}.
(Array a -> a) -> Buffer s a -> Buffer s a -> ST s ()
pushTo (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh) Buffer s (Tree a)
newSubtree Buffer s (Tree a)
newRoot
            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 <| :: forall a. a -> Vector a -> Vector a
<| Vector a
Empty = a -> Vector a
forall a. a -> Vector a
singleton a
x
a
x <| Root Int
size Int
sh Tree a
tree
    | Int
insertShift Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
sh = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
insertShift (Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
insertShift (let !new :: Tree a
new = a -> Int -> Tree a
forall a. a -> Int -> Tree a
newBranch a
x Int
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 = Int -> Int -> Tree a -> Vector a
forall a. Int -> Int -> Tree a -> Vector a
Root (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
sh (Int -> Tree a -> Tree a
consTree Int
sh Tree a
tree)
  where
    consTree :: Int -> Tree a -> Tree a
consTree Int
sh (Balanced Array (Tree a)
arr)
        | Int
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
insertShift = Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.cons Array (Tree a)
arr (Tree a -> Array (Tree a)) -> Tree a -> Array (Tree a)
forall a b. (a -> b) -> a -> b
$! a -> Int -> Tree a
forall a. a -> Int -> Tree a
newBranch a
x (Int -> Int
down Int
sh))
        | Bool
otherwise = Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Int
0 (Int -> Tree a -> Tree a
consTree (Int -> Int
down Int
sh)))
    consTree Int
sh (Unbalanced Array (Tree a)
arr PrimArray Int
_)
        | Int
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
insertShift = Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Tree a -> Array (Tree a)
forall a. Array a -> a -> Array a
A.cons Array (Tree a)
arr (Tree a -> Array (Tree a)) -> Tree a -> Array (Tree a)
forall a b. (a -> b) -> a -> b
$! a -> Int -> Tree a
forall a. a -> Int -> Tree a
newBranch a
x (Int -> Int
down Int
sh))
        | Bool
otherwise = Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
arr Int
0 (Int -> Tree a -> Tree a
consTree (Int -> Int
down Int
sh)))
    consTree Int
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Tree a) -> Array a -> Tree a
forall a b. (a -> b) -> a -> b
$ Array a -> a -> Array a
forall a. Array a -> a -> Array a
A.cons Array a
arr a
x

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

    -- compute the shift at which the new branch needs to be inserted (0 means there is space in the leaf)
    -- the 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)
_) =
        -- @sz - 1@ is the index of the last element
        let hiShift :: Int
hiShift = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max ((Int -> Int
log2 (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
blockShift) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
blockShift) Int
0 -- the shift of the root when normalizing
            hi :: Int
hi = (Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
hiShift -- the length of the root node when normalizing minus 1
            newShift :: Int
newShift = if Int
hi Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
blockMask then Int
hiShift else Int
hiShift Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
blockShift
        in if Int
newShift Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
sh then Int
min else Int
newShift
    computeShift Int
_ Int
sh Int
min (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) =
        let sz' :: Int
sz' = PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes Int
0 -- the size of the first subtree
            newMin :: Int
newMin = if Array (Tree a) -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
sh else Int
min
        in Int -> Int -> Int -> Tree a -> Int
computeShift Int
sz' (Int -> Int
down Int
sh) Int
newMin (Array (Tree a) -> Tree a
forall a. Array a -> a
A.head Array (Tree a)
arr)
    computeShift Int
_ Int
_ Int
min (Leaf Array a
arr) = if Array a -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
0 else Int
min

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

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

    -- compute the shift at which the new branch needs to be inserted (0 means there is space in the leaf)
    -- the 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
forall b. FiniteBits b => b -> Int
countTrailingZeros Int
sz Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
blockShift) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
blockShift
        in if Int
newShift Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
sh then Int
min else Int
newShift
    computeShift Int
_ Int
sh Int
min (Unbalanced Array (Tree a)
arr PrimArray Int
sizes) =
        let lastIdx :: Int
lastIdx = Array (Tree a) -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
            sz' :: Int
sz' = PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes Int
lastIdx Int -> Int -> Int
forall a. Num a => a -> a -> a
- PrimArray Int -> Int -> Int
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Int
sizes (Int
lastIdx Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) -- the size of the last subtree
            newMin :: Int
newMin = if Array (Tree a) -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array (Tree a)
arr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
sh else Int
min
        in Int -> Int -> Int -> Tree a -> Int
computeShift Int
sz' (Int -> Int
down Int
sh) Int
newMin (Array (Tree a) -> Tree a
forall a. Array a -> a
A.last Array (Tree a)
arr)
    computeShift Int
_ Int
_ Int
min (Leaf Array a
arr) = if Array a -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
blockSize then Int
0 else Int
min

-- create a new tree with shift @sh@
newBranch :: a -> 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 = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (a -> Array a
forall a. a -> Array a
A.singleton a
x)
    go Int
sh = Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (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
$! 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 = Array (Tree a) -> Int -> Array (Tree a)
forall a. Array a -> Int -> Array a
A.take Array (Tree a)
arr (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    in Array (Tree a) -> Tree a
forall a. Array (Tree a) -> Tree a
Balanced (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Int
idx (Int -> Int -> Tree a -> Tree a
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 = Array (Tree a) -> Int -> Array (Tree a)
forall a. Array a -> Int -> Array a
A.take Array (Tree a)
arr (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    in Int -> Array (Tree a) -> Tree a
forall a. Int -> Array (Tree a) -> Tree a
computeSizes Int
sh (Array (Tree a) -> Int -> (Tree a -> Tree a) -> Array (Tree a)
forall a. Array a -> Int -> (a -> a) -> Array a
A.adjust' Array (Tree a)
newArr Int
idx (Int -> Int -> Tree a -> Tree a
forall a. Int -> Int -> Tree a -> Tree a
takeTree Int
subIdx (Int -> Int
down Int
sh)))
takeTree Int
i Int
_ (Leaf Array a
arr) = Array a -> Tree a
forall a. Array a -> Tree a
Leaf (Array a -> Int -> Array a
forall a. Array a -> Int -> Array a
A.take Array a
arr ((Int
i Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
blockMask) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))

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

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