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

-- | This is an internal module.
--
-- It provides a thin wrapper over "Data.Primitive.SmallArray"
-- with \(O(1)\) slicing.
--
-- __Warning:__ No bound checks are performed!

module Data.RRBVector.Internal.Array
    ( Array, MutableArray
    , ifoldrStep, ifoldlStep, ifoldrStep', ifoldlStep'
    , empty, singleton, from2, wrap
    , replicate, replicateSnoc
    , index, head, last
    , update, adjust, adjust'
    , take, drop, splitAt
    , snoc, cons, (++)
    , map, map'
    , imapStep, imapStep'
    , unzipWith
    , traverse, traverse'
    , itraverseStep, itraverseStep'
    , new, read, write
    , freeze, thaw
    ) where

#if !(MIN_VERSION_base(4,18,0))
import Control.Applicative (liftA2)
#endif
import Control.DeepSeq (NFData(..))
import Control.Monad (when)
import Control.Monad.ST
import Data.Foldable (Foldable(..))
import Data.Primitive.SmallArray
import Prelude hiding (replicate, take, drop, splitAt, head, last, map, traverse, read, unzip, (++))

-- start length array
data Array a = Array !Int !Int !(SmallArray a)
data MutableArray s a = MutableArray !Int !Int !(SmallMutableArray s a)

instance Foldable Array where
    foldr :: forall a b. (a -> b -> b) -> b -> Array a -> b
foldr a -> b -> b
f b
z (Array Int
start Int
len SmallArray a
arr) =
        let end :: Int
end = Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len
            go :: Int -> b
go Int
i
                | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
end = b
z
                | (# a
x #) <- SmallArray a -> Int -> (# a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = a -> b -> b
f a
x (Int -> b
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
        in Int -> b
go Int
start

    foldl :: forall b a. (b -> a -> b) -> b -> Array a -> b
foldl b -> a -> b
f b
z (Array Int
start Int
len SmallArray a
arr) =
        let go :: Int -> b
go Int
i
                | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start = b
z
                | (# a
x #) <- SmallArray a -> Int -> (# a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = b -> a -> b
f (Int -> b
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) a
x
        in Int -> b
go (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

    foldr' :: forall a b. (a -> b -> b) -> b -> Array a -> b
foldr' a -> b -> b
f b
z (Array Int
start Int
len SmallArray a
arr) =
        let go :: Int -> b -> b
go Int
i !b
acc
                | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start = b
acc
                | (# a
x #) <- SmallArray a -> Int -> (# a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = Int -> b -> b
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (a -> b -> b
f a
x b
acc)
        in Int -> b -> b
go (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) b
z

    foldl' :: forall b a. (b -> a -> b) -> b -> Array a -> b
foldl' b -> a -> b
f b
z (Array Int
start Int
len SmallArray a
arr) =
        let end :: Int
end = Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len
            go :: Int -> b -> b
go Int
i !b
acc
                | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
end = b
acc
                | (# a
x #) <- SmallArray a -> Int -> (# a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = Int -> b -> b
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (b -> a -> b
f b
acc a
x)
        in Int -> b -> b
go Int
start b
z

    null :: forall a. Array a -> Bool
null Array a
arr = Array a -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0

    length :: forall a. Array a -> Int
length (Array Int
_ Int
len SmallArray a
_) = Int
len

instance (NFData a) => NFData (Array a) where
    rnf :: Array a -> ()
rnf = (() -> a -> ()) -> () -> Array a -> ()
forall b a. (b -> a -> b) -> b -> Array a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\()
_ a
x -> a -> ()
forall a. NFData a => a -> ()
rnf a
x) ()

ifoldrStep :: Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
ifoldrStep :: forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
ifoldrStep Int
i0 a -> Int
step Int -> a -> b -> b
f b
z (Array Int
start Int
len SmallArray a
arr) =
    let end :: Int
end = Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len
        go :: Int -> Int -> b
go !Int
i !Int
j -- i is the index in arr, j is the index for f
            | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
end = 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
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int
step a
x))
    in Int -> Int -> b
go Int
start Int
i0

ifoldlStep :: Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
ifoldlStep :: forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
ifoldlStep Int
i0 a -> Int
step Int -> b -> a -> b
f b
z (Array Int
start Int
len SmallArray a
arr) =
    let go :: Int -> Int -> b
go !Int
i !Int
j -- i is the index in arr, j is the index for f
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start = 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
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- a -> Int
step a
x)) a
x
    in Int -> Int -> b
go (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
i0

ifoldrStep' :: Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
ifoldrStep' :: forall a b.
Int -> (a -> Int) -> (Int -> a -> b -> b) -> b -> Array a -> b
ifoldrStep' Int
i0 a -> Int
step Int -> a -> b -> b
f b
z (Array Int
start Int
len SmallArray a
arr) =
    let go :: Int -> Int -> b -> b
go !Int
i !Int
j !b
acc -- i is the index in arr, j is the index for f
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
start = 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
j Int -> Int -> Int
forall a. Num a => a -> a -> a
- a -> Int
step a
x) (Int -> a -> b -> b
f Int
j a
x b
acc)
    in Int -> Int -> b -> b
go (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int
i0 b
z

ifoldlStep' :: Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
ifoldlStep' :: forall a b.
Int -> (a -> Int) -> (Int -> b -> a -> b) -> b -> Array a -> b
ifoldlStep' Int
i0 a -> Int
step Int -> b -> a -> b
f b
z (Array Int
start Int
len SmallArray a
arr) =
    let end :: Int
end = Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len
        go :: Int -> Int -> b -> b
go !Int
i !Int
j !b
acc -- i is the index in arr, j is the index for f
            | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
end = 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
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int
step a
x) (Int -> b -> a -> b
f Int
j b
acc a
x)
    in Int -> Int -> b -> b
go Int
start Int
i0 b
z

uninitialized :: a
uninitialized :: forall a. a
uninitialized = [Char] -> a
forall a. [Char] -> a
errorWithoutStackTrace [Char]
"uninitialized"

empty :: Array a
empty :: forall a. Array a
empty = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
0 (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray (Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
0 a
forall a. a
uninitialized)

singleton :: a -> Array a
singleton :: forall a. a -> Array a
singleton a
x = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
1 (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray (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
x)

from2 :: a -> a -> Array a
from2 :: forall a. a -> a -> Array a
from2 a
x a
y = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
2 (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- 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
sma Int
1 a
y
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma

wrap :: SmallArray a -> Array a
wrap :: forall a. SmallArray a -> Array a
wrap SmallArray a
arr = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 (SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
arr) SmallArray a
arr

replicate :: Int -> a -> Array a
replicate :: forall a. Int -> a -> Array a
replicate Int
n a
x = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
n (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray (Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
n a
x)

-- > replicateSnoc n x y = snoc (replicate n x) y
replicateSnoc :: Int -> a -> a -> Array a
replicateSnoc :: forall a. Int -> a -> a -> Array a
replicateSnoc Int
n a
x a
y = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len 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
sma Int
n a
y
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma
  where
    len :: Int
len = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

index :: Array a -> Int -> a
index :: forall a. Array a -> Int -> a
index (Array Int
start Int
_ SmallArray a
arr) Int
idx = SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray a
arr (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx)

update :: Array a -> Int -> a -> Array a
update :: forall a. Array a -> Int -> a -> Array a
update (Array Int
start Int
len SmallArray a
sa) Int
idx a
x = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- SmallArray 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 SmallArray a
sa Int
start Int
len
    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
sma Int
idx a
x
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma

adjust :: Array a -> Int -> (a -> a) -> Array a
adjust :: forall a. Array a -> Int -> (a -> a) -> Array a
adjust (Array Int
start Int
len SmallArray a
sa) Int
idx a -> a
f = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- SmallArray 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 SmallArray a
sa Int
start Int
len
    a
x <- SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
sa (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx)
    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
sma Int
idx (a -> a
f a
x)
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma

adjust' :: Array a -> Int -> (a -> a) -> Array a
adjust' :: forall a. Array a -> Int -> (a -> a) -> Array a
adjust' (Array Int
start Int
len SmallArray a
sa) Int
idx a -> a
f = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- SmallArray 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 SmallArray a
sa Int
start Int
len
    a
x <- SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
sa (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx)
    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
sma Int
idx (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a -> a
f a
x
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma

take :: Array a -> Int -> Array a
take :: forall a. Array a -> Int -> Array a
take (Array Int
start Int
_ SmallArray a
arr) Int
n = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
start Int
n SmallArray a
arr

drop :: Array a -> Int -> Array a
drop :: forall a. Array a -> Int -> Array a
drop (Array Int
start Int
len SmallArray a
arr) Int
n = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n) SmallArray a
arr

splitAt :: Array a -> Int -> (Array a, Array a)
splitAt :: forall a. Array a -> Int -> (Array a, Array a)
splitAt Array a
arr Int
idx = (Array a -> Int -> Array a
forall a. Array a -> Int -> Array a
take Array a
arr Int
idx, Array a -> Int -> Array a
forall a. Array a -> Int -> Array a
drop Array a
arr Int
idx)

head :: Array a -> a
head :: forall a. Array a -> a
head Array a
arr = Array a -> Int -> a
forall a. Array a -> Int -> a
index Array a
arr Int
0

last :: Array a -> a
last :: forall a. Array a -> a
last Array a
arr = Array a -> Int -> a
forall a. Array a -> Int -> a
index Array a
arr (Array a -> Int
forall a. Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

snoc :: Array a -> a -> Array a
snoc :: forall a. Array a -> a -> Array a
snoc (Array Int
start Int
len SmallArray a
arr) a
x = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len' (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len' a
x
    SmallMutableArray (PrimState (ST s)) a
-> Int -> SmallArray a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
sma Int
0 SmallArray a
arr Int
start Int
len
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma
  where
    !len' :: Int
len' = Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

cons :: Array a -> a -> Array a
cons :: forall a. Array a -> a -> Array a
cons (Array Int
start Int
len SmallArray a
arr) a
x = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len' (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len' a
x
    SmallMutableArray (PrimState (ST s)) a
-> Int -> SmallArray a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
sma Int
1 SmallArray a
arr Int
start Int
len
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma
  where
    !len' :: Int
len' = Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1

(++) :: Array a -> Array a -> Array a
Array Int
start1 Int
len1 SmallArray a
arr1 ++ :: forall a. Array a -> Array a -> Array a
++ Array Int
start2 Int
len2 SmallArray a
arr2 = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len' (SmallArray a -> Array a) -> SmallArray a -> Array a
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s a)) -> SmallArray a)
-> (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len' a
forall a. a
uninitialized
    SmallMutableArray (PrimState (ST s)) a
-> Int -> SmallArray a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
sma Int
0 SmallArray a
arr1 Int
start1 Int
len1
    SmallMutableArray (PrimState (ST s)) a
-> Int -> SmallArray a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
sma Int
len1 SmallArray a
arr2 Int
start2 Int
len2
    SmallMutableArray s a -> ST s (SmallMutableArray s a)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma
  where
    !len' :: Int
len' = Int
len1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len2

map :: (a -> b) -> Array a -> Array b
map :: forall a b. (a -> b) -> Array a -> Array b
map a -> b
f (Array Int
start Int
len SmallArray a
arr) = Int -> Int -> SmallArray b -> Array b
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len (SmallArray b -> Array b) -> SmallArray b -> Array b
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s b)) -> SmallArray b
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s b)) -> SmallArray b)
-> (forall s. ST s (SmallMutableArray s b)) -> SmallArray b
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray (PrimState (ST s)) b
sma <- Int -> b -> ST s (SmallMutableArray (PrimState (ST s)) b)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len b
forall a. a
uninitialized
    -- i is the index in arr, j is the index in sma
    let loop :: Int -> Int -> ST s ()
loop Int
i Int
j = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j 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. Applicative 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 (PrimState (ST s)) b
sma Int
j (a -> b
f a
x)
            Int -> Int -> ST s ()
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> Int -> ST s ()
loop Int
start Int
0
    SmallMutableArray s b -> ST s (SmallMutableArray s b)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s b
SmallMutableArray (PrimState (ST s)) b
sma

map' :: (a -> b) -> Array a -> Array b
map' :: forall a b. (a -> b) -> Array a -> Array b
map' a -> b
f (Array Int
start Int
len SmallArray a
arr) = Int -> Int -> SmallArray b -> Array b
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len (SmallArray b -> Array b) -> SmallArray b -> Array b
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s b)) -> SmallArray b
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s b)) -> SmallArray b)
-> (forall s. ST s (SmallMutableArray s b)) -> SmallArray b
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray (PrimState (ST s)) b
sma <- Int -> b -> ST s (SmallMutableArray (PrimState (ST s)) b)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len b
forall a. a
uninitialized
    -- i is the index in arr, j is the index in sma
    let loop :: Int -> Int -> ST s ()
loop Int
i Int
j = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j 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. Applicative 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 (PrimState (ST s)) b
sma Int
j (b -> ST s ()) -> b -> ST s ()
forall a b. (a -> b) -> a -> b
$! a -> b
f a
x
            Int -> Int -> ST s ()
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> Int -> ST s ()
loop Int
start Int
0
    SmallMutableArray s b -> ST s (SmallMutableArray s b)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s b
SmallMutableArray (PrimState (ST s)) b
sma

-- helper function for implementing imap
imapStep :: Int -> (a -> Int) -> (Int -> a -> b) -> Array a -> Array b
imapStep :: forall a b.
Int -> (a -> Int) -> (Int -> a -> b) -> Array a -> Array b
imapStep Int
i0 a -> Int
step Int -> a -> b
f (Array Int
start Int
len SmallArray a
arr) = Int -> Int -> SmallArray b -> Array b
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len (SmallArray b -> Array b) -> SmallArray b -> Array b
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s b)) -> SmallArray b
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s b)) -> SmallArray b)
-> (forall s. ST s (SmallMutableArray s b)) -> SmallArray b
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray (PrimState (ST s)) b
sma <- Int -> b -> ST s (SmallMutableArray (PrimState (ST s)) b)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len b
forall a. a
uninitialized
    -- i is the index in arr, j is the index in sma, k is the index for f
    let loop :: Int -> Int -> Int -> ST s ()
loop !Int
i !Int
j !Int
k = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j 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. Applicative 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 (PrimState (ST s)) b
sma Int
j (Int -> a -> b
f Int
k a
x)
            Int -> Int -> Int -> ST s ()
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int
step a
x)
    Int -> Int -> Int -> ST s ()
loop Int
start Int
0 Int
i0
    SmallMutableArray s b -> ST s (SmallMutableArray s b)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s b
SmallMutableArray (PrimState (ST s)) b
sma

-- helper function for implementing imap
imapStep' :: Int -> (a -> Int) -> (Int -> a -> b) -> Array a -> Array b
imapStep' :: forall a b.
Int -> (a -> Int) -> (Int -> a -> b) -> Array a -> Array b
imapStep' Int
i0 a -> Int
step Int -> a -> b
f (Array Int
start Int
len SmallArray a
arr) = Int -> Int -> SmallArray b -> Array b
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len (SmallArray b -> Array b) -> SmallArray b -> Array b
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallMutableArray s b)) -> SmallArray b
forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray ((forall s. ST s (SmallMutableArray s b)) -> SmallArray b)
-> (forall s. ST s (SmallMutableArray s b)) -> SmallArray b
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray (PrimState (ST s)) b
sma <- Int -> b -> ST s (SmallMutableArray (PrimState (ST s)) b)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len b
forall a. a
uninitialized
    -- i is the index in arr, j is the index in sma, k is the index for f
    let loop :: Int -> Int -> Int -> ST s ()
loop !Int
i !Int
j !Int
k = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j 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. Applicative 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 (PrimState (ST s)) b
sma Int
j (b -> ST s ()) -> b -> ST s ()
forall a b. (a -> b) -> a -> b
$! Int -> a -> b
f Int
k a
x
            Int -> Int -> Int -> ST s ()
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int
step a
x)
    Int -> Int -> Int -> ST s ()
loop Int
start Int
0 Int
i0
    SmallMutableArray s b -> ST s (SmallMutableArray s b)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s b
SmallMutableArray (PrimState (ST s)) b
sma

unzipWith :: (a -> (b, c)) -> Array a -> (Array b, Array c)
unzipWith :: forall a b c. (a -> (b, c)) -> Array a -> (Array b, Array c)
unzipWith a -> (b, c)
f (Array Int
start Int
len SmallArray a
arr) = (forall s. ST s (Array b, Array c)) -> (Array b, Array c)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array b, Array c)) -> (Array b, Array c))
-> (forall s. ST s (Array b, Array c)) -> (Array b, Array c)
forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray (PrimState (ST s)) b
sma1 <- Int -> b -> ST s (SmallMutableArray (PrimState (ST s)) b)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len b
forall a. a
uninitialized
    SmallMutableArray (PrimState (ST s)) c
sma2 <- Int -> c -> ST s (SmallMutableArray (PrimState (ST s)) c)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len c
forall a. a
uninitialized
    -- i is the index in arr, j is the index in sma1/sma2
    let loop :: Int -> Int -> ST s ()
loop Int
i Int
j = Bool -> ST s () -> ST s ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j 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
val <- SmallArray a -> Int -> ST s a
forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
arr Int
i
            let !(b
x, c
y) = a -> (b, c)
f a
val
            SmallMutableArray (PrimState (ST s)) b -> Int -> b -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState (ST s)) b
sma1 Int
j b
x
            SmallMutableArray (PrimState (ST s)) c -> Int -> c -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState (ST s)) c
sma2 Int
j c
y
            Int -> Int -> ST s ()
loop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
    Int -> Int -> ST s ()
loop Int
start Int
0
    SmallArray b
arr1 <- SmallMutableArray (PrimState (ST s)) b -> ST s (SmallArray b)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray (PrimState (ST s)) b
sma1
    SmallArray c
arr2 <- SmallMutableArray (PrimState (ST s)) c -> ST s (SmallArray c)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray (PrimState (ST s)) c
sma2
    (Array b, Array c) -> ST s (Array b, Array c)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Int -> SmallArray b -> Array b
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len SmallArray b
arr1, Int -> Int -> SmallArray c -> Array c
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len SmallArray c
arr2)

newtype STA a = STA (forall s. SmallMutableArray s a -> ST s (SmallArray a))

runSTA :: Int -> STA a -> Array a
runSTA :: forall a. Int -> STA a -> Array a
runSTA Int
len (STA forall s. SmallMutableArray s a -> ST s (SmallArray a)
m) = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len ((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
len a
forall a. a
uninitialized ST s (SmallMutableArray s a)
-> (SmallMutableArray s a -> ST s (SmallArray a))
-> ST s (SmallArray a)
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SmallMutableArray s a -> ST s (SmallArray a)
forall s. SmallMutableArray s a -> ST s (SmallArray a)
m)

traverse :: (Applicative f) => (a -> f b) -> Array a -> f (Array b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array a -> f (Array b)
traverse a -> f b
f (Array Int
start Int
len SmallArray a
arr) =
    -- i is the index in arr, j is the index in sma
    let go :: Int -> Int -> f (STA b)
go Int
i Int
j
            | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = STA b -> f (STA b)
forall a. a -> f a
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 SmallMutableArray s b -> ST s (SmallArray b)
SmallMutableArray (PrimState (ST s)) b -> ST s (SmallArray b)
forall s. SmallMutableArray s b -> ST s (SmallArray b)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray
            | (# 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 a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\b
y (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
sma -> 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
sma Int
j b
y ST s () -> ST s (SmallArray b) -> ST s (SmallArray b)
forall a b. ST s a -> ST s b -> ST s b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> SmallMutableArray s b -> ST s (SmallArray b)
forall s. SmallMutableArray s b -> ST s (SmallArray b)
m SmallMutableArray s b
sma) (a -> f b
f a
x) (Int -> Int -> f (STA b)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
    in Int -> STA b -> Array b
forall a. Int -> STA a -> Array a
runSTA Int
len (STA b -> Array b) -> f (STA b) -> f (Array b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> f (STA b)
go Int
start Int
0

traverse' :: (Applicative f) => (a -> f b) -> Array a -> f (Array b)
traverse' :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Array a -> f (Array b)
traverse' a -> f b
f (Array Int
start Int
len SmallArray a
arr) =
    -- i is the index in arr, j is the index in sma
    let go :: Int -> Int -> f (STA b)
go Int
i Int
j
            | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = STA b -> f (STA b)
forall a. a -> f a
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 SmallMutableArray s b -> ST s (SmallArray b)
SmallMutableArray (PrimState (ST s)) b -> ST s (SmallArray b)
forall s. SmallMutableArray s b -> ST s (SmallArray b)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray
            | (# 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 a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\ !b
y (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
sma -> 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
sma Int
j b
y ST s () -> ST s (SmallArray b) -> ST s (SmallArray b)
forall a b. ST s a -> ST s b -> ST s b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> SmallMutableArray s b -> ST s (SmallArray b)
forall s. SmallMutableArray s b -> ST s (SmallArray b)
m SmallMutableArray s b
sma) (a -> f b
f a
x) (Int -> Int -> f (STA b)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
    in Int -> STA b -> Array b
forall a. Int -> STA a -> Array a
runSTA Int
len (STA b -> Array b) -> f (STA b) -> f (Array b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> f (STA b)
go Int
start Int
0

-- helper function for implementing itraverse
itraverseStep :: (Applicative f) => Int -> (a -> Int) -> (Int -> a -> f b) -> Array a -> f (Array b)
itraverseStep :: forall (f :: * -> *) a b.
Applicative f =>
Int -> (a -> Int) -> (Int -> a -> f b) -> Array a -> f (Array b)
itraverseStep Int
i0 a -> Int
step Int -> a -> f b
f (Array Int
start Int
len SmallArray a
arr) =
    -- i is the index in arr, j is the index in sma, k is the index for f
    let go :: Int -> Int -> Int -> f (STA b)
go !Int
i !Int
j !Int
k
            | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = STA b -> f (STA b)
forall a. a -> f a
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 SmallMutableArray s b -> ST s (SmallArray b)
SmallMutableArray (PrimState (ST s)) b -> ST s (SmallArray b)
forall s. SmallMutableArray s b -> ST s (SmallArray b)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray
            | (# 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 a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\b
y (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
sma -> 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
sma Int
j b
y ST s () -> ST s (SmallArray b) -> ST s (SmallArray b)
forall a b. ST s a -> ST s b -> ST s b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> SmallMutableArray s b -> ST s (SmallArray b)
forall s. SmallMutableArray s b -> ST s (SmallArray b)
m SmallMutableArray s b
sma) (Int -> a -> f b
f Int
k a
x) (Int -> Int -> Int -> f (STA b)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int
step a
x))
    in Int -> STA b -> Array b
forall a. Int -> STA a -> Array a
runSTA Int
len (STA b -> Array b) -> f (STA b) -> f (Array b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> Int -> f (STA b)
go Int
start Int
0 Int
i0

-- helper function for implementing itraverse
itraverseStep' :: (Applicative f) => Int -> (a -> Int) -> (Int -> a -> f b) -> Array a -> f (Array b)
itraverseStep' :: forall (f :: * -> *) a b.
Applicative f =>
Int -> (a -> Int) -> (Int -> a -> f b) -> Array a -> f (Array b)
itraverseStep' Int
i0 a -> Int
step Int -> a -> f b
f (Array Int
start Int
len SmallArray a
arr) =
    -- i is the index in arr, j is the index in sma, k is the index for f
    let go :: Int -> Int -> Int -> f (STA b)
go !Int
i !Int
j !Int
k
            | Int
j Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
len = STA b -> f (STA b)
forall a. a -> f a
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 SmallMutableArray s b -> ST s (SmallArray b)
SmallMutableArray (PrimState (ST s)) b -> ST s (SmallArray b)
forall s. SmallMutableArray s b -> ST s (SmallArray b)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray
            | (# 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 a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (\ !b
y (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
sma -> 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
sma Int
j b
y ST s () -> ST s (SmallArray b) -> ST s (SmallArray b)
forall a b. ST s a -> ST s b -> ST s b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> SmallMutableArray s b -> ST s (SmallArray b)
forall s. SmallMutableArray s b -> ST s (SmallArray b)
m SmallMutableArray s b
sma) (Int -> a -> f b
f Int
k a
x) (Int -> Int -> Int -> f (STA b)
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ a -> Int
step a
x))
    in Int -> STA b -> Array b
forall a. Int -> STA a -> Array a
runSTA Int
len (STA b -> Array b) -> f (STA b) -> f (Array b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> Int -> f (STA b)
go Int
start Int
0 Int
i0

new :: Int -> ST s (MutableArray s a)
new :: forall s a. Int -> ST s (MutableArray s a)
new Int
len = Int -> Int -> SmallMutableArray s a -> MutableArray s a
forall s a. Int -> Int -> SmallMutableArray s a -> MutableArray s a
MutableArray Int
0 Int
len (SmallMutableArray s a -> MutableArray s a)
-> ST s (SmallMutableArray s a) -> ST s (MutableArray s a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len a
forall a. a
uninitialized

read :: MutableArray s a -> Int -> ST s a
read :: forall s a. MutableArray s a -> Int -> ST s a
read (MutableArray Int
start Int
_ SmallMutableArray s a
arr) Int
idx = SmallMutableArray (PrimState (ST s)) a -> Int -> ST s a
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
arr (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx)

write :: MutableArray s a -> Int -> a -> ST s ()
write :: forall s a. MutableArray s a -> Int -> a -> ST s ()
write (MutableArray Int
start Int
_ SmallMutableArray s a
arr) Int
idx = 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
arr (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx)

freeze :: MutableArray s a -> Int -> Int -> ST s (Array a)
freeze :: forall s a. MutableArray s a -> Int -> Int -> ST s (Array a)
freeze (MutableArray Int
start Int
_ SmallMutableArray s a
arr) Int
idx Int
len = Int -> Int -> SmallArray a -> Array a
forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len (SmallArray a -> Array a) -> ST s (SmallArray a) -> ST s (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SmallMutableArray (PrimState (ST s)) a
-> Int -> Int -> ST s (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a)
freezeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
arr (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx) Int
len

thaw :: Array a -> Int -> Int -> ST s (MutableArray s a)
thaw :: forall a s. Array a -> Int -> Int -> ST s (MutableArray s a)
thaw (Array Int
start Int
_ SmallArray a
arr) Int
idx Int
len = Int -> Int -> SmallMutableArray s a -> MutableArray s a
forall s a. Int -> Int -> SmallMutableArray s a -> MutableArray s a
MutableArray Int
0 Int
len (SmallMutableArray s a -> MutableArray s a)
-> ST s (SmallMutableArray s a) -> ST s (MutableArray s a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SmallArray 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 SmallArray a
arr (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
idx) Int
len