{-# 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
#endif

type role Vector representational

-- | A vector.
--
-- The instances are based on those of @Seq@s, which are in turn based on those of lists.
data Vector a
  = -- |
    -- Invariants: The only time tail can be empty is when init is empty.
    -- Otherwise tailOffset will give the wrong value.
    RootNode
    { forall a. Vector a -> Int
size :: !Int,
      -- | 1 << 'shift' is the maximum that each child can contain
      forall a. Vector a -> Int
shift :: !Int,
      forall a. Vector a -> Array (Node a)
init :: !(Array (Node a)),
      forall a. Vector a -> Array a
tail :: !(Array a)
    }

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

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

instance Eq a => Eq (Vector a) where
  == :: Vector a -> Vector a -> Bool
(==) = forall a. Eq a => Vector a -> Vector a -> Bool
persistentVectorEq
  {-# INLINE (==) #-}

instance Ord a => Ord (Vector a) where
  compare :: Vector a -> Vector a -> Ordering
compare = forall a. Ord a => Vector a -> Vector a -> Ordering
persistentVectorCompare
  {-# INLINE compare #-}

instance Functor Vector where
  fmap :: forall a b. (a -> b) -> Vector a -> Vector b
fmap = forall a b. (a -> b) -> Vector a -> Vector b
Data.Vector.Persistent.Internal.map
  {-# INLINE fmap #-}

instance Foldable Vector where
  foldr :: forall a b. (a -> b -> b) -> b -> Vector a -> b
foldr = forall a b. (a -> b -> b) -> b -> Vector a -> b
Data.Vector.Persistent.Internal.foldr
  {-# INLINE foldr #-}
  foldl :: forall b a. (b -> a -> b) -> b -> Vector a -> b
foldl = forall b a. (b -> a -> b) -> b -> Vector a -> b
Data.Vector.Persistent.Internal.foldl
  {-# INLINE foldl #-}
  foldl' :: forall b a. (b -> a -> b) -> b -> Vector a -> b
foldl' = forall b a. (b -> a -> b) -> b -> Vector a -> b
Data.Vector.Persistent.Internal.foldl'
  {-# INLINE foldl' #-}
  foldr' :: forall a b. (a -> b -> b) -> b -> Vector a -> b
foldr' = forall a b. (a -> b -> b) -> b -> Vector a -> b
Data.Vector.Persistent.Internal.foldr'
  {-# INLINE foldr' #-}
  length :: forall a. Vector a -> Int
length = forall a. Vector a -> Int
Data.Vector.Persistent.Internal.length
  {-# INLINE length #-}
  null :: forall a. Vector a -> Bool
null = forall a. Vector a -> Bool
Data.Vector.Persistent.Internal.null
  {-# INLINE null #-}

instance Traversable Vector where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Vector a -> f (Vector b)
traverse = forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Vector a -> f (Vector b)
Data.Vector.Persistent.Internal.traverse
  {-# INLINE traverse #-}

instance Semigroup (Vector a) where
  <> :: Vector a -> Vector a -> Vector a
(<>) = forall a. Vector a -> Vector a -> Vector a
(><)
  {-# INLINE (<>) #-}

instance Monoid (Vector a) where
  mempty :: Vector a
mempty = forall a. Vector a
empty
  {-# INLINE mempty #-}

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

instance Applicative Vector where
  pure :: forall a. a -> Vector a
pure = forall a. a -> Vector a
singleton
  {-# INLINE pure #-}
  Vector (a -> b)
fs <*> :: forall a b. Vector (a -> b) -> Vector a -> Vector b
<*> Vector a
xs = forall b a. (b -> a -> b) -> b -> Vector a -> b
foldl' (\Vector b
acc a -> b
f -> Vector b
acc forall a. Vector a -> Vector a -> Vector a
>< forall a b. (a -> b) -> Vector a -> Vector b
map a -> b
f Vector a
xs) forall a. Vector a
empty Vector (a -> b)
fs

instance Monad Vector where
  Vector a
xs >>= :: forall a b. Vector a -> (a -> Vector b) -> Vector b
>>= a -> Vector b
f = forall b a. (b -> a -> b) -> b -> Vector a -> b
foldl' (\Vector b
acc a
x -> Vector b
acc forall a. Vector a -> Vector a -> Vector a
>< a -> Vector b
f a
x) forall a. Vector a
empty Vector a
xs
  {-# INLINE (>>=) #-}

instance Fail.MonadFail Vector where
  fail :: forall a. String -> Vector a
fail String
_ = forall a. Vector a
empty
  {-# INLINE fail #-}

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

instance MonadPlus Vector

instance NFData1 Vector where
  liftRnf :: forall a. (a -> ()) -> Vector a -> ()
liftRnf a -> ()
f = forall b a. (b -> a -> b) -> b -> Vector a -> b
foldl' (\()
_ a
x -> a -> ()
f a
x) ()
  {-# INLINE liftRnf #-}

data Node a
  = InternalNode {forall a. Node a -> Array (Node a)
getInternalNode :: !(Array (Node a))}
  | DataNode {forall a. Node a -> Array a
getDataNode :: !(Array a)}
  deriving (Int -> Node a -> ShowS
forall a. Show a => Int -> Node a -> ShowS
forall a. Show a => [Node a] -> ShowS
forall a. Show a => Node a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Node a] -> ShowS
$cshowList :: forall a. Show a => [Node a] -> ShowS
show :: Node a -> String
$cshow :: forall a. Show a => Node a -> String
showsPrec :: Int -> Node a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Node a -> ShowS
Show)

instance Eq a => Eq (Node a) where
  == :: Node a -> Node a -> Bool
(==) = forall a. Eq a => Node a -> Node a -> Bool
nodeEq
  {-# INLINE (==) #-}

instance Ord a => Ord (Node a) where
  compare :: Node a -> Node a -> Ordering
compare = forall a. Ord a => Node a -> Node a -> Ordering
nodeCompare
  {-# INLINE compare #-}

instance IsList (Vector a) where
  type Item (Vector a) = a
  fromList :: [Item (Vector a)] -> Vector a
fromList = forall a. [a] -> Vector a
Data.Vector.Persistent.Internal.fromList
  {-# INLINE fromList #-}
  toList :: Vector a -> [Item (Vector a)]
toList = forall a. Vector a -> [a]
Data.Vector.Persistent.Internal.toList
  {-# INLINE toList #-}

-- | \(O(n)\) Lazy right fold.
foldr :: (a -> b -> b) -> b -> Vector a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Vector a -> b
foldr a -> b -> b
f b
z = forall a. Identity a -> a
runIdentity forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
Stream.foldr a -> b -> b
f b
z forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Vector a -> Stream m a
streamL
{-# INLINE foldr #-}

-- | \(O(n)\) Strict right fold.
foldr' :: (a -> b -> b) -> b -> Vector a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Vector a -> b
foldr' a -> b -> b
f b
z = forall a. Identity a -> a
runIdentity forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (m :: * -> *) a b.
Monad m =>
(a -> b -> a) -> a -> Stream m b -> m a
Stream.foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> b
f) b
z forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Vector a -> Stream m a
streamR
{-# INLINE foldr' #-}

-- | \(O(n)\) Lazy left fold.
foldl :: (b -> a -> b) -> b -> Vector a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Vector a -> b
foldl b -> a -> b
f b
z = forall a. Identity a -> a
runIdentity forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
Stream.foldr (forall a b c. (a -> b -> c) -> b -> a -> c
flip b -> a -> b
f) b
z forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Vector a -> Stream m a
streamR
{-# INLINE foldl #-}

-- | \(O(n)\) Strict left fold.
foldl' :: (b -> a -> b) -> b -> Vector a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Vector a -> b
foldl' b -> a -> b
f b
z = forall a. Identity a -> a
runIdentity forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (m :: * -> *) a b.
Monad m =>
(a -> b -> a) -> a -> Stream m b -> m a
Stream.foldl' b -> a -> b
f b
z forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Vector a -> Stream m a
streamL
{-# INLINE foldl' #-}

-- | \(O(n)\) Indexed lazy right fold.
ifoldr :: (Int -> a -> b -> b) -> b -> Vector a -> b
ifoldr :: forall a b. (Int -> a -> b -> b) -> b -> Vector a -> b
ifoldr Int -> a -> b -> b
f b
z = forall a. Identity a -> a
runIdentity forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
Stream.foldr (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> a -> b -> b
f) b
z forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Vector a -> Stream m (Int, a)
istreamL
{-# INLINE ifoldr #-}

-- | \(O(n)\) Indexed lazy left fold.
ifoldl :: (Int -> b -> a -> b) -> b -> Vector a -> b
ifoldl :: forall b a. (Int -> b -> a -> b) -> b -> Vector a -> b
ifoldl Int -> b -> a -> b
f b
z = forall a. Identity a -> a
runIdentity forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
Stream.foldr (\(Int
i, a
x) b
y -> Int -> b -> a -> b
f Int
i b
y a
x) b
z forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Vector a -> Stream m (Int, a)
istreamR
{-# INLINE ifoldl #-}

-- | \(O(n)\) Indexed strict right fold.
ifoldr' :: (Int -> a -> b -> b) -> b -> Vector a -> b
ifoldr' :: forall a b. (Int -> a -> b -> b) -> b -> Vector a -> b
ifoldr' Int -> a -> b -> b
f b
z = forall a. Identity a -> a
runIdentity forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (m :: * -> *) a b.
Monad m =>
(a -> b -> a) -> a -> Stream m b -> m a
Stream.foldl' (\b
y (Int
i, a
x) -> Int -> a -> b -> b
f Int
i a
x b
y) b
z forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Vector a -> Stream m (Int, a)
istreamR
{-# INLINE ifoldr' #-}

-- | \(O(n)\) Indexed strict left fold.
ifoldl' :: (Int -> b -> a -> b) -> b -> Vector a -> b
ifoldl' :: forall b a. (Int -> b -> a -> b) -> b -> Vector a -> b
ifoldl' Int -> b -> a -> b
f b
z = forall a. Identity a -> a
runIdentity forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. forall (m :: * -> *) a b.
Monad m =>
(a -> b -> a) -> a -> Stream m b -> m a
Stream.foldl' (\b
y (Int
i, a
x) -> Int -> b -> a -> b
f Int
i b
y a
x) b
z forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Vector a -> Stream m (Int, a)
istreamL
{-# INLINE ifoldl' #-}

persistentVectorEq :: Eq a => Vector a -> Vector a -> Bool
persistentVectorEq :: forall a. Eq a => Vector a -> Vector a -> Bool
persistentVectorEq
  RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Int
shift :: Int
$sel:shift:RootNode :: forall a. Vector a -> Int
shift, Array (Node a)
init :: Array (Node a)
$sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail}
  RootNode {$sel:size:RootNode :: forall a. Vector a -> Int
size = Int
size', $sel:shift:RootNode :: forall a. Vector a -> Int
shift = Int
shift', $sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init = Array (Node a)
init', $sel:tail:RootNode :: forall a. Vector a -> Array a
tail = Array a
tail'} =
    Int
size forall a. Eq a => a -> a -> Bool
== Int
size' Bool -> Bool -> Bool
&& (Int
size forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| (Int
shift forall a. Eq a => a -> a -> Bool
== Int
shift' Bool -> Bool -> Bool
&& Array a
tail forall a. Eq a => a -> a -> Bool
== Array a
tail' Bool -> Bool -> Bool
&& Array (Node a)
init forall a. Eq a => a -> a -> Bool
== Array (Node a)
init'))
{-# INLINEABLE persistentVectorEq #-}

nodeEq :: Eq a => Node a -> Node a -> Bool
nodeEq :: forall a. Eq a => Node a -> Node a -> Bool
nodeEq (InternalNode Array (Node a)
ns) (InternalNode Array (Node a)
ns') = Array (Node a)
ns forall a. Eq a => a -> a -> Bool
== Array (Node a)
ns'
nodeEq (DataNode Array a
as) (DataNode Array a
as') = Array a
as forall a. Eq a => a -> a -> Bool
== Array a
as'
nodeEq Node a
_ Node a
_ = Bool
False
{-# INLINEABLE nodeEq #-}

persistentVectorCompare :: Ord a => Vector a -> Vector a -> Ordering
persistentVectorCompare :: forall a. Ord a => Vector a -> Vector a -> Ordering
persistentVectorCompare
  RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Array (Node a)
init :: Array (Node a)
$sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail}
  RootNode {$sel:size:RootNode :: forall a. Vector a -> Int
size = Int
size', $sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init = Array (Node a)
init', $sel:tail:RootNode :: forall a. Vector a -> Array a
tail = Array a
tail'} =
    forall a. Ord a => a -> a -> Ordering
compare Int
size Int
size'
      forall a. Semigroup a => a -> a -> a
<> if Int
size forall a. Eq a => a -> a -> Bool
== Int
0
        then Ordering
EQ
        else forall a. Ord a => a -> a -> Ordering
compare Array (Node a)
init Array (Node a)
init' forall a. Semigroup a => a -> a -> a
<> forall a. Ord a => a -> a -> Ordering
compare Array a
tail Array a
tail'
{-# INLINEABLE persistentVectorCompare #-}

nodeCompare :: Ord a => Node a -> Node a -> Ordering
nodeCompare :: forall a. Ord a => Node a -> Node a -> Ordering
nodeCompare (DataNode Array a
as) (DataNode Array a
as') = forall a. Ord a => a -> a -> Ordering
compare Array a
as Array a
as'
nodeCompare (InternalNode Array (Node a)
ns) (InternalNode Array (Node a)
ns') = forall a. Ord a => a -> a -> Ordering
compare Array (Node a)
ns Array (Node a)
ns'
nodeCompare (DataNode Array a
_) (InternalNode Array (Node a)
_) = Ordering
LT
nodeCompare (InternalNode Array (Node a)
_) (DataNode Array a
_) = Ordering
GT
{-# INLINEABLE nodeCompare #-}

-- | \(O(1)\). A vector with a single element.
singleton :: a -> Vector a
singleton :: forall a. a -> Vector a
singleton a
a = RootNode {$sel:size:RootNode :: Int
size = Int
1, $sel:shift:RootNode :: Int
shift = Int
keyBits, $sel:tail:RootNode :: Array a
tail = forall a. a -> Array a
singletonSmallArray a
a, $sel:init:RootNode :: Array (Node a)
init = forall a. SmallArray a
emptySmallArray}
{-# INLINE singleton #-}

-- | \(O(1)\). The empty vector.
empty :: Vector a
empty :: forall a. Vector a
empty = RootNode {$sel:size:RootNode :: Int
size = Int
0, $sel:shift:RootNode :: Int
shift = Int
keyBits, $sel:init:RootNode :: Array (Node a)
init = forall a. SmallArray a
emptySmallArray, $sel:tail:RootNode :: Array a
tail = forall a. SmallArray a
emptySmallArray}
{-# NOINLINE empty #-}

-- | \(O(1)\) Return 'True' if the vector is empty, 'False' otherwise.
null :: Vector a -> Bool
null :: forall a. Vector a -> Bool
null Vector a
xs = forall a. Vector a -> Int
length Vector a
xs forall a. Eq a => a -> a -> Bool
== Int
0
{-# INLINE null #-}

-- | \(O(\log n)\). An alias for 'snoc'
-- Mnemonic: a triangle with the single element at the pointy end.
(|>) :: Vector a -> a -> Vector a
|> :: forall a. Vector a -> a -> Vector a
(|>) = forall a. Vector a -> a -> Vector a
snoc
{-# INLINE (|>) #-}

-- | \(O(\log n)\). A bidirectional pattern synonym viewing the rear of a non-empty
-- sequence.
pattern (:|>) :: Vector a -> a -> Vector a
pattern vec $b:|> :: forall a. Vector a -> a -> Vector a
$m:|> :: forall {r} {a}.
Vector a -> (Vector a -> a -> r) -> ((# #) -> r) -> r
:|> x <-
  (unsnoc -> Just (vec, x))
  where
    Vector a
vec :|> a
x = Vector a
vec forall a. Vector a -> a -> Vector a
|> a
x

infixl 5 :|>

-- | \(O(1)\). A bidirectional pattern synonym matching an empty sequence.
pattern Empty :: Vector a
pattern $bEmpty :: forall a. Vector a
$mEmpty :: forall {r} {a}. Vector a -> ((# #) -> r) -> ((# #) -> r) -> r
Empty <-
  (null -> True)
  where
    Empty = forall a. Vector a
empty

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

-- | \(O(\log n)\) Add an element to the end of the vector.
snoc :: Vector a -> a -> Vector a
snoc :: forall a. Vector a -> a -> Vector a
snoc vec :: Vector a
vec@RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail} a
a
  -- Room in tail, and vector non-empty
  | (Int
size forall a. Bits a => a -> a -> a
.&. Int
keyMask) forall a. Eq a => a -> a -> Bool
/= Int
0 =
      Vector a
vec
        { $sel:tail:RootNode :: Array a
tail = forall a. Array a -> a -> Array a
snocSmallArray Array a
tail a
a,
          $sel:size:RootNode :: Int
size = Int
size forall a. Num a => a -> a -> a
+ Int
1
        }
  | Bool
otherwise = forall a. Vector a -> Int -> Array a -> Vector a
snocArr Vector a
vec Int
1 forall a b. (a -> b) -> a -> b
$ forall a. a -> Array a
singletonSmallArray a
a
{-# INLINEABLE snoc #-}

-- Invariant: the tail must be large enough to mutably write to it
-- Best to use this with emptyMaxTail
-- After calling this many times you must run shrink
unsafeSnoc :: Vector a -> a -> Vector a
unsafeSnoc :: forall a. Vector a -> a -> Vector a
unsafeSnoc vec :: Vector a
vec@RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail} a
a
  -- Room in tail, and vector non-empty
  | (Int
size forall a. Bits a => a -> a -> a
.&. Int
keyMask) forall a. Eq a => a -> a -> Bool
/= Int
0 =
      Vector a
vec
        { $sel:tail:RootNode :: Array a
tail =
            -- update the array in place
            forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
              SmallMutableArray s a
marr <- forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> m (SmallMutableArray (PrimState m) a)
unsafeThawSmallArray Array a
tail
              forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
marr (Int
size forall a. Bits a => a -> a -> a
.&. Int
keyMask) a
a
              forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray s a
marr,
          $sel:size:RootNode :: Int
size = Int
size forall a. Num a => a -> a -> a
+ Int
1
        }
  | Bool
otherwise = forall a. Vector a -> Int -> Array a -> Vector a
snocArr Vector a
vec Int
1 forall a b. (a -> b) -> a -> b
$ forall a. a -> Array a
singletonSmallArray a
a
{-# INLINEABLE unsafeSnoc #-}

-- Invariant: arr cannot be empty
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 :: forall a. Vector a -> Int -> Array a -> Vector a
snocArr vec :: Vector a
vec@RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Int
shift :: Int
$sel:shift:RootNode :: forall a. Vector a -> Int
shift, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail} Int
addedSize Array a
arr
  | forall a. Vector a -> Bool
null Vector a
vec =
      RootNode
        { $sel:size:RootNode :: Int
size = Int
addedSize,
          $sel:shift:RootNode :: Int
shift = Int
keyBits,
          $sel:tail:RootNode :: Array a
tail = Array a
arr,
          $sel:init:RootNode :: Array (Node a)
init = forall a. SmallArray a
emptySmallArray
        }
  | Int
size forall a. Bits a => a -> Int -> a
!>>. Int
keyBits forall a. Ord a => a -> a -> Bool
> Int
1 forall a. Bits a => a -> Int -> a
!<<. Int
shift =
      RootNode
        { $sel:size:RootNode :: Int
size = Int
size forall a. Num a => a -> a -> a
+ Int
addedSize,
          $sel:shift:RootNode :: Int
shift = Int
shift forall a. Num a => a -> a -> a
+ Int
keyBits,
          $sel:init:RootNode :: Array (Node a)
init =
            let !path :: Node a
path = forall a. Int -> Array a -> Node a
newPath Int
shift Array a
tail
                !internal :: Node a
internal = forall a. Array (Node a) -> Node a
InternalNode forall a b. (a -> b) -> a -> b
$ forall a. Vector a -> Array (Node a)
init Vector a
vec
             in forall a. a -> a -> Array a
twoSmallArray Node a
internal Node a
path,
          $sel:tail:RootNode :: Array a
tail = Array a
arr
        }
  | Bool
otherwise =
      RootNode
        { $sel:size:RootNode :: Int
size = Int
size forall a. Num a => a -> a -> a
+ Int
addedSize,
          Int
shift :: Int
$sel:shift:RootNode :: Int
shift,
          $sel:init:RootNode :: Array (Node a)
init = forall a. Int -> Array a -> Int -> Array (Node a) -> Array (Node a)
snocTail Int
size Array a
tail Int
shift forall a b. (a -> b) -> a -> b
$ forall a. Vector a -> Array (Node a)
init Vector a
vec,
          $sel:tail:RootNode :: Array a
tail = Array a
arr
        }
{-# INLINE snocArr #-}

snocTail :: Int -> Array a -> Int -> Array (Node a) -> Array (Node a)
snocTail :: forall a. Int -> Array a -> Int -> Array (Node a) -> Array (Node a)
snocTail Int
size Array a
tail = Int -> Array (Node a) -> Array (Node a)
go
  where
    go :: Int -> Array (Node a) -> Array (Node a)
go !Int
level !Array (Node a)
parent = forall a. Array a -> Int -> a -> Array a
updateResizeSmallArray Array (Node a)
parent Int
subIx Node a
toInsert
      where
        toInsert :: Node a
toInsert
          | Int
level forall a. Eq a => a -> a -> Bool
== Int
keyBits = forall a. Array a -> Node a
DataNode Array a
tail
          | Int
subIx forall a. Ord a => a -> a -> Bool
< forall a. SmallArray a -> Int
sizeofSmallArray Array (Node a)
parent =
              let vec' :: Node a
vec' = forall a. SmallArray a -> Int -> a
indexSmallArray Array (Node a)
parent Int
subIx
               in forall a. Array (Node a) -> Node a
InternalNode forall a b. (a -> b) -> a -> b
$ Int -> Array (Node a) -> Array (Node a)
go (Int
level forall a. Num a => a -> a -> a
- Int
keyBits) (forall a. Node a -> Array (Node a)
getInternalNode Node a
vec')
          | Bool
otherwise = forall a. Int -> Array a -> Node a
newPath (Int
level forall a. Num a => a -> a -> a
- Int
keyBits) Array a
tail
        subIx :: Int
subIx = ((Int
size forall a. Num a => a -> a -> a
- Int
1) forall a. Bits a => a -> Int -> a
!>>. Int
level) forall a. Bits a => a -> a -> a
.&. Int
keyMask
{-# INLINE snocTail #-}

newPath :: Int -> Array a -> Node a
newPath :: forall a. Int -> Array a -> Node a
newPath Int
0 Array a
tail = forall a. Array a -> Node a
DataNode Array a
tail
newPath Int
level Array a
tail = forall a. Array (Node a) -> Node a
InternalNode forall a b. (a -> b) -> a -> b
$ forall a. a -> Array a
singletonSmallArray forall a b. (a -> b) -> a -> b
$! forall a. Int -> Array a -> Node a
newPath (Int
level forall a. Num a => a -> a -> a
- Int
keyBits) Array a
tail

unsafeIndex :: Vector a -> Int -> a
unsafeIndex :: forall a. Vector a -> Int -> a
unsafeIndex Vector a
vec Int
ix | (# a
a #) <- forall a. a -> a
Exts.inline forall a. Vector a -> Int -> (# a #)
unsafeIndex# Vector a
vec Int
ix = a
a
{-# NOINLINE unsafeIndex #-}

unsafeIndex# :: Vector a -> Int -> (# a #)
unsafeIndex# :: forall a. Vector a -> Int -> (# a #)
unsafeIndex# Vector a
vec Int
ix
  | Int
ix forall a. Ord a => a -> a -> Bool
>= forall a. Vector a -> Int
tailOffset Vector a
vec = forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## (forall a. Vector a -> Array a
tail Vector a
vec) (Int
ix forall a. Bits a => a -> a -> a
.&. Int
keyMask)
  -- no need to use keyMask here as we are at the top
  | Bool
otherwise = forall {a}. Int -> Int -> Node a -> (# a #)
go Int
ix (forall a. Vector a -> Int
shift Vector a
vec forall a. Num a => a -> a -> a
- Int
keyBits) (forall a. SmallArray a -> Int -> a
indexSmallArray (forall a. Vector a -> Array (Node a)
init Vector a
vec) (Int
ix forall a. Bits a => a -> Int -> a
!>>. forall a. Vector a -> Int
shift Vector a
vec))
  where
    go :: Int -> Int -> Node a -> (# a #)
go Int
ix Int
0 !Node a
node = forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## (forall a. Node a -> Array a
getDataNode Node a
node) (Int
ix forall a. Bits a => a -> a -> a
.&. Int
keyMask)
    go Int
ix Int
level !Node a
node = Int -> Int -> Node a -> (# a #)
go Int
ix (Int
level forall a. Num a => a -> a -> a
- Int
keyBits) Node a
node'
      where
        !(# Node a
node' #) = forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## (forall a. Node a -> Array (Node a)
getInternalNode Node a
node) ((Int
ix forall a. Bits a => a -> Int -> a
!>>. Int
level) forall a. Bits a => a -> a -> a
.&. Int
keyMask)
{-# NOINLINE unsafeIndex# #-}

lookup# :: Int -> Vector a -> (# (# #) | a #)
lookup# :: forall a. Int -> Vector a -> (# (# #) | a #)
lookup# Int
ix Vector a
vec
  | (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix :: Word) forall a. Ord a => a -> a -> Bool
>= forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Vector a -> Int
length Vector a
vec) = (# (# #) | #)
  | Bool
otherwise = case forall a. a -> a
Exts.inline forall a. Vector a -> Int -> (# a #)
unsafeIndex# Vector a
vec Int
ix of (# a
x #) -> (# | a
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 :: forall a. Int -> Vector a -> Maybe a
lookup Int
ix Vector a
vec
  | (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix :: Word) forall a. Ord a => a -> a -> Bool
>= forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Vector a -> Int
length Vector a
vec) = forall a. Maybe a
Nothing
  | Bool
otherwise = case forall a. Vector a -> Int -> (# a #)
unsafeIndex# Vector a
vec Int
ix of (# a
x #) -> forall a. a -> Maybe a
Just a
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 :: forall a. HasCallStack => Int -> Vector a -> a
index Int
ix Vector a
vec
  | Int
ix forall a. Ord a => a -> a -> Bool
< Int
0 = forall a. HasCallStack => String -> String -> a
moduleError String
"index" forall a b. (a -> b) -> a -> b
$ String
"negative index: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
ix
  | Int
ix forall a. Ord a => a -> a -> Bool
>= forall a. Vector a -> Int
length Vector a
vec = forall a. HasCallStack => String -> String -> a
moduleError String
"index" forall a b. (a -> b) -> a -> b
$ String
"index too large: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Int
ix
  | Bool
otherwise = forall a. a -> a
Exts.inline forall a. Vector a -> Int -> a
unsafeIndex Vector a
vec Int
ix
{-# INLINEABLE index #-}

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

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

-- | \(O(\log n)\). 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 :: (a -> a) -> Int -> Vector a -> Vector a
adjust :: forall a. (a -> a) -> Int -> Vector a -> Vector a
adjust a -> a
f = forall a. (a -> (# a #)) -> Int -> Vector a -> Vector a
adjust# forall a b. (a -> b) -> a -> b
$ \a
x -> (# a -> a
f a
x #)
{-# INLINE adjust #-}

-- needs better tests for this
adjust# :: (a -> (# a #)) -> Int -> Vector a -> Vector a
adjust# :: forall a. (a -> (# a #)) -> Int -> Vector a -> Vector a
adjust# a -> (# a #)
f Int
ix vec :: Vector a
vec@RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Int
shift :: Int
$sel:shift:RootNode :: forall a. Vector a -> Int
shift, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail}
  -- Invalid index. This funny business uses a single test to determine whether
  -- ix is too small (negative) or too large (at least sz).
  | (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix :: Word) forall a. Ord a => a -> a -> Bool
>= forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
size = Vector a
vec
  | Int
ix forall a. Ord a => a -> a -> Bool
>= forall a. Vector a -> Int
tailOffset Vector a
vec = Vector a
vec {$sel:tail:RootNode :: Array a
tail = forall a. Array a -> Int -> (a -> (# a #)) -> Array a
modifySmallArray# Array a
tail (Int
ix forall a. Bits a => a -> a -> a
.&. Int
keyMask) a -> (# a #)
f}
  | Bool
otherwise = Vector a
vec {$sel:init:RootNode :: Array (Node a)
init = Int -> Int -> Array (Node a) -> Array (Node a)
go Int
ix Int
shift (forall a. Vector a -> Array (Node a)
init Vector a
vec)}
  where
    go :: Int -> Int -> Array (Node a) -> Array (Node a)
go Int
ix Int
level Array (Node a)
vec
      | Int
level forall a. Eq a => a -> a -> Bool
== Int
keyBits =
          let !node :: Node a
node = forall a. Array a -> Node a
DataNode forall a b. (a -> b) -> a -> b
$ forall a. Array a -> Int -> (a -> (# a #)) -> Array a
modifySmallArray# (forall a. Node a -> Array a
getDataNode Node a
vec') (Int
ix forall a. Bits a => a -> a -> a
.&. Int
keyMask) a -> (# a #)
f
           in forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
vec Int
ix' Node a
node
      | Bool
otherwise =
          let !node :: Array (Node a)
node = Int -> Int -> Array (Node a) -> Array (Node a)
go Int
ix (Int
level forall a. Num a => a -> a -> a
- Int
keyBits) (forall a. Node a -> Array (Node a)
getInternalNode Node a
vec')
           in forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
vec Int
ix' forall a b. (a -> b) -> a -> b
$! forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
node
      where
        ix' :: Int
ix' = (Int
ix forall a. Bits a => a -> Int -> a
!>>. Int
level) forall a. Bits a => a -> a -> a
.&. Int
keyMask
        vec' :: Node a
vec' = forall a. SmallArray a -> Int -> a
indexSmallArray Array (Node a)
vec Int
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 :: forall (f :: * -> *) a.
Applicative f =>
(a -> f a) -> Int -> Vector a -> f (Vector a)
adjustF a -> f a
f Int
ix vec :: Vector a
vec@RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Int
shift :: Int
$sel:shift:RootNode :: forall a. Vector a -> Int
shift, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail}
  -- Invalid index. This funny business uses a single test to determine whether
  -- ix is too small (negative) or too large (at least sz).
  | (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix :: Word) forall a. Ord a => a -> a -> Bool
>= forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
size = forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector a
vec
  | Int
ix forall a. Ord a => a -> a -> Bool
>= forall a. Vector a -> Int
tailOffset Vector a
vec = (\Array a
tail -> Vector a
vec {Array a
tail :: Array a
$sel:tail:RootNode :: Array a
tail}) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a.
Functor f =>
Array a -> Int -> (a -> f a) -> f (Array a)
modifySmallArrayF Array a
tail (Int
ix forall a. Bits a => a -> a -> a
.&. Int
keyMask) a -> f a
f
  | Bool
otherwise = (\Array (Node a)
init -> Vector a
vec {Array (Node a)
init :: Array (Node a)
$sel:init:RootNode :: Array (Node a)
init}) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> Array (Node a) -> f (Array (Node a))
go Int
ix Int
shift (forall a. Vector a -> Array (Node a)
init Vector a
vec)
  where
    go :: Int -> Int -> Array (Node a) -> f (Array (Node a))
go Int
ix Int
level Array (Node a)
vec
      | Int
level forall a. Eq a => a -> a -> Bool
== Int
keyBits =
          (\Array a
node' -> forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
vec Int
ix' forall a b. (a -> b) -> a -> b
$! forall a. Array a -> Node a
DataNode Array a
node')
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a.
Functor f =>
Array a -> Int -> (a -> f a) -> f (Array a)
modifySmallArrayF (forall a. Node a -> Array a
getDataNode Node a
vec') (Int
ix forall a. Bits a => a -> a -> a
.&. Int
keyMask) a -> f a
f
      | Bool
otherwise =
          (\Array (Node a)
node -> forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
vec Int
ix' forall a b. (a -> b) -> a -> b
$! forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
node)
            forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> Array (Node a) -> f (Array (Node a))
go Int
ix (Int
level forall a. Num a => a -> a -> a
- Int
keyBits) (forall a. Node a -> Array (Node a)
getInternalNode Node a
vec')
      where
        ix' :: Int
ix' = (Int
ix forall a. Bits a => a -> Int -> a
!>>. Int
level) forall a. Bits a => a -> a -> a
.&. Int
keyMask
        vec' :: Node a
vec' = forall a. SmallArray a -> Int -> a
indexSmallArray Array (Node a)
vec Int
ix'
{-# INLINE adjustF #-}

-- | \(O(\log n)\). Replace the element at the specified position.
-- If the position is out of range, the original sequence is returned.
update :: Int -> a -> Vector a -> Vector a
-- we could use adjust# (\_ -> (# x #)) to implement this
-- and the const like function would get optimized out
-- but we don't because we don't want to create any closures for the go function
-- so we rewrite out the loop and also lambda lift some arguments
-- also: trees are very shallow, so the loop won't be called much.
-- So allocating a closure to not have pass the arguments on the stack has too much overhead
update :: forall a. Int -> a -> Vector a -> Vector a
update Int
ix a
x vec :: Vector a
vec@RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Int
shift :: Int
$sel:shift:RootNode :: forall a. Vector a -> Int
shift, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail}
  -- Invalid index. This funny business uses a single test to determine whether
  -- ix is too small (negative) or too large (at least sz).
  | (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix :: Word) forall a. Ord a => a -> a -> Bool
>= forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
size = Vector a
vec
  | Int
ix forall a. Ord a => a -> a -> Bool
>= forall a. Vector a -> Int
tailOffset Vector a
vec = Vector a
vec {$sel:tail:RootNode :: Array a
tail = forall a. Array a -> Int -> a -> Array a
updateSmallArray Array a
tail (Int
ix forall a. Bits a => a -> a -> a
.&. Int
keyMask) a
x}
  | Bool
otherwise = Vector a
vec {$sel:init:RootNode :: Array (Node a)
init = forall {a}. Int -> a -> Int -> Array (Node a) -> Array (Node a)
go Int
ix a
x Int
shift (forall a. Vector a -> Array (Node a)
init Vector a
vec)}
  where
    go :: Int -> a -> Int -> Array (Node a) -> Array (Node a)
go Int
ix a
x Int
level Array (Node a)
vec
      | Int
level forall a. Eq a => a -> a -> Bool
== Int
keyBits =
          let !node :: Node a
node = forall a. Array a -> Node a
DataNode forall a b. (a -> b) -> a -> b
$ forall a. Array a -> Int -> a -> Array a
updateSmallArray (forall a. Node a -> Array a
getDataNode Node a
vec') (Int
ix forall a. Bits a => a -> a -> a
.&. Int
keyMask) a
x
           in forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
vec Int
ix' Node a
node
      | Bool
otherwise =
          let !node :: Array (Node a)
node = Int -> a -> Int -> Array (Node a) -> Array (Node a)
go Int
ix a
x (Int
level forall a. Num a => a -> a -> a
- Int
keyBits) (forall a. Node a -> Array (Node a)
getInternalNode Node a
vec')
           in forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
vec Int
ix' forall a b. (a -> b) -> a -> b
$! forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
node
      where
        ix' :: Int
ix' = (Int
ix forall a. Bits a => a -> Int -> a
!>>. Int
level) forall a. Bits a => a -> a -> a
.&. Int
keyMask
        vec' :: Node a
vec' = forall a. SmallArray a -> Int -> a
indexSmallArray Array (Node a)
vec Int
ix'
{-# INLINEABLE update #-}

-- | \(O(\log n)\). Decompose a list into its head and tail.
--
-- * If the list is empty, returns 'Nothing'.
-- * If the list is non-empty, returns @'Just' (x, xs)@,
-- where @x@ is the head of the list and @xs@ its tail.
unsnoc :: Vector a -> Maybe (Vector a, a)
unsnoc :: forall a. Vector a -> Maybe (Vector a, a)
unsnoc vec :: Vector a
vec@RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail, Array (Node a)
init :: Array (Node a)
$sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init, Int
shift :: Int
$sel:shift:RootNode :: forall a. Vector a -> Int
shift}
  | Int
0 <- Int
size = forall a. Maybe a
Nothing
  -- we need to have this case because we can't run unsnocTail, there is nothing left in the tail
  | Int
1 <- Int
size, (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## Array a
tail Int
0 = forall a. a -> Maybe a
Just (forall a. Vector a
empty, a
x)
  | Bool
otherwise = case forall a. SmallArray a -> Maybe (SmallArray a, a)
unsnocSmallArray Array a
tail of
      Maybe (Array a, a)
Nothing ->
        let !(# Array (Node a)
init', Array a
tail' #) = forall a.
Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
unsnocTail# Int
size Int
shift Array (Node a)
init
            !(# a
x #) = forall a. SmallArray a -> (# a #)
lastSmallArray# Array a
tail'
         in forall a. a -> Maybe a
Just (Vector a
vec {$sel:size:RootNode :: Int
size = Int
size forall a. Num a => a -> a -> a
- Int
1, $sel:init:RootNode :: Array (Node a)
init = Array (Node a)
init', $sel:tail:RootNode :: Array a
tail = forall a. Array a -> Array a
popSmallArray Array a
tail'}, a
x)
      Just (Array a
tail', a
a) -> forall a. a -> Maybe a
Just (Vector a
vec {$sel:size:RootNode :: Int
size = Int
size forall a. Num a => a -> a -> a
- Int
1, $sel:tail:RootNode :: Array a
tail = Array a
tail'}, a
a)
{-# INLINEABLE unsnoc #-}

unsnocTail# :: Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
unsnocTail# :: forall a.
Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
unsnocTail# = forall a.
Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
go
  where
    go :: Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
go Int
size !Int
level !Array (Node a)
parent
      | Int
level forall a. Eq a => a -> a -> Bool
== Int
keyBits = (# forall a. Array a -> Array a
popSmallArray Array (Node a)
parent, forall a. Node a -> Array a
getDataNode Node a
child #)
      | Bool
otherwise = do
          let (# Array (Node a)
child', Array a
tail #) = Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
go Int
size (Int
level forall a. Num a => a -> a -> a
- Int
keyBits) (forall a. Node a -> Array (Node a)
getInternalNode Node a
child)
          if forall a. SmallArray a -> Bool
nullSmallArray Array (Node a)
child'
            then (# forall a. Array a -> Array a
popSmallArray Array (Node a)
parent, Array a
tail #)
            else (# forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
parent Int
subIx forall a b. (a -> b) -> a -> b
$ forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
child', Array a
tail #)
      where
        child :: Node a
child = forall a. SmallArray a -> Int -> a
indexSmallArray Array (Node a)
parent Int
subIx
        subIx :: Int
subIx = ((Int
size forall a. Num a => a -> a -> a
- Int
1) forall a. Bits a => a -> Int -> a
!>>. Int
level) forall a. Bits a => a -> a -> a
.&. Int
keyMask
{-# INLINE unsnocTail# #-}

-- | The index of the first element of the tail of the vector (that is, the
-- *last* element of the list representing the tail). This is also the number
-- of elements stored in the array tree.
--
-- Caution: Only gives a sensible result if the vector is nonempty.
tailOffset :: Vector a -> Int
tailOffset :: forall a. Vector a -> Int
tailOffset Vector a
vec = (forall a. Vector a -> Int
length Vector a
vec forall a. Num a => a -> a -> a
- Int
1) forall a. Bits a => a -> a -> a
.&. ((-Int
1) forall a. Bits a => a -> Int -> a
!<<. Int
keyBits)
{-# INLINE tailOffset #-}

-- | \(O(1)\) Get the length of the vector.
length :: Vector a -> Int
length :: forall a. Vector a -> Int
length = forall a. Vector a -> Int
size
{-# INLINE length #-}

impossibleError :: forall a. a
impossibleError :: forall a. a
impossibleError = forall a. HasCallStack => String -> a
error String
"this should be impossible"

moduleError :: forall a. HasCallStack => String -> String -> a
moduleError :: forall a. HasCallStack => String -> String -> a
moduleError String
fun String
msg = forall a. HasCallStack => String -> a
error (String
"Data.Vector.Persistent.Internal" forall a. [a] -> [a] -> [a]
++ String
fun forall a. [a] -> [a] -> [a]
++ Char
':' forall a. a -> [a] -> [a]
: Char
' ' forall a. a -> [a] -> [a]
: String
msg)
{-# NOINLINE moduleError #-}

toList :: Vector a -> [a]
toList :: forall a. Vector a -> [a]
toList = forall a. Stream Identity a -> [a]
pureStreamToList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Vector a -> Stream m a
streamL
{-# INLINE toList #-}

-- | Convert a 'Stream' to a list
pureStreamToList :: Stream Identity a -> [a]
pureStreamToList :: forall a. Stream Identity a -> [a]
pureStreamToList Stream Identity a
s = forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
Exts.build (\a -> b -> b
c b
n -> forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
Stream.foldr a -> b -> b
c b
n Stream Identity a
s)
{-# INLINE pureStreamToList #-}

-- | \(O(n)\). Apply a function to all values in the vector.
map :: (a -> b) -> Vector a -> Vector b
map :: forall a b. (a -> b) -> Vector a -> Vector b
map a -> b
f vec :: Vector a
vec@RootNode {Array (Node a)
init :: Array (Node a)
$sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail} = Vector a
vec {$sel:tail:RootNode :: Array b
tail = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Array a
tail, $sel:init:RootNode :: Array (Node b)
init = forall a b. (a -> b) -> SmallArray a -> SmallArray b
mapSmallArray' Node a -> Node b
go Array (Node a)
init}
  where
    go :: Node a -> Node b
go (DataNode Array a
as) = forall a. Array a -> Node a
DataNode forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f Array a
as
    go (InternalNode Array (Node a)
ns) = forall a. Array (Node a) -> Node a
InternalNode forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> SmallArray a -> SmallArray b
mapSmallArray' Node a -> Node b
go Array (Node a)
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 :: forall a b. (Int -> a -> b) -> Vector a -> Vector b
imap Int -> a -> b
f vec :: Vector a
vec@RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Int
shift :: Int
$sel:shift:RootNode :: forall a. Vector a -> Int
shift, Array (Node a)
init :: Array (Node a)
$sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail}
  | Int
size forall a. Eq a => a -> a -> Bool
== Int
0 = forall a. Vector a
empty
  | Bool
otherwise =
      Vector a
vec
        { $sel:init:RootNode :: Array (Node b)
init = forall a b.
Int -> Int -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray Int
0 (Int
1 forall a. Bits a => a -> Int -> a
!<<. Int
shift) (Int -> Int -> Node a -> Node b
go forall a b. (a -> b) -> a -> b
$! Int
shift forall a. Num a => a -> a -> a
- Int
keyBits) Array (Node a)
init,
          $sel:tail:RootNode :: Array b
tail = forall a b.
Int -> Int -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray (forall a. Vector a -> Int
tailOffset Vector a
vec) Int
1 Int -> a -> b
f Array a
tail
        }
  where
    go :: Int -> Int -> Node a -> Node b
go Int
_shift Int
i0 (DataNode Array a
as) = forall a. Array a -> Node a
DataNode forall a b. (a -> b) -> a -> b
$ forall a b.
Int -> Int -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray Int
i0 Int
1 Int -> a -> b
f Array a
as
    go Int
shift Int
i0 (InternalNode Array (Node a)
ns) = forall a. Array (Node a) -> Node a
InternalNode forall a b. (a -> b) -> a -> b
$ forall a b.
Int -> Int -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray Int
i0 (Int
1 forall a. Bits a => a -> Int -> a
!<<. Int
shift) (Int -> Int -> Node a -> Node b
go forall a b. (a -> b) -> a -> b
$! Int
shift forall a. Num a => a -> a -> a
- Int
keyBits) Array (Node a)
ns
{-# INLINE imap #-}

traverse :: Applicative f => (a -> f b) -> Vector a -> f (Vector b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Vector a -> f (Vector b)
traverse a -> f b
f vec :: Vector a
vec@RootNode {Array (Node a)
init :: Array (Node a)
$sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail} =
  forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2
    (\Array (Node b)
init Array b
tail -> Vector a
vec {Array (Node b)
init :: Array (Node b)
$sel:init:RootNode :: Array (Node b)
init, Array b
tail :: Array b
$sel:tail:RootNode :: Array b
tail})
    (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
Traversable.traverse Node a -> f (Node b)
go Array (Node a)
init)
    (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
Traversable.traverse a -> f b
f Array a
tail)
  where
    go :: Node a -> f (Node b)
go (DataNode Array a
as) = forall a. Array a -> Node a
DataNode forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
Traversable.traverse a -> f b
f Array a
as
    go (InternalNode Array (Node a)
ns) = forall a. Array (Node a) -> Node a
InternalNode forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
Traversable.traverse Node a -> f (Node b)
go Array (Node a)
ns
{-# INLINE traverse #-}

itraverse :: Applicative f => (Int -> a -> f b) -> Vector a -> f (Vector b)
itraverse :: forall (f :: * -> *) a b.
Applicative f =>
(Int -> a -> f b) -> Vector a -> f (Vector b)
itraverse Int -> a -> f b
f vec :: Vector a
vec@RootNode {Int
size :: Int
$sel:size:RootNode :: forall a. Vector a -> Int
size, Int
shift :: Int
$sel:shift:RootNode :: forall a. Vector a -> Int
shift, Array (Node a)
init :: Array (Node a)
$sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail}
  | Int
size forall a. Eq a => a -> a -> Bool
== Int
0 = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Vector a
empty
  | Bool
otherwise =
      forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2
        (\Array (Node b)
init Array b
tail -> Vector a
vec {Array (Node b)
init :: Array (Node b)
$sel:init:RootNode :: Array (Node b)
init, Array b
tail :: Array b
$sel:tail:RootNode :: Array b
tail})
        (forall (f :: * -> *) a b.
Applicative f =>
Int -> Int -> (Int -> a -> f b) -> SmallArray a -> f (SmallArray b)
itraverseStepSmallArray Int
0 (Int
1 forall a. Bits a => a -> Int -> a
!<<. Int
shift) (Int -> Int -> Node a -> f (Node b)
go forall a b. (a -> b) -> a -> b
$! Int
shift forall a. Num a => a -> a -> a
- Int
keyBits) Array (Node a)
init)
        (forall (f :: * -> *) a b.
Applicative f =>
Int -> Int -> (Int -> a -> f b) -> SmallArray a -> f (SmallArray b)
itraverseStepSmallArray (forall a. Vector a -> Int
tailOffset Vector a
vec) Int
1 Int -> a -> f b
f Array a
tail)
  where
    go :: Int -> Int -> Node a -> f (Node b)
go Int
_shift Int
i0 (DataNode Array a
as) = forall a. Array a -> Node a
DataNode forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a b.
Applicative f =>
Int -> Int -> (Int -> a -> f b) -> SmallArray a -> f (SmallArray b)
itraverseStepSmallArray Int
i0 Int
1 Int -> a -> f b
f Array a
as
    go Int
shift Int
i0 (InternalNode Array (Node a)
ns) = forall a. Array (Node a) -> Node a
InternalNode forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (f :: * -> *) a b.
Applicative f =>
Int -> Int -> (Int -> a -> f b) -> SmallArray a -> f (SmallArray b)
itraverseStepSmallArray Int
i0 (Int
1 forall a. Bits a => a -> Int -> a
!<<. Int
shift) (Int -> Int -> Node a -> f (Node b)
go forall a b. (a -> b) -> a -> b
$! Int
shift forall a. Num a => a -> a -> a
- Int
keyBits) Array (Node a)
ns
{-# INLINE itraverse #-}

-- | \(O(n)\). For each pair @(i,a)@ from the vector of index/value pairs,
-- replace the vector element at position @i@ by @a@.
--
-- > update <5,9,2,7> <(2,1),(0,3),(2,8)> = <3,9,8,7>
(//) :: Vector a -> [(Int, a)] -> Vector a
// :: forall a. Vector a -> [(Int, a)] -> Vector a
(//) = forall a. a -> a
Exts.inline forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. Int -> a -> Vector a -> Vector a
update

-- | \(O(n)\). Concatenate two vectors.
(><) :: Vector a -> Vector a -> Vector a
>< :: forall a. Vector a -> Vector a -> Vector a
(><) = forall a. a -> a
Exts.inline forall b a. (b -> a -> b) -> b -> Vector a -> b
foldl' forall a. Vector a -> a -> Vector a
snoc

-- | Check the invariant of the vector
invariant :: Vector a -> Bool
invariant :: forall a. Vector a -> Bool
invariant Vector a
_vec = Bool
True

-- | \(O(n)\). Create a vector from a list.
fromList :: [a] -> Vector a
fromList :: forall a. [a] -> Vector a
fromList = forall a. Stream Identity a -> Vector a
unstream forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => [a] -> Stream m a
Stream.fromList

keyBits :: Int
#ifdef TEST
keyBits = 2
#else
keyBits :: Int
keyBits = Int
5
#endif

nodeWidth :: Int
nodeWidth :: Int
nodeWidth = Int
1 forall a. Bits a => a -> Int -> a
!<<. Int
keyBits

keyMask :: Int
keyMask :: Int
keyMask = Int
nodeWidth forall a. Num a => a -> a -> a
- Int
1

(!<<.) :: Bits a => a -> Int -> a
!<<. :: forall a. Bits a => a -> Int -> a
(!<<.) = forall a. Bits a => a -> Int -> a
unsafeShiftL
{-# INLINE (!<<.) #-}

(!>>.) :: Bits a => a -> Int -> a
!>>. :: forall a. Bits a => a -> Int -> a
(!>>.) = forall a. Bits a => a -> Int -> a
unsafeShiftR
{-# INLINE (!>>.) #-}

infixl 8 !<<., !>>.

unstream :: Stream Identity a -> Vector a
unstream :: forall a. Stream Identity a -> Vector a
unstream Stream Identity a
stream = forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
  forall (m :: * -> *) a.
PrimMonad m =>
Stream Identity a -> m (Int, SmallArray a, [Node a])
streamToContents Stream Identity a
stream forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    (Int
size, SmallArray a
tail, [Node a
tree]) ->
      forall (f :: * -> *) a. Applicative f => a -> f a
pure RootNode {Int
size :: Int
$sel:size:RootNode :: Int
size, $sel:shift:RootNode :: Int
shift = Int
keyBits, SmallArray a
tail :: SmallArray a
$sel:tail:RootNode :: SmallArray a
tail, $sel:init:RootNode :: Array (Node a)
init = forall (f :: * -> *) a. Applicative f => a -> f a
pure Node a
tree}
    (Int
size, SmallArray a
tail, [Node a]
ls') -> do
      let iterateNodes :: Int -> [Node a] -> ST s (Vector a)
iterateNodes !Int
shift [Node a]
trees =
            forall {m :: * -> *} {t :: * -> *} {a}.
(PrimMonad m, Foldable t) =>
t (Node a) -> m [Node a]
nodes (forall a. [a] -> [a]
Prelude.reverse [Node a]
trees) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
              [Node a
tree] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure RootNode {Int
size :: Int
$sel:size:RootNode :: Int
size, Int
shift :: Int
$sel:shift:RootNode :: Int
shift, SmallArray a
tail :: SmallArray a
$sel:tail:RootNode :: SmallArray a
tail, $sel:init:RootNode :: Array (Node a)
init = forall a. Node a -> Array (Node a)
getInternalNode Node a
tree}
              [Node a]
trees' -> Int -> [Node a] -> ST s (Vector a)
iterateNodes (Int
shift forall a. Num a => a -> a -> a
+ Int
keyBits) [Node a]
trees'
      Int -> [Node a] -> ST s (Vector a)
iterateNodes Int
keyBits [Node a]
ls'
  where
    nodes :: t (Node a) -> m [Node a]
nodes t (Node a)
trees = do
      Buffer (PrimState m) (Node a)
buffer <- forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Int -> m (Buffer s a)
Buffer.newWithCapacity Int
nodeWidth
      (Buffer (PrimState m) (Node a)
buffer, [Node a]
acc) <-
        forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
Foldable.foldlM
          ( \(!Buffer (PrimState m) (Node a)
buffer, [Node a]
acc) Node a
t ->
              if forall s a. Buffer s a -> Int
Buffer.length Buffer (PrimState m) (Node a)
buffer forall a. Eq a => a -> a -> Bool
== Int
nodeWidth
                then do
                  SmallArray (Node a)
result <- forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Buffer s a -> m (SmallArray a)
Buffer.freeze Buffer (PrimState m) (Node a)
buffer
                  Buffer (PrimState m) (Node a)
buffer <- forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
a -> Buffer s a -> m (Buffer s a)
Buffer.push Node a
t forall a b. (a -> b) -> a -> b
$ forall s a. Buffer s a -> Buffer s a
Buffer.clear Buffer (PrimState m) (Node a)
buffer
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure (Buffer (PrimState m) (Node a)
buffer, forall a. Array (Node a) -> Node a
InternalNode SmallArray (Node a)
result forall a. a -> [a] -> [a]
: [Node a]
acc)
                else do
                  Buffer (PrimState m) (Node a)
buffer <- forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
a -> Buffer s a -> m (Buffer s a)
Buffer.push Node a
t Buffer (PrimState m) (Node a)
buffer
                  forall (f :: * -> *) a. Applicative f => a -> f a
pure (Buffer (PrimState m) (Node a)
buffer, [Node a]
acc)
          )
          (Buffer (PrimState m) (Node a)
buffer, [])
          t (Node a)
trees
      SmallArray (Node a)
final <- forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Buffer s a -> m (SmallArray a)
Buffer.unsafeFreeze Buffer (PrimState m) (Node a)
buffer
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. Array (Node a) -> Node a
InternalNode SmallArray (Node a)
final forall a. a -> [a] -> [a]
: [Node a]
acc
{-# INLINE unstream #-}

streamToContents :: PrimMonad m => Stream Identity a -> m (Int, SmallArray a, [Node a])
streamToContents :: forall (m :: * -> *) a.
PrimMonad m =>
Stream Identity a -> m (Int, SmallArray a, [Node a])
streamToContents (Stream s -> Identity (Step s a)
step s
s) = do
  Buffer (PrimState m) a
buffer <- forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Int -> m (Buffer s a)
Buffer.newWithCapacity Int
nodeWidth
  Int
-> Buffer (PrimState m) a
-> [Node a]
-> s
-> m (Int, SmallArray a, [Node a])
loop (Int
0 :: Int) Buffer (PrimState m) a
buffer [] s
s
  where
    loop :: Int
-> Buffer (PrimState m) a
-> [Node a]
-> s
-> m (Int, SmallArray a, [Node a])
loop !Int
size !Buffer (PrimState m) a
buffer [Node a]
acc s
s = do
      case forall a. Identity a -> a
runIdentity forall a b. (a -> b) -> a -> b
$ s -> Identity (Step s a)
step s
s of
        Stream.Yield a
x s
s' -> do
          if forall s a. Buffer s a -> Int
Buffer.length Buffer (PrimState m) a
buffer forall a. Eq a => a -> a -> Bool
== Int
nodeWidth
            then do
              SmallArray a
result <- forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Buffer s a -> m (SmallArray a)
Buffer.freeze Buffer (PrimState m) a
buffer
              Buffer (PrimState m) a
buffer <- forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
a -> Buffer s a -> m (Buffer s a)
Buffer.push a
x forall a b. (a -> b) -> a -> b
$ forall s a. Buffer s a -> Buffer s a
Buffer.clear Buffer (PrimState m) a
buffer
              Int
-> Buffer (PrimState m) a
-> [Node a]
-> s
-> m (Int, SmallArray a, [Node a])
loop (Int
size forall a. Num a => a -> a -> a
+ Int
1) Buffer (PrimState m) a
buffer (forall a. Array a -> Node a
DataNode SmallArray a
result forall a. a -> [a] -> [a]
: [Node a]
acc) s
s'
            else do
              Buffer (PrimState m) a
buffer <- forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
a -> Buffer s a -> m (Buffer s a)
Buffer.push a
x Buffer (PrimState m) a
buffer
              Int
-> Buffer (PrimState m) a
-> [Node a]
-> s
-> m (Int, SmallArray a, [Node a])
loop (Int
size forall a. Num a => a -> a -> a
+ Int
1) Buffer (PrimState m) a
buffer [Node a]
acc s
s'
        Stream.Skip s
s' -> Int
-> Buffer (PrimState m) a
-> [Node a]
-> s
-> m (Int, SmallArray a, [Node a])
loop Int
size Buffer (PrimState m) a
buffer [Node a]
acc s
s'
        Step s a
Stream.Done -> do
          SmallArray a
tail <- forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Buffer s a -> m (SmallArray a)
Buffer.unsafeFreeze Buffer (PrimState m) a
buffer
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
size, SmallArray a
tail, [Node a]
acc)
{-# INLINE streamToContents #-}

streamL :: Monad m => Vector a -> Stream m a
streamL :: forall (m :: * -> *) a. Monad m => Vector a -> Stream m a
streamL RootNode {Array (Node a)
init :: Array (Node a)
$sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail} = forall (m :: * -> *) a s. (s -> m (Step s a)) -> s -> Stream m a
Stream forall {f :: * -> *} {a}.
Applicative f =>
[(Node a, Int)] -> f (Step [(Node a, Int)] a)
step [(forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
init, Int
0 :: Int), (forall a. Array a -> Node a
DataNode Array a
tail, Int
0)]
  where
    step :: [(Node a, Int)] -> f (Step [(Node a, Int)] a)
step [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall s a. Step s a
Stream.Done
    step ((Node a
n, Int
i) : [(Node a, Int)]
rest) = case Node a
n of
      InternalNode Array (Node a)
ns
        | Int
i forall a. Ord a => a -> a -> Bool
< forall a. SmallArray a -> Int
sizeofSmallArray Array (Node a)
ns -> do
            let !(# Node a
ns' #) = forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## Array (Node a)
ns Int
i
                !i' :: Int
i' = Int
i forall a. Num a => a -> a -> a
+ Int
1
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Stream.Skip forall a b. (a -> b) -> a -> b
$ (Node a
ns', Int
0) forall a. a -> [a] -> [a]
: (Node a
n, Int
i') forall a. a -> [a] -> [a]
: [(Node a, Int)]
rest
        | Bool
otherwise -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Stream.Skip [(Node a, Int)]
rest
      DataNode Array a
xs
        | Int
i forall a. Ord a => a -> a -> Bool
< forall a. SmallArray a -> Int
sizeofSmallArray Array a
xs -> do
            let !(# a
x #) = forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## Array a
xs Int
i
                !i' :: Int
i' = Int
i forall a. Num a => a -> a -> a
+ Int
1
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a s. a -> s -> Step s a
Stream.Yield a
x forall a b. (a -> b) -> a -> b
$ (Node a
n, Int
i') forall a. a -> [a] -> [a]
: [(Node a, Int)]
rest
        | Bool
otherwise -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Stream.Skip [(Node a, Int)]
rest
    {-# INLINE step #-}
{-# INLINE streamL #-}

streamR :: Monad m => Vector a -> Stream m a
streamR :: forall (m :: * -> *) a. Monad m => Vector a -> Stream m a
streamR RootNode {Array (Node a)
init :: Array (Node a)
$sel:init:RootNode :: forall a. Vector a -> Array (Node a)
init, Array a
tail :: Array a
$sel:tail:RootNode :: forall a. Vector a -> Array a
tail} = forall (m :: * -> *) a s. (s -> m (Step s a)) -> s -> Stream m a
Stream forall {f :: * -> *} {a}.
Applicative f =>
[(Node a, Int)] -> f (Step [(Node a, Int)] a)
step [(forall a. Array a -> Node a
DataNode Array a
tail, Int
tailSize), (forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
init, Int
initSize)]
  where
    !tailSize :: Int
tailSize = forall a. SmallArray a -> Int
sizeofSmallArray Array a
tail forall a. Num a => a -> a -> a
- Int
1
    !initSize :: Int
initSize = forall a. SmallArray a -> Int
sizeofSmallArray Array (Node a)
init forall a. Num a => a -> a -> a
- Int
1

    step :: [(Node a, Int)] -> f (Step [(Node a, Int)] a)
step [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall s a. Step s a
Stream.Done
    step ((Node a
n, Int
i) : [(Node a, Int)]
rest) = case Node a
n of
      InternalNode Array (Node a)
ns
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Stream.Skip [(Node a, Int)]
rest
        | Bool
otherwise -> do
            let !(# Node a
n' #) = forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## Array (Node a)
ns Int
i
                !i' :: Int
i' = Int
i forall a. Num a => a -> a -> a
- Int
1
            case Node a
n' of
              InternalNode Array (Node a)
ns -> do
                let !z :: Int
z = forall a. SmallArray a -> Int
sizeofSmallArray Array (Node a)
ns forall a. Num a => a -> a -> a
- Int
1
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Stream.Skip forall a b. (a -> b) -> a -> b
$ (Node a
n', Int
z) forall a. a -> [a] -> [a]
: (Node a
n, Int
i') forall a. a -> [a] -> [a]
: [(Node a, Int)]
rest
              DataNode Array a
xs -> do
                let !z :: Int
z = forall a. SmallArray a -> Int
sizeofSmallArray Array a
xs forall a. Num a => a -> a -> a
- Int
1
                forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Stream.Skip forall a b. (a -> b) -> a -> b
$ (Node a
n', Int
z) forall a. a -> [a] -> [a]
: (Node a
n, Int
i') forall a. a -> [a] -> [a]
: [(Node a, Int)]
rest
      DataNode Array a
xs
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Stream.Skip [(Node a, Int)]
rest
        | Bool
otherwise -> do
            let !(# a
x #) = forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## Array a
xs Int
i
                !i' :: Int
i' = Int
i forall a. Num a => a -> a -> a
- Int
1
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a s. a -> s -> Step s a
Stream.Yield a
x forall a b. (a -> b) -> a -> b
$ (Node a
n, Int
i') forall a. a -> [a] -> [a]
: [(Node a, Int)]
rest
    {-# INLINE step #-}
{-# INLINE streamR #-}

istreamL :: Monad m => Vector a -> Stream m (Int, a)
istreamL :: forall (m :: * -> *) a. Monad m => Vector a -> Stream m (Int, a)
istreamL = forall (m :: * -> *) a. Monad m => Stream m a -> Stream m (Int, a)
Stream.indexed forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Vector a -> Stream m a
streamL
{-# INLINE istreamL #-}

istreamR :: Monad m => Vector a -> Stream m (Int, a)
istreamR :: forall (m :: * -> *) a. Monad m => Vector a -> Stream m (Int, a)
istreamR Vector a
vec = forall (m :: * -> *) a.
Monad m =>
Int -> Stream m a -> Stream m (Int, a)
Stream.indexedR (forall a. Vector a -> Int
length Vector a
vec forall a. Num a => a -> a -> a
- Int
1) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => Vector a -> Stream m a
streamR Vector a
vec
{-# INLINE istreamR #-}