{-# 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
  { Vector a -> Int
size :: !Int,
    -- | 1 << 'shift' is the maximum that each child can contain
    Vector a -> Int
shift :: !Int,
    Vector a -> Array (Node a)
init :: !(Array (Node a)),
    Vector a -> Array a
tail :: !(Array a)
  }

instance Show1 Vector where
  liftShowsPrec :: (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) -> String -> Int -> [a] -> ShowS
forall a. (Int -> a -> ShowS) -> String -> Int -> a -> ShowS
showsUnaryWith ((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) String
"fromList" Int
p (Vector a -> [a]
forall a. Vector 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 Eq a => Eq (Vector a) where
  == :: Vector a -> Vector a -> Bool
(==) = 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 = Vector a -> Vector a -> Ordering
forall a. Ord a => Vector a -> Vector a -> Ordering
persistentVectorCompare
  {-# INLINE compare #-}

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

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

instance Traversable Vector where
  traverse :: (a -> f b) -> Vector a -> f (Vector b)
traverse = (a -> f b) -> Vector a -> f (Vector b)
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
(<>) = Vector a -> Vector a -> Vector a
forall a. Vector a -> Vector a -> Vector a
(><)
  {-# INLINE (<>) #-}

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

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

instance Applicative Vector where
  pure :: a -> Vector a
pure = a -> Vector a
forall a. a -> Vector a
singleton
  {-# INLINE pure #-}
  Vector (a -> b)
fs <*> :: Vector (a -> b) -> Vector a -> Vector b
<*> Vector a
xs = (Vector b -> (a -> b) -> Vector b)
-> Vector b -> Vector (a -> b) -> Vector b
forall b a. (b -> a -> b) -> b -> Vector 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

instance Monad Vector where
  Vector a
xs >>= :: Vector a -> (a -> Vector b) -> Vector b
>>= a -> Vector b
f = (Vector b -> a -> Vector b) -> Vector b -> Vector a -> Vector b
forall b a. (b -> a -> b) -> b -> Vector 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
  {-# INLINE (>>=) #-}

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

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

instance MonadPlus Vector

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

data Node a
  = InternalNode {Node a -> Array (Node a)
getInternalNode :: !(Array (Node a))}
  | DataNode {Node a -> Array a
getDataNode :: !(Array a)}
  deriving (Int -> Node a -> ShowS
[Node a] -> ShowS
Node a -> String
(Int -> Node a -> ShowS)
-> (Node a -> String) -> ([Node a] -> ShowS) -> Show (Node a)
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
(==) = 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 = Node a -> Node a -> Ordering
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 = [Item (Vector a)] -> Vector a
forall a. [a] -> Vector a
Data.Vector.Persistent.Internal.fromList
  {-# INLINE fromList #-}
  toList :: Vector a -> [Item (Vector a)]
toList = Vector a -> [Item (Vector a)]
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 :: (a -> b -> b) -> b -> Vector a -> b
foldr a -> b -> b
f b
z = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> (Vector a -> Identity b) -> Vector a -> b
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. (a -> b -> b) -> b -> Stream Identity a -> Identity b
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
Stream.foldr a -> b -> b
f b
z (Stream Identity a -> Identity b)
-> (Vector a -> Stream Identity a) -> Vector a -> Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Stream Identity a
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' :: (a -> b -> b) -> b -> Vector a -> b
foldr' a -> b -> b
f b
z = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> (Vector a -> Identity b) -> Vector a -> b
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. (b -> a -> b) -> b -> Stream Identity a -> Identity b
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> a) -> a -> Stream m b -> m a
Stream.foldl' ((a -> b -> b) -> b -> a -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> b -> b
f) b
z (Stream Identity a -> Identity b)
-> (Vector a -> Stream Identity a) -> Vector a -> Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Stream Identity a
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 :: (b -> a -> b) -> b -> Vector a -> b
foldl b -> a -> b
f b
z = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> (Vector a -> Identity b) -> Vector a -> b
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. (a -> b -> b) -> b -> Stream Identity a -> Identity b
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
Stream.foldr ((b -> a -> b) -> a -> b -> b
forall a b c. (a -> b -> c) -> b -> a -> c
flip b -> a -> b
f) b
z (Stream Identity a -> Identity b)
-> (Vector a -> Stream Identity a) -> Vector a -> Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Stream Identity a
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' :: (b -> a -> b) -> b -> Vector a -> b
foldl' b -> a -> b
f b
z = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> (Vector a -> Identity b) -> Vector a -> b
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. (b -> a -> b) -> b -> Stream Identity a -> Identity b
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> a) -> a -> Stream m b -> m a
Stream.foldl' b -> a -> b
f b
z (Stream Identity a -> Identity b)
-> (Vector a -> Stream Identity a) -> Vector a -> Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Stream Identity a
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 :: (Int -> a -> b -> b) -> b -> Vector a -> b
ifoldr Int -> a -> b -> b
f b
z = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> (Vector a -> Identity b) -> Vector a -> b
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. ((Int, a) -> b -> b) -> b -> Stream Identity (Int, a) -> Identity b
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
Stream.foldr ((Int -> a -> b -> b) -> (Int, a) -> b -> b
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> a -> b -> b
f) b
z (Stream Identity (Int, a) -> Identity b)
-> (Vector a -> Stream Identity (Int, a)) -> Vector a -> Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Stream Identity (Int, a)
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 :: (Int -> b -> a -> b) -> b -> Vector a -> b
ifoldl Int -> b -> a -> b
f b
z = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> (Vector a -> Identity b) -> Vector a -> b
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. ((Int, a) -> b -> b) -> b -> Stream Identity (Int, a) -> Identity b
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 (Stream Identity (Int, a) -> Identity b)
-> (Vector a -> Stream Identity (Int, a)) -> Vector a -> Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Stream Identity (Int, a)
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' :: (Int -> a -> b -> b) -> b -> Vector a -> b
ifoldr' Int -> a -> b -> b
f b
z = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> (Vector a -> Identity b) -> Vector a -> b
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. (b -> (Int, a) -> b) -> b -> Stream Identity (Int, a) -> Identity b
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 (Stream Identity (Int, a) -> Identity b)
-> (Vector a -> Stream Identity (Int, a)) -> Vector a -> Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Stream Identity (Int, a)
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' :: (Int -> b -> a -> b) -> b -> Vector a -> b
ifoldl' Int -> b -> a -> b
f b
z = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> (Vector a -> Identity b) -> Vector a -> b
forall b c a. Coercible b c => (b -> c) -> (a -> b) -> a -> c
#. (b -> (Int, a) -> b) -> b -> Stream Identity (Int, a) -> Identity b
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 (Stream Identity (Int, a) -> Identity b)
-> (Vector a -> Stream Identity (Int, a)) -> Vector a -> Identity b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Stream Identity (Int, a)
forall (m :: * -> *) a. Monad m => Vector a -> Stream m (Int, a)
istreamL
{-# INLINE ifoldl' #-}

persistentVectorEq :: Eq a => Vector a -> Vector a -> Bool
persistentVectorEq :: 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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
size' Bool -> Bool -> Bool
&& (Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| (Int
shift Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
shift' Bool -> Bool -> Bool
&& Array a
tail Array a -> Array a -> Bool
forall a. Eq a => a -> a -> Bool
== Array a
tail' Bool -> Bool -> Bool
&& Array (Node a)
init Array (Node a) -> Array (Node a) -> Bool
forall a. Eq a => a -> a -> Bool
== Array (Node a)
init'))
{-# INLINEABLE persistentVectorEq #-}

nodeEq :: Eq a => Node a -> Node a -> Bool
nodeEq :: Node a -> Node a -> Bool
nodeEq (InternalNode Array (Node a)
ns) (InternalNode Array (Node a)
ns') = Array (Node a)
ns Array (Node a) -> Array (Node a) -> Bool
forall a. Eq a => a -> a -> Bool
== Array (Node a)
ns'
nodeEq (DataNode Array a
as) (DataNode Array a
as') = Array a
as Array a -> Array a -> Bool
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 :: 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'} =
    Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
size Int
size'
      Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> if Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then Ordering
EQ
        else Array (Node a) -> Array (Node a) -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Array (Node a)
init Array (Node a)
init' Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Array a -> Array a -> Ordering
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 :: Node a -> Node a -> Ordering
nodeCompare (DataNode Array a
as) (DataNode Array a
as') = Array a -> Array a -> Ordering
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') = Array (Node a) -> Array (Node a) -> Ordering
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 :: a -> Vector a
singleton a
a = RootNode :: forall a. Int -> Int -> Array (Node a) -> Array a -> Vector a
RootNode {$sel:size:RootNode :: Int
size = Int
1, $sel:shift:RootNode :: Int
shift = Int
keyBits, $sel:tail:RootNode :: Array a
tail = a -> Array a
forall a. a -> Array a
singletonSmallArray a
a, $sel:init:RootNode :: Array (Node a)
init = Array (Node a)
forall a. SmallArray a
emptySmallArray}
{-# INLINE singleton #-}

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

-- | \(O(1)\) Return 'True' if the vector is empty, 'False' otherwise.
null :: Vector a -> Bool
null :: Vector a -> Bool
null Vector a
xs = Vector a -> Int
forall a. Vector a -> Int
length Vector a
xs Int -> Int -> Bool
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
|> :: Vector a -> a -> Vector 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:|> :: Vector a -> a -> Vector a
$m:|> :: forall r a. Vector a -> (Vector a -> a -> r) -> (Void# -> r) -> r
:|> x <-
  (unsnoc -> Just (vec, x))
  where
    Vector a
vec :|> a
x = Vector a
vec Vector a -> a -> Vector a
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 :: Vector a
$mEmpty :: forall r a. Vector a -> (Void# -> r) -> (Void# -> r) -> r
Empty <-
  (null -> True)
  where
    Empty = Vector a
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 :: 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 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyMask) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 =
      Vector a
vec
        { $sel:tail:RootNode :: Array a
tail = Array a -> Int -> a -> Array a
forall a. Array a -> Int -> a -> Array a
updateResizeSmallArray Array a
tail (Int
size Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyMask) a
a,
          $sel:size:RootNode :: Int
size = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        }
  | Bool
otherwise = Vector a -> Int -> Array a -> Vector a
forall a. Vector a -> Int -> Array a -> Vector a
snocArr Vector a
vec Int
1 (Array a -> Vector a) -> Array a -> Vector a
forall a b. (a -> b) -> a -> b
$ a -> Array a
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 :: 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 Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyMask) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 =
      Vector a
vec
        { $sel:tail:RootNode :: Array a
tail =
            -- update the array in place
            (forall s. ST s (Array a)) -> Array a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array a)) -> Array a)
-> (forall s. ST s (Array a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
              SmallMutableArray s a
marr <- Array a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> m (SmallMutableArray (PrimState m) a)
unsafeThawSmallArray Array a
tail
              SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr (Int
size Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyMask) a
a
              SmallMutableArray (PrimState (ST s)) a -> ST s (Array a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr,
          $sel:size:RootNode :: Int
size = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
        }
  | Bool
otherwise = Vector a -> Int -> Array a -> Vector a
forall a. Vector a -> Int -> Array a -> Vector a
snocArr Vector a
vec Int
1 (Array a -> Vector a) -> Array a -> Vector a
forall a b. (a -> b) -> a -> b
$ a -> Array a
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 :: 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
  | Vector a -> Bool
forall a. Vector a -> Bool
null Vector a
vec =
      RootNode :: forall a. Int -> Int -> Array (Node a) -> Array a -> Vector a
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 = Array (Node a)
forall a. SmallArray a
emptySmallArray
        }
  | Int
size Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!>>. Int
keyBits Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!<<. Int
shift =
      RootNode :: forall a. Int -> Int -> Array (Node a) -> Array a -> Vector a
RootNode
        { $sel:size:RootNode :: Int
size = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
addedSize,
          $sel:shift:RootNode :: Int
shift = Int
shift Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
keyBits,
          $sel:init:RootNode :: Array (Node a)
init =
            let !path :: Node a
path = Int -> Array a -> Node a
forall a. Int -> Array a -> Node a
newPath Int
shift Array a
tail
                !internal :: Node a
internal = Array (Node a) -> Node a
forall a. Array (Node a) -> Node a
InternalNode (Array (Node a) -> Node a) -> Array (Node a) -> Node a
forall a b. (a -> b) -> a -> b
$ Vector a -> Array (Node a)
forall a. Vector a -> Array (Node a)
init Vector a
vec
             in Node a -> Node a -> Array (Node a)
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 :: forall a. Int -> Int -> Array (Node a) -> Array a -> Vector a
RootNode
        { $sel:size:RootNode :: Int
size = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
addedSize,
          Int
shift :: Int
$sel:shift:RootNode :: Int
shift,
          $sel:init:RootNode :: Array (Node a)
init = Int -> Array a -> Int -> Array (Node a) -> Array (Node a)
forall a. Int -> Array a -> Int -> Array (Node a) -> Array (Node a)
snocTail Int
size Array a
tail Int
shift (Array (Node a) -> Array (Node a))
-> Array (Node a) -> Array (Node a)
forall a b. (a -> b) -> a -> b
$ Vector a -> Array (Node a)
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 :: 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 = Array (Node a) -> Int -> Node a -> Array (Node a)
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
keyBits = Array a -> Node a
forall a. Array a -> Node a
DataNode Array a
tail
          | Int
subIx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Array (Node a) -> Int
forall a. SmallArray a -> Int
sizeofSmallArray Array (Node a)
parent =
              let vec' :: Node a
vec' = Array (Node a) -> Int -> Node a
forall a. SmallArray a -> Int -> a
indexSmallArray Array (Node a)
parent Int
subIx
               in Array (Node a) -> Node a
forall a. Array (Node a) -> Node a
InternalNode (Array (Node a) -> Node a) -> Array (Node a) -> Node a
forall a b. (a -> b) -> a -> b
$ Int -> Array (Node a) -> Array (Node a)
go (Int
level Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
keyBits) (Node a -> Array (Node a)
forall a. Node a -> Array (Node a)
getInternalNode Node a
vec')
          | Bool
otherwise = Int -> Array a -> Node a
forall a. Int -> Array a -> Node a
newPath (Int
level Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
keyBits) Array a
tail
        subIx :: Int
subIx = ((Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!>>. Int
level) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyMask
{-# INLINE snocTail #-}

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

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

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

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

adjust# :: (a -> (# a #)) -> Int -> Vector a -> Vector a
adjust# :: (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).
  | (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix :: Word) Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
size = Vector a
vec
  | Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector a -> Int
forall a. Vector a -> Int
tailOffset Vector a
vec = Vector a
vec {$sel:tail:RootNode :: Array a
tail = Array a -> Int -> (a -> (# a #)) -> Array a
forall a. Array a -> Int -> (a -> (# a #)) -> Array a
modifySmallArray# Array a
tail (Int
ix Int -> Int -> Int
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 (Vector a -> Array (Node a)
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
keyBits,
        let !node :: Node a
node = Array a -> Node a
forall a. Array a -> Node a
DataNode (Array a -> Node a) -> Array a -> Node a
forall a b. (a -> b) -> a -> b
$ Array a -> Int -> (a -> (# a #)) -> Array a
forall a. Array a -> Int -> (a -> (# a #)) -> Array a
modifySmallArray# (Node a -> Array a
forall a. Node a -> Array a
getDataNode Node a
vec') (Int
ix Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyMask) a -> (# a #)
f =
          Array (Node a) -> Int -> Node a -> Array (Node a)
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
keyBits) (Node a -> Array (Node a)
forall a. Node a -> Array (Node a)
getInternalNode Node a
vec') =
          Array (Node a) -> Int -> Node a -> Array (Node a)
forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
vec Int
ix' (Node a -> Array (Node a)) -> Node a -> Array (Node a)
forall a b. (a -> b) -> a -> b
$! Array (Node a) -> Node a
forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
node
      where
        ix' :: Int
ix' = (Int
ix Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!>>. Int
level) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyBits
        vec' :: Node a
vec' = Array (Node a) -> Int -> Node a
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 :: (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).
  | (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix :: Word) Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
size = Vector a -> f (Vector a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector a
vec
  | Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector a -> Int
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}) (Array a -> Vector a) -> f (Array a) -> f (Vector a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array a -> Int -> (a -> f a) -> f (Array a)
forall (f :: * -> *) a.
Functor f =>
Array a -> Int -> (a -> f a) -> f (Array a)
modifySmallArrayF Array a
tail (Int
ix Int -> Int -> Int
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}) (Array (Node a) -> Vector a) -> f (Array (Node a)) -> f (Vector a)
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 (Vector a -> Array (Node a)
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
keyBits =
          (\Array a
node' -> Array (Node a) -> Int -> Node a -> Array (Node a)
forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
vec Int
ix' (Node a -> Array (Node a)) -> Node a -> Array (Node a)
forall a b. (a -> b) -> a -> b
$! Array a -> Node a
forall a. Array a -> Node a
DataNode Array a
node')
            (Array a -> Array (Node a)) -> f (Array a) -> f (Array (Node a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array a -> Int -> (a -> f a) -> f (Array a)
forall (f :: * -> *) a.
Functor f =>
Array a -> Int -> (a -> f a) -> f (Array a)
modifySmallArrayF (Node a -> Array a
forall a. Node a -> Array a
getDataNode Node a
vec') (Int
ix Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyMask) a -> f a
f
      | Bool
otherwise =
          (\Array (Node a)
node -> Array (Node a) -> Int -> Node a -> Array (Node a)
forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
vec Int
ix' (Node a -> Array (Node a)) -> Node a -> Array (Node a)
forall a b. (a -> b) -> a -> b
$! Array (Node a) -> Node a
forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
node)
            (Array (Node a) -> Array (Node a))
-> f (Array (Node a)) -> f (Array (Node a))
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
keyBits) (Node a -> Array (Node a)
forall a. Node a -> Array (Node a)
getInternalNode Node a
vec')
      where
        ix' :: Int
ix' = (Int
ix Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!>>. Int
level) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyBits
        vec' :: Node a
vec' = Array (Node a) -> Int -> Node a
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 :: 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).
  | (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
ix :: Word) Word -> Word -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
size = Vector a
vec
  | Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Vector a -> Int
forall a. Vector a -> Int
tailOffset Vector a
vec = Vector a
vec {$sel:tail:RootNode :: Array a
tail = Array a -> Int -> a -> Array a
forall a. Array a -> Int -> a -> Array a
updateSmallArray Array a
tail (Int
ix Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyMask) a
x}
  | Bool
otherwise = Vector a
vec {$sel:init:RootNode :: Array (Node a)
init = Int -> a -> Int -> Array (Node a) -> Array (Node a)
forall a. Int -> a -> Int -> Array (Node a) -> Array (Node a)
go Int
ix a
x Int
shift (Vector a -> Array (Node a)
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
keyBits =
          let !node :: Node a
node = Array a -> Node a
forall a. Array a -> Node a
DataNode (Array a -> Node a) -> Array a -> Node a
forall a b. (a -> b) -> a -> b
$ Array a -> Int -> a -> Array a
forall a. Array a -> Int -> a -> Array a
updateSmallArray (Node a -> Array a
forall a. Node a -> Array a
getDataNode Node a
vec') (Int
ix Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyMask) a
x
           in Array (Node a) -> Int -> Node a -> Array (Node a)
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
keyBits) (Node a -> Array (Node a)
forall a. Node a -> Array (Node a)
getInternalNode Node a
vec')
           in Array (Node a) -> Int -> Node a -> Array (Node a)
forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
vec Int
ix' (Node a -> Array (Node a)) -> Node a -> Array (Node a)
forall a b. (a -> b) -> a -> b
$! Array (Node a) -> Node a
forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
node
      where
        ix' :: Int
ix' = (Int
ix Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!>>. Int
level) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
keyMask
        vec' :: Node a
vec' = Array (Node a) -> Int -> Node a
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 :: 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 = Maybe (Vector a, a)
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 #) <- Array a -> Int -> (# a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## Array a
tail Int
0 = (Vector a, a) -> Maybe (Vector a, a)
forall a. a -> Maybe a
Just (Vector a
forall a. Vector a
empty, a
x)
  | Array a -> Bool
forall a. SmallArray a -> Bool
nullSmallArray Array a
tail',
    (# Array (Node a)
init', Array a
tail' #) <- Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
forall a.
Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
unsnocTail# Int
size Int
shift Array (Node a)
init =
      (Vector a, a) -> Maybe (Vector a, a)
forall a. a -> Maybe a
Just (Vector a
vec {$sel:size:RootNode :: Int
size = Int
size Int -> Int -> Int
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 = Array a
tail'}, a
a)
  | Bool
otherwise = (Vector a, a) -> Maybe (Vector a, a)
forall a. a -> Maybe a
Just (Vector a
vec {$sel:size:RootNode :: Int
size = Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, $sel:tail:RootNode :: Array a
tail = Array a
tail'}, a
a)
  where
    a :: a
a = Array a -> a
forall a. SmallArray a -> a
lastSmallArray Array a
tail
    tail' :: Array a
tail' = Array a -> Array a
forall a. Array a -> Array a
popSmallArray Array a
tail
{-# INLINEABLE unsnoc #-}

unsnocTail# :: Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
unsnocTail# :: Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
unsnocTail# = Int -> Int -> Array (Node a) -> (# Array (Node a), Array a #)
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
keyBits = (# Array (Node a) -> Array (Node a)
forall a. Array a -> Array a
popSmallArray Array (Node a)
parent, Node a -> Array a
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
keyBits) (Node a -> Array (Node a)
forall a. Node a -> Array (Node a)
getInternalNode Node a
child)
          if Array (Node a) -> Bool
forall a. SmallArray a -> Bool
nullSmallArray Array (Node a)
child'
            then (# Array (Node a) -> Array (Node a)
forall a. Array a -> Array a
popSmallArray Array (Node a)
parent, Array a
tail #)
            else (# Array (Node a) -> Int -> Node a -> Array (Node a)
forall a. Array a -> Int -> a -> Array a
updateSmallArray Array (Node a)
parent Int
subIx (Node a -> Array (Node a)) -> Node a -> Array (Node a)
forall a b. (a -> b) -> a -> b
$ Array (Node a) -> Node a
forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
child', Array a
tail #)
      where
        child :: Node a
child = Array (Node a) -> Int -> Node a
forall a. SmallArray a -> Int -> a
indexSmallArray Array (Node a)
parent Int
subIx
        -- we need to subtract 2 because the first subtraction gets us to the tail element
        -- the second subtraction gets to the last element in the tree
        subIx :: Int
subIx = ((Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!>>. Int
level) Int -> Int -> Int
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 :: Vector a -> Int
tailOffset Vector a
vec = (Vector a -> Int
forall a. Vector a -> Int
length Vector a
vec Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. ((-Int
1) Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!<<. Int
keyBits)
{-# INLINE tailOffset #-}

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

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

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

toList :: Vector a -> [a]
toList :: Vector a -> [a]
toList = Stream Identity a -> [a]
forall a. Stream Identity a -> [a]
pureStreamToList (Stream Identity a -> [a])
-> (Vector a -> Stream Identity a) -> Vector a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Stream Identity a
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 :: Stream Identity a -> [a]
pureStreamToList Stream Identity a
s = (forall b. (a -> b -> b) -> b -> b) -> [a]
forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
Exts.build (\a -> b -> b
c b
n -> Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> Identity b -> b
forall a b. (a -> b) -> a -> b
$ (a -> b -> b) -> b -> Stream Identity a -> Identity 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 :: (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 = (a -> b) -> Array a -> Array b
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 = (Node a -> Node b) -> Array (Node a) -> Array (Node b)
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) = Array b -> Node b
forall a. Array a -> Node a
DataNode (Array b -> Node b) -> Array b -> Node b
forall a b. (a -> b) -> a -> b
$ (a -> b) -> Array a -> Array 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) = Array (Node b) -> Node b
forall a. Array (Node a) -> Node a
InternalNode (Array (Node b) -> Node b) -> Array (Node b) -> Node b
forall a b. (a -> b) -> a -> b
$ (Node a -> Node b) -> Array (Node a) -> Array (Node 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 :: (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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Vector b
forall a. Vector a
empty
  | Bool
otherwise =
      Vector a
vec
        { $sel:init:RootNode :: Array (Node b)
init = Int
-> Int
-> (Int -> Node a -> Node b)
-> Array (Node a)
-> Array (Node b)
forall a b.
Int -> Int -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray Int
0 (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!<<. Int
shift) (Int -> Int -> Node a -> Node b
go (Int -> Int -> Node a -> Node b) -> Int -> Int -> Node a -> Node b
forall a b. (a -> b) -> a -> b
$! Int
shift Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
keyBits) Array (Node a)
init,
          $sel:tail:RootNode :: Array b
tail = Int -> Int -> (Int -> a -> b) -> Array a -> Array b
forall a b.
Int -> Int -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray (Vector a -> Int
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) = Array b -> Node b
forall a. Array a -> Node a
DataNode (Array b -> Node b) -> Array b -> Node b
forall a b. (a -> b) -> a -> b
$ Int -> Int -> (Int -> a -> b) -> Array a -> Array 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) = Array (Node b) -> Node b
forall a. Array (Node a) -> Node a
InternalNode (Array (Node b) -> Node b) -> Array (Node b) -> Node b
forall a b. (a -> b) -> a -> b
$ Int
-> Int
-> (Int -> Node a -> Node b)
-> Array (Node a)
-> Array (Node b)
forall a b.
Int -> Int -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray Int
i0 (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!<<. Int
shift) (Int -> Int -> Node a -> Node b
go (Int -> Int -> Node a -> Node b) -> Int -> Int -> Node a -> Node b
forall a b. (a -> b) -> a -> b
$! Int
shift Int -> Int -> Int
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 :: (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} =
  (Array (Node b) -> Array b -> Vector b)
-> f (Array (Node b)) -> f (Array b) -> f (Vector b)
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})
    ((Node a -> f (Node b)) -> Array (Node a) -> f (Array (Node 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)
init)
    ((a -> f b) -> Array a -> f (Array 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
tail)
  where
    go :: Node a -> f (Node b)
go (DataNode Array a
as) = Array b -> Node b
forall a. Array a -> Node a
DataNode (Array b -> Node b) -> f (Array b) -> f (Node b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Array a -> f (Array 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) = Array (Node b) -> Node b
forall a. Array (Node a) -> Node a
InternalNode (Array (Node b) -> Node b) -> f (Array (Node b)) -> f (Node b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Node a -> f (Node b)) -> Array (Node a) -> f (Array (Node 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 :: (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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Vector b -> f (Vector b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Vector b
forall a. Vector a
empty
  | Bool
otherwise =
      (Array (Node b) -> Array b -> Vector b)
-> f (Array (Node b)) -> f (Array b) -> f (Vector b)
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})
        (Int
-> Int
-> (Int -> Node a -> f (Node b))
-> Array (Node a)
-> f (Array (Node b))
forall (f :: * -> *) a b.
Applicative f =>
Int -> Int -> (Int -> a -> f b) -> SmallArray a -> f (SmallArray b)
itraverseStepSmallArray Int
0 (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!<<. Int
shift) (Int -> Int -> Node a -> f (Node b)
go (Int -> Int -> Node a -> f (Node b))
-> Int -> Int -> Node a -> f (Node b)
forall a b. (a -> b) -> a -> b
$! Int
shift Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
keyBits) Array (Node a)
init)
        (Int -> Int -> (Int -> a -> f b) -> Array a -> f (Array b)
forall (f :: * -> *) a b.
Applicative f =>
Int -> Int -> (Int -> a -> f b) -> SmallArray a -> f (SmallArray b)
itraverseStepSmallArray (Vector a -> Int
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) = Array b -> Node b
forall a. Array a -> Node a
DataNode (Array b -> Node b) -> f (Array b) -> f (Node b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> (Int -> a -> f b) -> Array a -> f (Array 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) = Array (Node b) -> Node b
forall a. Array (Node a) -> Node a
InternalNode (Array (Node b) -> Node b) -> f (Array (Node b)) -> f (Node b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int
-> Int
-> (Int -> Node a -> f (Node b))
-> Array (Node a)
-> f (Array (Node b))
forall (f :: * -> *) a b.
Applicative f =>
Int -> Int -> (Int -> a -> f b) -> SmallArray a -> f (SmallArray b)
itraverseStepSmallArray Int
i0 (Int
1 Int -> Int -> Int
forall a. Bits a => a -> Int -> a
!<<. Int
shift) (Int -> Int -> Node a -> f (Node b)
go (Int -> Int -> Node a -> f (Node b))
-> Int -> Int -> Node a -> f (Node b)
forall a b. (a -> b) -> a -> b
$! Int
shift Int -> Int -> Int
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
// :: Vector a -> [(Int, a)] -> Vector a
(//) = ((Vector a -> (Int, a) -> Vector a)
 -> Vector a -> [(Int, a)] -> Vector a)
-> (Vector a -> (Int, a) -> Vector a)
-> Vector a
-> [(Int, a)]
-> Vector a
forall a. a -> a
Exts.inline (Vector a -> (Int, a) -> Vector a)
-> Vector a -> [(Int, a)] -> Vector a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
Foldable.foldl' ((Vector a -> (Int, a) -> Vector a)
 -> Vector a -> [(Int, a)] -> Vector a)
-> (Vector a -> (Int, a) -> Vector a)
-> Vector a
-> [(Int, a)]
-> Vector a
forall a b. (a -> b) -> a -> b
$ ((Int, a) -> Vector a -> Vector a)
-> Vector a -> (Int, a) -> Vector a
forall a b c. (a -> b -> c) -> b -> a -> c
flip (((Int, a) -> Vector a -> Vector a)
 -> Vector a -> (Int, a) -> Vector a)
-> ((Int, a) -> Vector a -> Vector a)
-> Vector a
-> (Int, a)
-> Vector a
forall a b. (a -> b) -> a -> b
$ (Int -> a -> Vector a -> Vector a)
-> (Int, a) -> Vector a -> Vector a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Int -> a -> Vector a -> Vector a
forall a. Int -> a -> Vector a -> Vector a
update

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

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

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

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

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

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

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

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

infixl 8 !<<., !>>.

unstream :: Stream Identity a -> Vector a
unstream :: Stream Identity a -> Vector a
unstream Stream Identity a
stream = (forall s. ST s (Vector a)) -> Vector a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector a)) -> Vector a)
-> (forall s. ST s (Vector a)) -> Vector a
forall a b. (a -> b) -> a -> b
$ do
  Stream Identity a -> ST s (Int, SmallArray a, [Node a])
forall (m :: * -> *) a.
PrimMonad m =>
Stream Identity a -> m (Int, SmallArray a, [Node a])
streamToContents Stream Identity a
stream ST s (Int, SmallArray a, [Node a])
-> ((Int, SmallArray a, [Node a]) -> ST s (Vector a))
-> ST s (Vector a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    (Int
size, SmallArray a
tail, [Node a
tree]) ->
      Vector a -> ST s (Vector a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure RootNode :: forall a. Int -> Int -> Array (Node a) -> Array a -> Vector a
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 = Node a -> Array (Node a)
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 =
            [Node a] -> ST s [Node a]
forall (m :: * -> *) (t :: * -> *) a.
(PrimMonad m, Foldable t) =>
t (Node a) -> m [Node a]
nodes ([Node a] -> [Node a]
forall a. [a] -> [a]
Prelude.reverse [Node a]
trees) ST s [Node a] -> ([Node a] -> ST s (Vector a)) -> ST s (Vector a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
              [Node a
tree] -> Vector a -> ST s (Vector a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure RootNode :: forall a. Int -> Int -> Array (Node a) -> Array a -> Vector a
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 = Node a -> Array (Node a)
forall a. Node a -> Array (Node a)
getInternalNode Node a
tree}
              [Node a]
trees' -> Int -> [Node a] -> ST s (Vector a)
iterateNodes (Int
shift Int -> Int -> Int
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 <- Int -> m (Buffer (PrimState m) (Node a))
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) <-
        ((Buffer (PrimState m) (Node a), [Node a])
 -> Node a -> m (Buffer (PrimState m) (Node a), [Node a]))
-> (Buffer (PrimState m) (Node a), [Node a])
-> t (Node a)
-> m (Buffer (PrimState m) (Node a), [Node a])
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 Buffer (PrimState m) (Node a) -> Int
forall s a. Buffer s a -> Int
Buffer.length Buffer (PrimState m) (Node a)
buffer Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nodeWidth
                then do
                  SmallArray (Node a)
result <- Buffer (PrimState m) (Node a) -> m (SmallArray (Node a))
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 <- Node a
-> Buffer (PrimState m) (Node a)
-> m (Buffer (PrimState m) (Node a))
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)
 -> m (Buffer (PrimState m) (Node a)))
-> Buffer (PrimState m) (Node a)
-> m (Buffer (PrimState m) (Node a))
forall a b. (a -> b) -> a -> b
$ Buffer (PrimState m) (Node a) -> Buffer (PrimState m) (Node a)
forall s a. Buffer s a -> Buffer s a
Buffer.clear Buffer (PrimState m) (Node a)
buffer
                  (Buffer (PrimState m) (Node a), [Node a])
-> m (Buffer (PrimState m) (Node a), [Node a])
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Buffer (PrimState m) (Node a)
buffer, SmallArray (Node a) -> Node a
forall a. Array (Node a) -> Node a
InternalNode SmallArray (Node a)
result Node a -> [Node a] -> [Node a]
forall a. a -> [a] -> [a]
: [Node a]
acc)
                else do
                  Buffer (PrimState m) (Node a)
buffer <- Node a
-> Buffer (PrimState m) (Node a)
-> m (Buffer (PrimState m) (Node a))
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
                  (Buffer (PrimState m) (Node a), [Node a])
-> m (Buffer (PrimState m) (Node a), [Node a])
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 <- Buffer (PrimState m) (Node a) -> m (SmallArray (Node a))
forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Buffer s a -> m (SmallArray a)
Buffer.unsafeFreeze Buffer (PrimState m) (Node a)
buffer
      [Node a] -> m [Node a]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Node a] -> m [Node a]) -> [Node a] -> m [Node a]
forall a b. (a -> b) -> a -> b
$ SmallArray (Node a) -> Node a
forall a. Array (Node a) -> Node a
InternalNode SmallArray (Node a)
final Node a -> [Node a] -> [Node a]
forall a. a -> [a] -> [a]
: [Node a]
acc
{-# INLINE unstream #-}

streamToContents :: PrimMonad m => Stream Identity a -> m (Int, SmallArray a, [Node a])
streamToContents :: 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 <- Int -> m (Buffer (PrimState m) a)
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 Identity (Step s a) -> Step s a
forall a. Identity a -> a
runIdentity (Identity (Step s a) -> Step s a)
-> Identity (Step s a) -> Step s a
forall a b. (a -> b) -> a -> b
$ s -> Identity (Step s a)
step s
s of
        Stream.Yield a
x s
s' -> do
          if Buffer (PrimState m) a -> Int
forall s a. Buffer s a -> Int
Buffer.length Buffer (PrimState m) a
buffer Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nodeWidth
            then do
              SmallArray a
result <- Buffer (PrimState m) a -> m (SmallArray a)
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 <- a -> Buffer (PrimState m) a -> m (Buffer (PrimState m) a)
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 -> m (Buffer (PrimState m) a))
-> Buffer (PrimState m) a -> m (Buffer (PrimState m) a)
forall a b. (a -> b) -> a -> b
$ Buffer (PrimState m) a -> Buffer (PrimState m) a
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Buffer (PrimState m) a
buffer (SmallArray a -> Node a
forall a. Array a -> Node a
DataNode SmallArray a
result Node a -> [Node a] -> [Node a]
forall a. a -> [a] -> [a]
: [Node a]
acc) s
s'
            else do
              Buffer (PrimState m) a
buffer <- a -> Buffer (PrimState m) a -> m (Buffer (PrimState m) a)
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 Int -> Int -> Int
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 <- Buffer (PrimState m) a -> m (SmallArray a)
forall (m :: * -> *) s a.
(PrimMonad m, s ~ PrimState m) =>
Buffer s a -> m (SmallArray a)
Buffer.unsafeFreeze Buffer (PrimState m) a
buffer
          (Int, SmallArray a, [Node a]) -> m (Int, SmallArray a, [Node a])
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 :: 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} = ([(Node a, Int)] -> m (Step [(Node a, Int)] a))
-> [(Node a, Int)] -> Stream m a
forall (m :: * -> *) a s. (s -> m (Step s a)) -> s -> Stream m a
Stream [(Node a, Int)] -> m (Step [(Node a, Int)] a)
forall (f :: * -> *) a.
Applicative f =>
[(Node a, Int)] -> f (Step [(Node a, Int)] a)
step [(Array (Node a) -> Node a
forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
init, Int
0 :: Int), (Array a -> Node a
forall a. Array a -> Node a
DataNode Array a
tail, Int
0)]
  where
    step :: [(Node a, Int)] -> f (Step [(Node a, Int)] a)
step [] = Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Step [(Node a, Int)] a
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Array (Node a) -> Int
forall a. SmallArray a -> Int
sizeofSmallArray Array (Node a)
ns -> Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a))
-> Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a)
forall a b. (a -> b) -> a -> b
$ [(Node a, Int)] -> Step [(Node a, Int)] a
forall s a. s -> Step s a
Stream.Skip [(Node a, Int)]
rest
        | Bool
otherwise -> do
            let !(# Node a
ns' #) = Array (Node a) -> Int -> (# Node a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## Array (Node a)
ns Int
i
                !i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a))
-> Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a)
forall a b. (a -> b) -> a -> b
$ [(Node a, Int)] -> Step [(Node a, Int)] a
forall s a. s -> Step s a
Stream.Skip ([(Node a, Int)] -> Step [(Node a, Int)] a)
-> [(Node a, Int)] -> Step [(Node a, Int)] a
forall a b. (a -> b) -> a -> b
$ (Node a
ns', Int
0) (Node a, Int) -> [(Node a, Int)] -> [(Node a, Int)]
forall a. a -> [a] -> [a]
: (Node a
n, Int
i') (Node a, Int) -> [(Node a, Int)] -> [(Node a, Int)]
forall a. a -> [a] -> [a]
: [(Node a, Int)]
rest
      DataNode Array a
xs
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Array a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray Array a
xs -> Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a))
-> Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a)
forall a b. (a -> b) -> a -> b
$ [(Node a, Int)] -> Step [(Node a, Int)] a
forall s a. s -> Step s a
Stream.Skip [(Node a, Int)]
rest
        | Bool
otherwise -> do
            let !(# a
x #) = Array a -> Int -> (# a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## Array a
xs Int
i
                !i' :: Int
i' = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a))
-> Step [(Node a, Int)] a -> f (Step [(Node a, Int)] a)
forall a b. (a -> b) -> a -> b
$ a -> [(Node a, Int)] -> Step [(Node a, Int)] a
forall a s. a -> s -> Step s a
Stream.Yield a
x ([(Node a, Int)] -> Step [(Node a, Int)] a)
-> [(Node a, Int)] -> Step [(Node a, Int)] a
forall a b. (a -> b) -> a -> b
$ (Node a
n, Int
i') (Node a, Int) -> [(Node a, Int)] -> [(Node a, Int)]
forall a. a -> [a] -> [a]
: [(Node a, Int)]
rest
    {-# INLINE step #-}
{-# INLINE streamL #-}

streamR :: Monad m => Vector a -> Stream m a
streamR :: 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} = ([(Node a, Int)] -> m (Step [(Node a, Int)] a))
-> [(Node a, Int)] -> Stream m a
forall (m :: * -> *) a s. (s -> m (Step s a)) -> s -> Stream m a
Stream [(Node a, Int)] -> m (Step [(Node a, Int)] a)
forall (f :: * -> *) a.
Applicative f =>
[(Node a, Int)] -> f (Step [(Node a, Int)] a)
step [(Array a -> Node a
forall a. Array a -> Node a
DataNode Array a
tail, Int
tailSize), (Array (Node a) -> Node a
forall a. Array (Node a) -> Node a
InternalNode Array (Node a)
init, Int
initSize)]
  where
    !tailSize :: Int
tailSize = Array a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray Array a
tail Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    !initSize :: Int
initSize = Array (Node a) -> Int
forall a. SmallArray a -> Int
sizeofSmallArray Array (Node a)
init Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1

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

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

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