{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

module Data.Vector.Persistent.Internal.Array
  ( Array,
    MArray,
    nullSmallArray,
    lastSmallArray#,
    singletonSmallArray,
    twoSmallArray,
    updateSmallArray,
    modifySmallArray,
    modifySmallArrayF,
    modifySmallArray',
    updateResizeSmallArray,
    popSmallArray,
    undefinedElem,
    ifoldrStepSmallArray,
    ifoldlStepSmallArray,
    ifoldrStepSmallArray',
    ifoldlStepSmallArray',
    imapStepSmallArray,
    imapStepSmallArray',
    itraverseStepSmallArray,
    modifySmallArray#,
    mapSmallArray#,
    shrinkSmallMutableArray_,
    snocSmallArray,
    unsnocSmallArray,
  )
where

import Control.Applicative (liftA2)
import Control.Monad (when)
import Control.Monad.Primitive (PrimMonad, PrimState)
import Control.Monad.ST (ST, runST)
import Data.Coerce (coerce)
import Data.Functor (($>))
import Data.Functor.Identity (Identity (..))
import qualified Data.Primitive as Primitive
import Data.Primitive.SmallArray
import GHC.Exts (SmallMutableArray#)

type Array = SmallArray

type MArray = SmallMutableArray

-- | Used to support older ghcs.
shrinkSmallMutableArray_ :: PrimMonad m => MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a)
#if __GLASGOW_HASKELL__ >= 810
shrinkSmallMutableArray_ :: forall (m :: * -> *) a.
PrimMonad m =>
MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a)
shrinkSmallMutableArray_ MArray (PrimState m) a
marr Int
n = forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m ()
Primitive.shrinkSmallMutableArray MArray (PrimState m) a
marr Int
n forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> MArray (PrimState m) a
marr
#else
shrinkSmallMutableArray_ mary n = Primitive.cloneSmallMutableArray mary 0 n
#endif 
{-# INLINE shrinkSmallMutableArray_ #-}

mapSmallArray# :: (a -> (# b #)) -> SmallArray a -> SmallArray b
mapSmallArray# :: forall a b. (a -> (# b #)) -> SmallArray a -> SmallArray b
mapSmallArray# a -> (# b #)
f SmallArray a
sa = forall a.
Int
-> a
-> (forall s. SmallMutableArray s a -> ST s ())
-> SmallArray a
createSmallArray (forall (t :: * -> *) a. Foldable t => t a -> Int
length SmallArray a
sa) (forall a. HasCallStack => [Char] -> a
error [Char]
"mapSmallArray#") forall a b. (a -> b) -> a -> b
$ \SmallMutableArray s b
smb -> do
  let go :: Int -> ST s ()
go Int
i =
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i forall a. Ord a => a -> a -> Bool
< forall (t :: * -> *) a. Foldable t => t a -> Int
length SmallArray a
sa) forall a b. (a -> b) -> a -> b
$ do
          a
x <- forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
sa Int
i
          let !(# b
y #) = a -> (# b #)
f a
x
          forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s b
smb Int
i b
y forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> ST s ()
go (Int
i forall a. Num a => a -> a -> a
+ Int
1)
  Int -> ST s ()
go Int
0
{-# INLINE mapSmallArray# #-}

nullSmallArray :: SmallArray a -> Bool
nullSmallArray :: forall a. SmallArray a -> Bool
nullSmallArray SmallArray a
arr = forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr forall a. Eq a => a -> a -> Bool
== Int
0
{-# INLINE nullSmallArray #-}

unsnocSmallArray :: SmallArray a -> Maybe (SmallArray a, a)
unsnocSmallArray :: forall a. SmallArray a -> Maybe (SmallArray a, a)
unsnocSmallArray SmallArray a
arr = do
  let len :: Int
len = forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
len forall a. Eq a => a -> a -> Bool
== Int
0) forall a. Maybe a
Nothing
  a
x <- forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
arr (Int
len forall a. Num a => a -> a -> a
- Int
1)
  let !arr' :: SmallArray a
arr' = forall a. SmallArray a -> Int -> Int -> SmallArray a
cloneSmallArray SmallArray a
arr Int
0 (Int
len forall a. Num a => a -> a -> a
- Int
1)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (SmallArray a
arr', a
x)
{-# INLINE unsnocSmallArray #-}

lastSmallArray# :: SmallArray a -> (# a #)
lastSmallArray# :: forall a. SmallArray a -> (# a #)
lastSmallArray# SmallArray a
arr = forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr forall a b. (a -> b) -> a -> b
$ forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr

singletonSmallArray :: a -> Array a
singletonSmallArray :: forall a. a -> Array a
singletonSmallArray a
a = forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
1 a
a
{-# INLINE singletonSmallArray #-}

twoSmallArray :: a -> a -> Array a
twoSmallArray :: forall a. a -> a -> Array a
twoSmallArray a
x a
y = forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
  SmallMutableArray s a
marr <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
2 a
x
  forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
marr Int
1 a
y
  forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
marr
{-# INLINE twoSmallArray #-}

updateSmallArray :: Array a -> Int -> a -> Array a
updateSmallArray :: forall a. Array a -> Int -> a -> Array a
updateSmallArray Array a
arr Int
i a
x = forall a. Array a -> Int -> (a -> (# a #)) -> Array a
modifySmallArray# Array a
arr Int
i forall a b. (a -> b) -> a -> b
$ \a
_ -> (# a
x #)
{-# INLINE updateSmallArray #-}

modifySmallArray :: Array a -> Int -> (a -> a) -> Array a
modifySmallArray :: forall a. Array a -> Int -> (a -> a) -> Array a
modifySmallArray Array a
arr Int
i a -> a
f = forall a. Array a -> Int -> (a -> (# a #)) -> Array a
modifySmallArray# Array a
arr Int
i forall a b. (a -> b) -> a -> b
$ \a
x -> (# a -> a
f a
x #)
{-# INLINE modifySmallArray #-}

modifySmallArrayF :: Functor f => Array a -> Int -> (a -> f a) -> f (Array a)
modifySmallArrayF :: forall (f :: * -> *) a.
Functor f =>
Array a -> Int -> (a -> f a) -> f (Array a)
modifySmallArrayF Array a
arr Int
i a -> f a
f | (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## Array a
arr Int
i = forall a. Array a -> Int -> a -> Array a
updateSmallArray Array a
arr Int
i forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f a
x
{-# INLINE modifySmallArrayF #-}

modifySmallArray' :: Array a -> Int -> (a -> a) -> Array a
modifySmallArray' :: forall a. Array a -> Int -> (a -> a) -> Array a
modifySmallArray' Array a
arr Int
i a -> a
f = forall a. Array a -> Int -> (a -> (# a #)) -> Array a
modifySmallArray# Array a
arr Int
i forall a b. (a -> b) -> a -> b
$ \a
x -> let !x' :: a
x' = a -> a
f a
x in (# a
x' #)
{-# INLINE modifySmallArray' #-}

modifySmallArray# :: Array a -> Int -> (a -> (# a #)) -> Array a
modifySmallArray# :: forall a. Array a -> Int -> (a -> (# a #)) -> Array a
modifySmallArray# Array a
arr Int
i a -> (# a #)
f = forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
  SmallMutableArray s a
marr <- forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray Array a
arr Int
0 forall a b. (a -> b) -> a -> b
$ forall a. SmallArray a -> Int
sizeofSmallArray Array a
arr
  a
x <- forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
indexSmallArrayM Array a
arr Int
i
  let !(# a
x' #) = a -> (# a #)
f a
x
  forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
marr Int
i a
x'
  forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
marr
{-# INLINE modifySmallArray# #-}

snocSmallArray :: Array a -> a -> Array a
snocSmallArray :: forall a. Array a -> a -> Array a
snocSmallArray Array a
arr a
x = forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
  SmallMutableArray s a
marr <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray (forall a. SmallArray a -> Int
sizeofSmallArray Array a
arr forall a. Num a => a -> a -> a
+ Int
1) a
x
  forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
marr Int
0 Array a
arr Int
0 (forall a. SmallArray a -> Int
sizeofSmallArray Array a
arr)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
marr
{-# INLINE snocSmallArray #-}

-- this is wrong
updateResizeSmallArray :: Array a -> Int -> a -> Array a
updateResizeSmallArray :: forall a. Array a -> Int -> a -> Array a
updateResizeSmallArray Array a
arr Int
i a
a =
  if Int
i forall a. Eq a => a -> a -> Bool
== Int
len
    then forall a. Array a -> a -> Array a
snocSmallArray Array a
arr a
a
    else forall a. Array a -> Int -> a -> Array a
updateSmallArray Array a
arr Int
i a
a
  where
    len :: Int
len = forall a. SmallArray a -> Int
sizeofSmallArray Array a
arr
{-# INLINE updateResizeSmallArray #-}

popSmallArray :: Array a -> Array a
popSmallArray :: forall a. Array a -> Array a
popSmallArray Array a
arr = forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray Array a
arr Int
0 (forall a. SmallArray a -> Int
sizeofSmallArray Array a
arr forall a. Num a => a -> a -> a
- Int
1)
{-# INLINE popSmallArray #-}

undefinedElem :: forall a. a
undefinedElem :: forall a. a
undefinedElem = forall a. HasCallStack => [Char] -> a
error [Char]
"undefined element"
{-# NOINLINE undefinedElem #-}

ifoldrStepSmallArray :: Int -> Int -> (Int -> a -> b -> b) -> b -> SmallArray a -> b
ifoldrStepSmallArray :: forall a b.
Int -> Int -> (Int -> a -> b -> b) -> b -> SmallArray a -> b
ifoldrStepSmallArray Int
i0 Int
step Int -> a -> b -> b
f b
z SmallArray a
arr = do
  let len :: Int
len = forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr
      go :: Int -> Int -> b
go Int
i Int
j
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
len = b
z
        | (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = Int -> a -> b -> b
f Int
j a
x (Int -> Int -> b
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$! Int
j forall a. Num a => a -> a -> a
+ Int
step)
  Int -> Int -> b
go Int
0 Int
i0
{-# INLINE ifoldrStepSmallArray #-}

ifoldlStepSmallArray :: Int -> Int -> (Int -> b -> a -> b) -> b -> SmallArray a -> b
ifoldlStepSmallArray :: forall b a.
Int -> Int -> (Int -> b -> a -> b) -> b -> SmallArray a -> b
ifoldlStepSmallArray Int
i0 Int
step Int -> b -> a -> b
f b
z SmallArray a
arr = do
  let len :: Int
len = forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr
      go :: Int -> Int -> b
go Int
i Int
j
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = b
z
        | (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = Int -> b -> a -> b
f Int
j (Int -> Int -> b
go (Int
i forall a. Num a => a -> a -> a
- Int
1) forall a b. (a -> b) -> a -> b
$! Int
j forall a. Num a => a -> a -> a
- Int
step) a
x
  Int -> Int -> b
go (Int
len forall a. Num a => a -> a -> a
- Int
1) Int
i0
{-# INLINE ifoldlStepSmallArray #-}

ifoldrStepSmallArray' :: Int -> Int -> (Int -> a -> b -> b) -> b -> SmallArray a -> b
ifoldrStepSmallArray' :: forall a b.
Int -> Int -> (Int -> a -> b -> b) -> b -> SmallArray a -> b
ifoldrStepSmallArray' Int
i0 Int
step Int -> a -> b -> b
f b
z SmallArray a
arr = do
  let go :: Int -> Int -> b -> b
go Int
i Int
j b
acc
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
0 = b
acc
        | (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = (Int -> Int -> b -> b
go (Int
i forall a. Num a => a -> a -> a
- Int
1) forall a b. (a -> b) -> a -> b
$! (Int
j forall a. Num a => a -> a -> a
- Int
step)) forall a b. (a -> b) -> a -> b
$! Int -> a -> b -> b
f Int
j a
x b
acc
  Int -> Int -> b -> b
go (forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr) Int
i0 b
z
{-# INLINE ifoldrStepSmallArray' #-}

ifoldlStepSmallArray' :: Int -> Int -> (Int -> b -> a -> b) -> b -> SmallArray a -> b
ifoldlStepSmallArray' :: forall b a.
Int -> Int -> (Int -> b -> a -> b) -> b -> SmallArray a -> b
ifoldlStepSmallArray' Int
i0 Int
step Int -> b -> a -> b
f b
z SmallArray a
arr = do
  let go :: Int -> Int -> b -> b
go Int
i Int
j b
acc
        | Int
i forall a. Eq a => a -> a -> Bool
== forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr = b
acc
        | (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = (Int -> Int -> b -> b
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$! (Int
j forall a. Num a => a -> a -> a
+ Int
step)) forall a b. (a -> b) -> a -> b
$! Int -> b -> a -> b
f Int
j b
acc a
x
  Int -> Int -> b -> b
go Int
0 Int
i0 b
z
{-# INLINE ifoldlStepSmallArray' #-}

imapStepSmallArray :: Int -> Int -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray :: forall a b.
Int -> Int -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray Int
i0 Int
step Int -> a -> b
f SmallArray a
arr = forall a.
Int
-> a
-> (forall s. SmallMutableArray s a -> ST s ())
-> SmallArray a
createSmallArray Int
len forall a. a
undefinedElem forall a b. (a -> b) -> a -> b
$ \SmallMutableArray s b
marr -> do
  let go :: Int -> Int -> ST s ()
go Int
i Int
k = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i forall a. Ord a => a -> a -> Bool
< Int
len) forall a b. (a -> b) -> a -> b
$ do
        a
x <- forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
arr Int
i
        forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s b
marr Int
i (Int -> a -> b
f Int
k a
x)
        Int -> Int -> ST s ()
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$! Int
k forall a. Num a => a -> a -> a
+ Int
step
  Int -> Int -> ST s ()
go Int
0 Int
i0
  where
    len :: Int
len = forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr
{-# INLINE imapStepSmallArray #-}

imapStepSmallArray' :: Int -> (a -> Int) -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray' :: forall a b.
Int
-> (a -> Int) -> (Int -> a -> b) -> SmallArray a -> SmallArray b
imapStepSmallArray' Int
i0 a -> Int
step Int -> a -> b
f SmallArray a
arr = forall a.
Int
-> a
-> (forall s. SmallMutableArray s a -> ST s ())
-> SmallArray a
createSmallArray Int
len forall a. a
undefinedElem forall a b. (a -> b) -> a -> b
$ \SmallMutableArray s b
marr -> do
  let go :: Int -> Int -> ST s ()
go Int
i Int
k = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i forall a. Ord a => a -> a -> Bool
< Int
len) forall a b. (a -> b) -> a -> b
$ do
        a
x <- forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
arr Int
i
        forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s b
marr Int
i forall a b. (a -> b) -> a -> b
$! Int -> a -> b
f Int
k a
x
        Int -> Int -> ST s ()
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$! Int
k forall a. Num a => a -> a -> a
+ a -> Int
step a
x
  Int -> Int -> ST s ()
go Int
0 Int
i0
  where
    len :: Int
len = forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr
{-# INLINE imapStepSmallArray' #-}

newtype STA a = STA {forall a.
STA a -> forall s. SmallMutableArray# s a -> ST s (SmallArray a)
_runSTA :: forall s. SmallMutableArray# s a -> ST s (SmallArray a)}

runSTA :: Int -> STA a -> SmallArray a
runSTA :: forall a. Int -> STA a -> SmallArray a
runSTA !Int
sz = \(STA forall s. SmallMutableArray# s a -> ST s (SmallArray a)
m) ->
  forall a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
sz forall a. a
undefinedElem
      forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(SmallMutableArray SmallMutableArray# s a
ar#) -> forall s. SmallMutableArray# s a -> ST s (SmallArray a)
m SmallMutableArray# s a
ar#

itraverseStepSmallArray :: Applicative f => Int -> Int -> (Int -> a -> f b) -> SmallArray a -> f (SmallArray b)
itraverseStepSmallArray :: forall (f :: * -> *) a b.
Applicative f =>
Int -> Int -> (Int -> a -> f b) -> SmallArray a -> f (SmallArray b)
itraverseStepSmallArray Int
i0 Int
step Int -> a -> f b
f = \ !SmallArray a
arr -> do
  let len :: Int
len = forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr
      go :: Int -> Int -> f (STA b)
go Int
i Int
k
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
len =
            forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a.
(forall s. SmallMutableArray# s a -> ST s (SmallArray a)) -> STA a
STA forall a b. (a -> b) -> a -> b
$ \SmallMutableArray# s b
marr -> forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray (forall s a. SmallMutableArray# s a -> SmallMutableArray s a
SmallMutableArray SmallMutableArray# s b
marr)
        | (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i =
            forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2
              (\b
b (STA forall s. SmallMutableArray# s b -> ST s (SmallArray b)
m) -> forall a.
(forall s. SmallMutableArray# s a -> ST s (SmallArray a)) -> STA a
STA forall a b. (a -> b) -> a -> b
$ \SmallMutableArray# s b
marr -> forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray (forall s a. SmallMutableArray# s a -> SmallMutableArray s a
SmallMutableArray SmallMutableArray# s b
marr) Int
i b
b forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall s. SmallMutableArray# s b -> ST s (SmallArray b)
m SmallMutableArray# s b
marr)
              (Int -> a -> f b
f Int
k a
x)
              (Int -> Int -> f (STA b)
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$! Int
k forall a. Num a => a -> a -> a
+ Int
step)
  if Int
len forall a. Eq a => a -> a -> Bool
== Int
0
    then forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. SmallArray a
emptySmallArray
    else forall a. Int -> STA a -> SmallArray a
runSTA Int
len forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> f (STA b)
go Int
0 Int
i0
{-# INLINE [1] itraverseStepSmallArray #-}

{-# RULES
"itraverseStepSmallArray/ST" forall i0 step (f :: Int -> a -> ST s b).
  itraverseStepSmallArray i0 step f =
    itraverseStepSmallArrayP i0 step f
"itraverseStepSmallArray/IO" forall i0 step (f :: Int -> a -> IO b).
  itraverseStepSmallArray i0 step f =
    itraverseStepSmallArrayP i0 step f
"itraverseStepSmallArray/Id" forall i0 step (f :: Int -> a -> Identity b).
  itraverseStepSmallArray i0 step f =
    ( coerce ::
        (SmallArray a -> SmallArray (Identity b)) ->
        SmallArray a ->
        Identity (SmallArray b)
    )
      (imapStepSmallArray i0 step f)
  #-}

-- | This is the fastest, most straightforward way to traverse
-- an array, but it only works correctly with a sufficiently
-- "affine" 'PrimMonad' instance. In particular, it must only produce
-- /one/ result array. 'Control.Monad.Trans.List.ListT'-transformed
-- monads, for example, will not work right at all.
itraverseStepSmallArrayP :: PrimMonad m => Int -> Int -> (Int -> a -> m b) -> SmallArray a -> m (SmallArray b)
itraverseStepSmallArrayP :: forall (m :: * -> *) a b.
PrimMonad m =>
Int -> Int -> (Int -> a -> m b) -> SmallArray a -> m (SmallArray b)
itraverseStepSmallArrayP Int
i0 Int
step Int -> a -> m b
f = \ !SmallArray a
ary -> do
  let len :: Int
len = forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
ary
      go :: Int -> Int -> SmallMutableArray (PrimState m) b -> m (SmallArray b)
go Int
i Int
k SmallMutableArray (PrimState m) b
marr
        | Int
i forall a. Eq a => a -> a -> Bool
== Int
len = forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray (PrimState m) b
marr
        | Bool
otherwise = do
            a
a <- forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
ary Int
i
            b
b <- Int -> a -> m b
f Int
k a
a
            forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState m) b
marr Int
i b
b
            (Int -> Int -> SmallMutableArray (PrimState m) b -> m (SmallArray b)
go (Int
i forall a. Num a => a -> a -> a
+ Int
1) forall a b. (a -> b) -> a -> b
$! Int
k forall a. Num a => a -> a -> a
+ Int
step) SmallMutableArray (PrimState m) b
marr
  SmallMutableArray (PrimState m) b
marr <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len forall a. a
undefinedElem
  Int -> Int -> SmallMutableArray (PrimState m) b -> m (SmallArray b)
go Int
0 Int
i0 SmallMutableArray (PrimState m) b
marr
{-# INLINE itraverseStepSmallArrayP #-}