{-# LANGUAGE CPP #-}

{-# LANGUAGE MagicHash #-}

{-# LANGUAGE PartialTypeSignatures #-}

{-# LANGUAGE RoleAnnotations #-}

{-# LANGUAGE UnboxedTuples #-}

module Data.Vector.Persistent.Internal where

import Control.Applicative (Alternative, liftA2)

import qualified Control.Applicative

import Control.DeepSeq (NFData (rnf), NFData1, rnf1)

import qualified Control.DeepSeq

import Control.Monad (MonadPlus)

import qualified Control.Monad.Fail as Fail

import Control.Monad.Primitive (PrimMonad)

import Control.Monad.ST (runST)

import Data.Bits (Bits, unsafeShiftL, unsafeShiftR, (.&.))

import qualified Data.Foldable as Foldable

import Data.Functor.Classes

( Show1,

liftShowsPrec,

showsPrec1,

showsUnaryWith,

)

import Data.Functor.Identity (Identity (..))

import Data.Primitive.SmallArray

import Data.Stream.Monadic (Stream (Stream))

import qualified Data.Stream.Monadic as Stream

import qualified Data.Traversable as Traversable

import Data.Vector.Persistent.Internal.Array

import qualified Data.Vector.Persistent.Internal.Buffer as Buffer

import Data.Vector.Persistent.Internal.CoercibleUtils

import GHC.Exts (IsList)

import qualified GHC.Exts as Exts

import GHC.Stack (HasCallStack)

import Prelude hiding (init, length, lookup, map, null, tail)

#ifdef INSPECTION

{-# LANGUAGE TemplateHaskell #-}

import Test.Inspection

type role Vector representational

-- | A vector.

data Vector a

= -- |

-- Invariants: The only time tail can be empty is when init is empty.

-- Otherwise tailOffset will give the wrong value.

RootNode

{ size :: !Int,

-- | 1 << 'shift' is the maximum that each child can contain

shift :: !Int,

init :: !(Array (Node a)),

tail :: !(Array a)

}

instance Show1 Vector where

liftShowsPrec sp sl p v = showsUnaryWith (liftShowsPrec sp sl) "fromList" p (toList v)

instance Show a => Show (Vector a) where

showsPrec = showsPrec1

instance Eq a => Eq (Vector a) where

(==) = persistentVectorEq

{-# INLINE (==) #-}

instance Ord a => Ord (Vector a) where

compare = persistentVectorCompare

{-# INLINE compare #-}

instance Functor Vector where

fmap = Data.Vector.Persistent.Internal.map

{-# INLINE fmap #-}

instance Foldable Vector where

foldr = Data.Vector.Persistent.Internal.foldr

{-# INLINE foldr #-}

foldl = Data.Vector.Persistent.Internal.foldl

{-# INLINE foldl #-}

foldl' = Data.Vector.Persistent.Internal.foldl'

{-# INLINE foldl' #-}

foldr' = Data.Vector.Persistent.Internal.foldr'

{-# INLINE foldr' #-}

length = Data.Vector.Persistent.Internal.length

{-# INLINE length #-}

null = Data.Vector.Persistent.Internal.null

{-# INLINE null #-}

instance Traversable Vector where

traverse = Data.Vector.Persistent.Internal.traverse

{-# INLINE traverse #-}

instance Semigroup (Vector a) where

(<>) = (><)

{-# INLINE (<>) #-}

instance Monoid (Vector a) where

mempty = empty

{-# INLINE mempty #-}

instance NFData a => NFData (Vector a) where

rnf = rnf1

{-# INLINE rnf #-}

instance Applicative Vector where

pure = singleton

{-# INLINE pure #-}

fs <*> xs = foldl' (\acc f -> acc >< map f xs) empty fs

instance Monad Vector where

xs >>= f = foldl' (\acc x -> acc >< f x) empty xs

{-# INLINE (>>=) #-}

instance Fail.MonadFail Vector where

fail _ = empty

{-# INLINE fail #-}

instance Alternative Vector where

empty = empty

{-# INLINE empty #-}

(<|>) = (><)

{-# INLINE (<|>) #-}

instance MonadPlus Vector

instance NFData1 Vector where

liftRnf f = foldl' (\_ x -> f x) ()

{-# INLINE liftRnf #-}

data Node a

= InternalNode {getInternalNode :: !(Array (Node a))}

| DataNode {getDataNode :: !(Array a)}

deriving (Show)

instance Eq a => Eq (Node a) where

(==) = nodeEq

{-# INLINE (==) #-}

instance Ord a => Ord (Node a) where

compare = nodeCompare

{-# INLINE compare #-}

instance IsList (Vector a) where

type Item (Vector a) = a

fromList = Data.Vector.Persistent.Internal.fromList

{-# INLINE fromList #-}

toList = Data.Vector.Persistent.Internal.toList

{-# INLINE toList #-}

-- | \(O(n)\) Lazy right fold.

foldr :: (a -> b -> b) -> b -> Vector a -> b

foldr f z = runIdentity #. Stream.foldr f z . streamL

{-# INLINE foldr #-}

-- | \(O(n)\) Strict right fold.

foldr' :: (a -> b -> b) -> b -> Vector a -> b

foldr' f z = runIdentity #. Stream.foldl' (flip f) z . streamR

{-# INLINE foldr' #-}

-- | \(O(n)\) Lazy left fold.

foldl :: (b -> a -> b) -> b -> Vector a -> b

foldl f z = runIdentity #. Stream.foldr (flip f) z . streamR

{-# INLINE foldl #-}

-- | \(O(n)\) Strict left fold.

foldl' :: (b -> a -> b) -> b -> Vector a -> b

foldl' f z = runIdentity #. Stream.foldl' f z . streamL

{-# INLINE foldl' #-}

-- | \(O(n)\) Indexed lazy right fold.

ifoldr :: (Int -> a -> b -> b) -> b -> Vector a -> b

ifoldr f z = runIdentity #. Stream.foldr (uncurry f) z . istreamL

{-# INLINE ifoldr #-}

-- | \(O(n)\) Indexed lazy left fold.

ifoldl :: (Int -> b -> a -> b) -> b -> Vector a -> b

ifoldl f z = runIdentity #. Stream.foldr (\(i, x) y -> f i y x) z . istreamR

{-# INLINE ifoldl #-}

-- | \(O(n)\) Indexed strict right fold.

ifoldr' :: (Int -> a -> b -> b) -> b -> Vector a -> b

ifoldr' f z = runIdentity #. Stream.foldl' (\y (i, x) -> f i x y) z . istreamR

{-# INLINE ifoldr' #-}

-- | \(O(n)\) Indexed strict left fold.

ifoldl' :: (Int -> b -> a -> b) -> b -> Vector a -> b

ifoldl' f z = runIdentity #. Stream.foldl' (\y (i, x) -> f i y x) z . istreamL

{-# INLINE ifoldl' #-}

persistentVectorEq :: Eq a => Vector a -> Vector a -> Bool

persistentVectorEq

RootNode {size, shift, init, tail}

RootNode {size = size', shift = shift', init = init', tail = tail'} =

size == size' && (size == 0 || (shift == shift' && tail == tail' && init == init'))

{-# INLINEABLE persistentVectorEq #-}

nodeEq :: Eq a => Node a -> Node a -> Bool

nodeEq (InternalNode ns) (InternalNode ns') = ns == ns'

nodeEq (DataNode as) (DataNode as') = as == as'

nodeEq _ _ = False

{-# INLINEABLE nodeEq #-}

persistentVectorCompare :: Ord a => Vector a -> Vector a -> Ordering

persistentVectorCompare

RootNode {size, init, tail}

RootNode {size = size', init = init', tail = tail'} =

compare size size'

<> if size == 0

then EQ

else compare init init' <> compare tail tail'

{-# INLINEABLE persistentVectorCompare #-}

nodeCompare :: Ord a => Node a -> Node a -> Ordering

nodeCompare (DataNode as) (DataNode as') = compare as as'

nodeCompare (InternalNode ns) (InternalNode ns') = compare ns ns'

nodeCompare (DataNode _) (InternalNode _) = LT

nodeCompare (InternalNode _) (DataNode _) = GT

{-# INLINEABLE nodeCompare #-}

-- | \(O(1)\). A vector with a single element.

singleton :: a -> Vector a

singleton a = RootNode {size = 1, shift = keyBits, tail = singletonSmallArray a, init = emptySmallArray}

{-# INLINE singleton #-}

-- | \(O(1)\). The empty vector.

empty :: Vector a

empty = RootNode {size = 0, shift = keyBits, init = emptySmallArray, tail = emptySmallArray}

{-# NOINLINE empty #-}

-- | \(O(1)\) Return 'True' if the vector is empty, 'False' otherwise.

null :: Vector a -> Bool

null xs = length xs == 0

{-# INLINE null #-}

-- | \(O(\log n)\). An alias for 'snoc'

(|>) :: Vector a -> a -> Vector a

(|>) = snoc

{-# INLINE (|>) #-}

-- | \(O(\log n)\). A bidirectional pattern synonym viewing the rear of a non-empty

pattern (:|>) :: Vector a -> a -> Vector a

pattern vec :|> x <-

(unsnoc -> Just (vec, x))

where

vec :|> x = vec |> x

infixl 5 :|>

-- | \(O(1)\). A bidirectional pattern synonym matching an empty sequence.

pattern Empty :: Vector a

pattern Empty <-

(null -> True)

where

Empty = empty

{-# COMPLETE (:|>), Empty #-}

-- | \(O(\log n)\) Add an element to the end of the vector.

snoc :: Vector a -> a -> Vector a

snoc vec@RootNode {size, tail} a

-- Room in tail, and vector non-empty

| (size .&. keyMask) /= 0 =

vec

{ tail = snocSmallArray tail a,

size = size + 1

}

| otherwise = snocArr vec 1 $ singletonSmallArray a

{-# INLINEABLE snoc #-}

unsafeSnoc :: Vector a -> a -> Vector a

unsafeSnoc vec@RootNode {size, tail} a

-- Room in tail, and vector non-empty

| (size .&. keyMask) /= 0 =

vec

{ tail =

-- update the array in place

runST $ do

marr <- unsafeThawSmallArray tail

writeSmallArray marr (size .&. keyMask) a

unsafeFreezeSmallArray marr,

size = size + 1

}

| otherwise = snocArr vec 1 $ singletonSmallArray a

{-# INLINEABLE unsafeSnoc #-}

snocArr ::

-- | The Vector to perform the operation on

Vector a ->

-- | The the added size. We can't find this from the array because the array might have bogus size due to undefined elements

Int ->

-- | The array to add as the new tail

Array a ->

Vector a

snocArr vec@RootNode {size, shift, tail} addedSize arr

| null vec =

RootNode

{ size = addedSize,

shift = keyBits,

tail = arr,

init = emptySmallArray

}

| size !>>. keyBits > 1 !<<. shift =

RootNode

{ size = size + addedSize,

shift = shift + keyBits,

init =

let !path = newPath shift tail

!internal = InternalNode $ init vec

in twoSmallArray internal path,

tail = arr

}

| otherwise =

RootNode

{ size = size + addedSize,

shift,

init = snocTail size tail shift $ init vec,

tail = arr

}

{-# INLINE snocArr #-}

snocTail :: Int -> Array a -> Int -> Array (Node a) -> Array (Node a)

snocTail size tail = go

where

go !level !parent = updateResizeSmallArray parent subIx toInsert

where

toInsert

| level == keyBits = DataNode tail

| subIx < sizeofSmallArray parent =

let vec' = indexSmallArray parent subIx

in InternalNode $ go (level - keyBits) (getInternalNode vec')

| otherwise = newPath (level - keyBits) tail

subIx = ((size - 1) !>>. level) .&. keyMask

{-# INLINE snocTail #-}

newPath :: Int -> Array a -> Node a

newPath 0 tail = DataNode tail

newPath level tail = InternalNode $ singletonSmallArray $! newPath (level - keyBits) tail

unsafeIndex :: Vector a -> Int -> a

unsafeIndex vec ix | (# a #) <- Exts.inline unsafeIndex# vec ix = a

{-# NOINLINE unsafeIndex #-}

unsafeIndex# :: Vector a -> Int -> (# a #)

unsafeIndex# vec ix

| ix >= tailOffset vec = indexSmallArray## (tail vec) (ix .&. keyMask)

-- no need to use keyMask here as we are at the top

| otherwise = go ix (shift vec - keyBits) (indexSmallArray (init vec) (ix !>>. shift vec))

where

go ix 0 !node = indexSmallArray## (getDataNode node) (ix .&. keyMask)

go ix level !node = go ix (level - keyBits) node'

where

!(# node' #) = indexSmallArray## (getInternalNode node) ((ix !>>. level) .&. keyMask)

{-# NOINLINE unsafeIndex# #-}

lookup# :: Int -> Vector a -> (# (# #) | a #)

lookup# ix vec

| (fromIntegral ix :: Word) >= fromIntegral (length vec) = (# (# #) | #)

| otherwise = case Exts.inline unsafeIndex# vec ix of (# x #) -> (# | x #)

{-# NOINLINE lookup# #-}

-- | \(O(\log n)\). The element at the index or 'Nothing' if the index is out of range.

lookup :: Int -> Vector a -> Maybe a

lookup ix vec

| (fromIntegral ix :: Word) >= fromIntegral (length vec) = Nothing

| otherwise = case unsafeIndex# vec ix of (# x #) -> Just x

{-# 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 ix vec

| ix < 0 = moduleError "index" $ "negative index: " ++ show ix

| ix >= length vec = moduleError "index" $ "index too large: " ++ show ix

| otherwise = Exts.inline unsafeIndex vec ix

{-# INLINEABLE index #-}

-- | \(O(\log n)\). A flipped version of 'index'.

(!) :: HasCallStack => Vector a -> Int -> a

(!) = flip index

{-# INLINE (!) #-}

-- | \(O(\log n)\). A flipped version of 'lookup'.

(!?) :: Vector a -> Int -> Maybe a

(!?) = flip lookup

{-# INLINE (!?) #-}

-- | \(O(\log n)\). Adjust the element at the index by applying the function to it.

adjust :: (a -> a) -> Int -> Vector a -> Vector a

adjust f = adjust# $ \x -> (# f x #)

{-# INLINE adjust #-}

adjust# :: (a -> (# a #)) -> Int -> Vector a -> Vector a

adjust# f ix vec@RootNode {size, shift, tail}

-- Invalid index. This funny business uses a single test to determine whether

-- ix is too small (negative) or too large (at least sz).

| (fromIntegral ix :: Word) >= fromIntegral size = vec

| ix >= tailOffset vec = vec {tail = modifySmallArray# tail (ix .&. keyMask) f}

| otherwise = vec {init = go ix shift (init vec)}

where

go ix level vec

| level == keyBits =

let !node = DataNode $ modifySmallArray# (getDataNode vec') (ix .&. keyMask) f

in updateSmallArray vec ix' node

| otherwise =

let !node = go ix (level - keyBits) (getInternalNode vec')

in updateSmallArray vec ix' $! InternalNode node

where

ix' = (ix !>>. level) .&. keyMask

vec' = indexSmallArray vec ix'

{-# INLINE adjust# #-}

-- | \(O(\log n)\). Same as 'adjust' but can have effects through 'Applicative'

adjustF :: Applicative f => (a -> f a) -> Int -> Vector a -> f (Vector a)

adjustF f ix vec@RootNode {size, shift, tail}

-- Invalid index. This funny business uses a single test to determine whether

-- ix is too small (negative) or too large (at least sz).

| (fromIntegral ix :: Word) >= fromIntegral size = pure vec

| ix >= tailOffset vec = (\tail -> vec {tail}) <$> modifySmallArrayF tail (ix .&. keyMask) f

| otherwise = (\init -> vec {init}) <$> go ix shift (init vec)

where

go ix level vec

| level == keyBits =

(\node' -> updateSmallArray vec ix' $! DataNode node')

<$> modifySmallArrayF (getDataNode vec') (ix .&. keyMask) f

| otherwise =

(\node -> updateSmallArray vec ix' $! InternalNode node)

<$> go ix (level - keyBits) (getInternalNode vec')

where

ix' = (ix !>>. level) .&. keyMask

vec' = indexSmallArray vec ix'

{-# INLINE adjustF #-}

-- | \(O(\log n)\). Replace the element at the specified position.

update :: Int -> a -> Vector a -> Vector a

update ix x vec@RootNode {size, shift, tail}

-- Invalid index. This funny business uses a single test to determine whether

-- ix is too small (negative) or too large (at least sz).

| (fromIntegral ix :: Word) >= fromIntegral size = vec

| ix >= tailOffset vec = vec {tail = updateSmallArray tail (ix .&. keyMask) x}

| otherwise = vec {init = go ix x shift (init vec)}

where

go ix x level vec

| level == keyBits =

let !node = DataNode $ updateSmallArray (getDataNode vec') (ix .&. keyMask) x

in updateSmallArray vec ix' node

| otherwise =

let !node = go ix x (level - keyBits) (getInternalNode vec')

in updateSmallArray vec ix' $! InternalNode node

where

ix' = (ix !>>. level) .&. keyMask

vec' = indexSmallArray vec ix'

{-# INLINEABLE update #-}

-- | \(O(\log n)\). Decompose a list into its head and tail.

unsnoc :: Vector a -> Maybe (Vector a, a)

unsnoc vec@RootNode {size, tail, init, shift}

| 0 <- size = Nothing

-- we need to have this case because we can't run unsnocTail, there is nothing left in the tail

| 1 <- size, (# x #) <- indexSmallArray## tail 0 = Just (empty, x)

| otherwise = case unsnocSmallArray tail of

Nothing ->

let !(# init', tail' #) = unsnocTail# size shift init

!(# x #) = lastSmallArray# tail'

in Just (vec {size = size - 1, init = init', tail = popSmallArray tail'}, x)

Just (tail', a) -> Just (vec {size = size - 1, tail = tail'}, a)

{-# INLINEABLE unsnoc #-}

unsnocTail# :: Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)

unsnocTail# = go

where

go size !level !parent

| level == keyBits = (# popSmallArray parent, getDataNode child #)

| otherwise = do

let (# child', tail #) = go size (level - keyBits) (getInternalNode child)

if nullSmallArray child'

then (# popSmallArray parent, tail #)

else (# updateSmallArray parent subIx $ InternalNode child', tail #)

where

child = indexSmallArray parent subIx

subIx = ((size - 1) !>>. level) .&. keyMask

{-# INLINE unsnocTail# #-}

-- | The index of the first element of the tail of the vector (that is, the

tailOffset :: Vector a -> Int

tailOffset vec = (length vec - 1) .&. ((-1) !<<. keyBits)

{-# INLINE tailOffset #-}

-- | \(O(1)\) Get the length of the vector.

length :: Vector a -> Int

length = size

{-# INLINE length #-}

impossibleError :: forall a. a

impossibleError = error "this should be impossible"

moduleError :: forall a. HasCallStack => String -> String -> a

moduleError fun msg = error ("Data.Vector.Persistent.Internal" ++ fun ++ ':' : ' ' : msg)

{-# NOINLINE moduleError #-}

toList :: Vector a -> [a]

toList = pureStreamToList . streamL

{-# INLINE toList #-}

-- | Convert a 'Stream' to a list

pureStreamToList :: Stream Identity a -> [a]

pureStreamToList s = Exts.build (\c n -> runIdentity $ Stream.foldr c n s)

{-# INLINE pureStreamToList #-}

-- | \(O(n)\). Apply a function to all values in the vector.

map :: (a -> b) -> Vector a -> Vector b

map f vec@RootNode {init, tail} = vec {tail = fmap f tail, init = mapSmallArray' go init}

where

go (DataNode as) = DataNode $ fmap f as

go (InternalNode ns) = InternalNode $ mapSmallArray' go ns

{-# INLINE map #-}

-- | \(O(n)\). Apply a function to all values of a vector and its index.

imap :: (Int -> a -> b) -> Vector a -> Vector b

imap f vec@RootNode {size, shift, init, tail}

| size == 0 = empty

| otherwise =

vec

{ init = imapStepSmallArray 0 (1 !<<. shift) (go $! shift - keyBits) init,

tail = imapStepSmallArray (tailOffset vec) 1 f tail

}

where

go _shift i0 (DataNode as) = DataNode $ imapStepSmallArray i0 1 f as

go shift i0 (InternalNode ns) = InternalNode $ imapStepSmallArray i0 (1 !<<. shift) (go $! shift - keyBits) ns

{-# INLINE imap #-}

traverse :: Applicative f => (a -> f b) -> Vector a -> f (Vector b)

traverse f vec@RootNode {init, tail} =

liftA2

(\init tail -> vec {init, tail})

(Traversable.traverse go init)

(Traversable.traverse f tail)

where

go (DataNode as) = DataNode <$> Traversable.traverse f as

go (InternalNode ns) = InternalNode <$> Traversable.traverse go ns

{-# INLINE traverse #-}

itraverse :: Applicative f => (Int -> a -> f b) -> Vector a -> f (Vector b)

itraverse f vec@RootNode {size, shift, init, tail}

| size == 0 = pure empty

| otherwise =

liftA2

(\init tail -> vec {init, tail})

(itraverseStepSmallArray 0 (1 !<<. shift) (go $! shift - keyBits) init)

(itraverseStepSmallArray (tailOffset vec) 1 f tail)

where

go _shift i0 (DataNode as) = DataNode <$> itraverseStepSmallArray i0 1 f as

go shift i0 (InternalNode ns) = InternalNode <$> itraverseStepSmallArray i0 (1 !<<. shift) (go $! shift - keyBits) ns

{-# INLINE itraverse #-}

-- | \(O(n)\). For each pair @(i,a)@ from the vector of index/value pairs,

(//) :: Vector a -> [(Int, a)] -> Vector a

(//) = Exts.inline Foldable.foldl' $ flip $ uncurry update

-- | \(O(n)\). Concatenate two vectors.

(><) :: Vector a -> Vector a -> Vector a

(><) = Exts.inline foldl' snoc

-- | Check the invariant of the vector

invariant :: Vector a -> Bool

invariant _vec = True

-- | \(O(n)\). Create a vector from a list.

fromList :: [a] -> Vector a

fromList = unstream . Stream.fromList

keyBits :: Int

#ifdef TEST

keyBits = 2

keyBits = 5

nodeWidth :: Int

nodeWidth = 1 !<<. keyBits

keyMask :: Int

keyMask = nodeWidth - 1

(!<<.) :: Bits a => a -> Int -> a

(!<<.) = unsafeShiftL

{-# INLINE (!<<.) #-}

(!>>.) :: Bits a => a -> Int -> a

(!>>.) = unsafeShiftR

{-# INLINE (!>>.) #-}

infixl 8 !<<., !>>.

unstream :: Stream Identity a -> Vector a

unstream stream = runST $ do

streamToContents stream >>= \case

(size, tail, [tree]) ->

pure RootNode {size, shift = keyBits, tail, init = pure tree}

(size, tail, ls') -> do

let iterateNodes !shift trees =

nodes (Prelude.reverse trees) >>= \case

[tree] -> pure RootNode {size, shift, tail, init = getInternalNode tree}

trees' -> iterateNodes (shift + keyBits) trees'

iterateNodes keyBits ls'

where

nodes trees = do

buffer <- Buffer.newWithCapacity nodeWidth

(buffer, acc) <-

Foldable.foldlM

( \(!buffer, acc) t ->

if Buffer.length buffer == nodeWidth

then do

result <- Buffer.freeze buffer

buffer <- Buffer.push t $ Buffer.clear buffer

pure (buffer, InternalNode result : acc)

else do

buffer <- Buffer.push t buffer

pure (buffer, acc)

)

(buffer, [])

trees

final <- Buffer.unsafeFreeze buffer

pure $ InternalNode final : acc

{-# INLINE unstream #-}

streamToContents :: PrimMonad m => Stream Identity a -> m (Int, SmallArray a, [Node a])

streamToContents (Stream step s) = do

buffer <- Buffer.newWithCapacity nodeWidth

loop (0 :: Int) buffer [] s

where

loop !size !buffer acc s = do

case runIdentity $ step s of

Stream.Yield x s' -> do

if Buffer.length buffer == nodeWidth

then do

result <- Buffer.freeze buffer

buffer <- Buffer.push x $ Buffer.clear buffer

loop (size + 1) buffer (DataNode result : acc) s'

else do

buffer <- Buffer.push x buffer

loop (size + 1) buffer acc s'

Stream.Skip s' -> loop size buffer acc s'

Stream.Done -> do

tail <- Buffer.unsafeFreeze buffer

pure (size, tail, acc)

{-# INLINE streamToContents #-}

streamL :: Monad m => Vector a -> Stream m a

streamL RootNode {init, tail} = Stream step [(InternalNode init, 0 :: Int), (DataNode tail, 0)]

where

step [] = pure Stream.Done

step ((n, i) : rest) = case n of

InternalNode ns

| i < sizeofSmallArray ns -> do

let !(# ns' #) = indexSmallArray## ns i

!i' = i + 1

pure $ Stream.Skip $ (ns', 0) : (n, i') : rest

| otherwise -> pure $ Stream.Skip rest

DataNode xs

| i < sizeofSmallArray xs -> do

let !(# x #) = indexSmallArray## xs i

!i' = i + 1

pure $ Stream.Yield x $ (n, i') : rest

| otherwise -> pure $ Stream.Skip rest

{-# INLINE step #-}

{-# INLINE streamL #-}

streamR :: Monad m => Vector a -> Stream m a

streamR RootNode {init, tail} = Stream step [(DataNode tail, tailSize), (InternalNode init, initSize)]

where

!tailSize = sizeofSmallArray tail - 1

!initSize = sizeofSmallArray init - 1

step [] = pure Stream.Done

step ((n, i) : rest) = case n of

InternalNode ns

| i < 0 -> pure $ Stream.Skip rest

| otherwise -> do

let !(# n' #) = indexSmallArray## ns i

!i' = i - 1

case n' of

InternalNode ns -> do

let !z = sizeofSmallArray ns - 1

pure $ Stream.Skip $ (n', z) : (n, i') : rest

DataNode xs -> do

let !z = sizeofSmallArray xs - 1

pure $ Stream.Skip $ (n', z) : (n, i') : rest

DataNode xs

| i < 0 -> pure $ Stream.Skip rest

| otherwise -> do

let !(# x #) = indexSmallArray## xs i

!i' = i - 1

pure $ Stream.Yield x $ (n, i') : rest

{-# INLINE step #-}

{-# INLINE streamR #-}

istreamL :: Monad m => Vector a -> Stream m (Int, a)

istreamL = Stream.indexed . streamL

{-# INLINE istreamL #-}

istreamR :: Monad m => Vector a -> Stream m (Int, a)

istreamR vec = Stream.indexedR (length vec - 1) $ streamR vec

{-# INLINE istreamR #-}