{-# 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
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 #-}
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)
#-}
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 #-}