{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedSums #-}
{-# 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_,
  )
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_ :: MArray (PrimState m) a -> Int -> m (MArray (PrimState m) a)
shrinkSmallMutableArray_ MArray (PrimState m) a
marr Int
n = MArray (PrimState m) a -> Int -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m ()
Primitive.shrinkSmallMutableArray MArray (PrimState m) a
marr Int
n m () -> MArray (PrimState m) a -> m (MArray (PrimState m) a)
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# :: (a -> (# b #)) -> SmallArray a -> SmallArray b
mapSmallArray# a -> (# b #)
f SmallArray a
sa = Int
-> b
-> (forall s. SmallMutableArray s b -> ST s ())
-> SmallArray b
forall a.
Int
-> a
-> (forall s. SmallMutableArray s a -> ST s ())
-> SmallArray a
createSmallArray (SmallArray a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SmallArray a
sa) ([Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"mapSmallArray#") ((forall s. SmallMutableArray s b -> ST s ()) -> SmallArray b)
-> (forall s. SmallMutableArray s b -> ST s ()) -> SmallArray b
forall a b. (a -> b) -> a -> b
$ \SmallMutableArray s b
smb -> do
  let go :: Int -> ST s ()
go Int
i =
        Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< SmallArray a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SmallArray a
sa) (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          a
x <- SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Monad m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
sa Int
i
          let !(# b
y #) = a -> (# b #)
f a
x
          SmallMutableArray (PrimState (ST s)) b -> Int -> b -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s b
SmallMutableArray (PrimState (ST s)) b
smb Int
i b
y ST s () -> ST s () -> ST s ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Int -> ST s ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> ST s ()
go Int
0
{-# INLINE mapSmallArray# #-}

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

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

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

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

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

updateResizeSmallArray :: Array a -> Int -> a -> Array a
updateResizeSmallArray :: Array a -> Int -> a -> Array a
updateResizeSmallArray Array a
arr Int
i a
a = (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> Array a)
-> (forall s. ST s (SmallMutableArray s a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
  SmallMutableArray s a
marr <- Array a
-> Int -> Int -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray Array a
arr Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
len (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
  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
i a
a
  SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
marr
  where
    len :: Int
len = Array a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray Array a
arr
{-# INLINE updateResizeSmallArray #-}

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

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

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

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

newtype STA a = STA {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 :: Int -> STA a -> SmallArray a
runSTA !Int
sz = \(STA forall s. SmallMutableArray# s a -> ST s (SmallArray a)
m) ->
  (forall s. ST s (SmallArray a)) -> SmallArray a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (SmallArray a)) -> SmallArray a)
-> (forall s. ST s (SmallArray a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$
    Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
sz a
forall a. a
undefinedElem
      ST s (SmallMutableArray s a)
-> (SmallMutableArray s a -> ST s (SmallArray a))
-> ST s (SmallArray a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \(SmallMutableArray SmallMutableArray# s a
ar#) -> SmallMutableArray# s a -> ST s (SmallArray a)
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 :: 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 = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr
      go :: Int -> Int -> f (STA b)
go Int
i Int
k
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len =
            STA b -> f (STA b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (STA b -> f (STA b)) -> STA b -> f (STA b)
forall a b. (a -> b) -> a -> b
$ (forall s. SmallMutableArray# s b -> ST s (SmallArray b)) -> STA b
forall a.
(forall s. SmallMutableArray# s a -> ST s (SmallArray a)) -> STA a
STA ((forall s. SmallMutableArray# s b -> ST s (SmallArray b))
 -> STA b)
-> (forall s. SmallMutableArray# s b -> ST s (SmallArray b))
-> STA b
forall a b. (a -> b) -> a -> b
$ \SmallMutableArray# s b
marr -> SmallMutableArray (PrimState (ST s)) b -> ST s (SmallArray b)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray (SmallMutableArray# s b -> SmallMutableArray s b
forall s a. SmallMutableArray# s a -> SmallMutableArray s a
SmallMutableArray SmallMutableArray# s b
marr)
        | (# a
x #) <- SmallArray a -> Int -> (# a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i =
            (b -> STA b -> STA b) -> f b -> f (STA b) -> f (STA b)
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 s. SmallMutableArray# s b -> ST s (SmallArray b)) -> STA b
forall a.
(forall s. SmallMutableArray# s a -> ST s (SmallArray a)) -> STA a
STA ((forall s. SmallMutableArray# s b -> ST s (SmallArray b))
 -> STA b)
-> (forall s. SmallMutableArray# s b -> ST s (SmallArray b))
-> STA b
forall a b. (a -> b) -> a -> b
$ \SmallMutableArray# s b
marr -> SmallMutableArray (PrimState (ST s)) b -> Int -> b -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray (SmallMutableArray# s b -> SmallMutableArray s b
forall s a. SmallMutableArray# s a -> SmallMutableArray s a
SmallMutableArray SmallMutableArray# s b
marr) Int
i b
b ST s () -> ST s (SmallArray b) -> ST s (SmallArray b)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> SmallMutableArray# s b -> ST s (SmallArray 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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> f (STA b)) -> Int -> f (STA b)
forall a b. (a -> b) -> a -> b
$! Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
step)
  if Int
len Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    then SmallArray b -> f (SmallArray b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallArray b
forall a. SmallArray a
emptySmallArray
    else Int -> STA b -> SmallArray b
forall a. Int -> STA a -> SmallArray a
runSTA Int
len (STA b -> SmallArray b) -> f (STA b) -> f (SmallArray b)
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 :: 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 = SmallArray a -> Int
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = SmallMutableArray (PrimState m) b -> m (SmallArray b)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray (PrimState m) b
marr
        | Bool
otherwise = do
            a
a <- SmallArray a -> Int -> m 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
            SmallMutableArray (PrimState m) b -> Int -> b -> m ()
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> SmallMutableArray (PrimState m) b -> m (SmallArray b))
-> Int -> SmallMutableArray (PrimState m) b -> m (SmallArray b)
forall a b. (a -> b) -> a -> b
$! Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
step) SmallMutableArray (PrimState m) b
marr
  SmallMutableArray (PrimState m) b
marr <- Int -> b -> m (SmallMutableArray (PrimState m) b)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len b
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 #-}