{-# 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
    , 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 forall a. Num a => a -> a -> a
+ Int
len
            go :: Int -> b
go Int
i
                | Int
i forall a. Eq a => a -> a -> Bool
== Int
end = b
z
                | (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = a -> b -> b
f a
x (Int -> b
go (Int
i 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 forall a. Ord a => a -> a -> Bool
< Int
start = b
z
                | (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = b -> a -> b
f (Int -> b
go (Int
i forall a. Num a => a -> a -> a
- Int
1)) a
x
        in Int -> b
go (Int
start forall a. Num a => a -> a -> a
+ Int
len 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 forall a. Ord a => a -> a -> Bool
< Int
start = b
acc
                | (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = Int -> b -> b
go (Int
i forall a. Num a => a -> a -> a
- Int
1) (a -> b -> b
f a
x b
acc)
        in Int -> b -> b
go (Int
start forall a. Num a => a -> a -> a
+ Int
len 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 forall a. Num a => a -> a -> a
+ Int
len
            go :: Int -> b -> b
go Int
i !b
acc
                | Int
i forall a. Eq a => a -> a -> Bool
== Int
end = b
acc
                | (# a
x #) <- forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i = Int -> b -> b
go (Int
i 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 = forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr 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 = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\()
_ a
x -> 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 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 forall a. Eq a => a -> a -> Bool
== Int
end = 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) (Int
j 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 forall a. Ord a => a -> a -> Bool
< Int
start = 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) (Int
j forall a. Num a => a -> a -> a
- a -> Int
step a
x)) a
x
    in Int -> Int -> b
go (Int
start forall a. Num a => a -> a -> a
+ Int
len 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 forall a. Ord a => a -> a -> Bool
< Int
start = 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) (Int
j 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 forall a. Num a => a -> a -> a
+ Int
len 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 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 forall a. Eq a => a -> a -> Bool
== Int
end = 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) (Int
j 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 = forall a. [Char] -> a
errorWithoutStackTrace [Char]
"uninitialized"

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

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

replicate :: Int -> a -> Array a
replicate :: forall a. Int -> a -> Array a
replicate Int
n a
x = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
n forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray (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 = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len a
x
    forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
sma Int
n a
y
    forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma
  where
    len :: Int
len = Int
n 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 = forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray a
arr (Int
start 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 = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray SmallArray a
sa Int
start Int
len
    forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
sma Int
idx a
x
    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 = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray SmallArray a
sa Int
start Int
len
    a
x <- forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
sa (Int
start forall a. Num a => a -> a -> a
+ Int
idx)
    forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
sma Int
idx (a -> a
f a
x)
    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 = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray SmallArray a
sa Int
start Int
len
    a
x <- forall (m :: * -> *) a. Applicative m => SmallArray a -> Int -> m a
indexSmallArrayM SmallArray a
sa (Int
start forall a. Num a => a -> a -> a
+ Int
idx)
    forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
sma Int
idx forall a b. (a -> b) -> a -> b
$! a -> a
f a
x
    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 = 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 = forall a. Int -> Int -> SmallArray a -> Array a
Array (Int
start forall a. Num a => a -> a -> a
+ Int
n) (Int
len 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 = (forall a. Array a -> Int -> Array a
take Array a
arr Int
idx, 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 = 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 = forall a. Array a -> Int -> a
index Array a
arr (forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
arr 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 = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len' forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len' a
x
    forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
sma Int
0 SmallArray a
arr Int
start Int
len
    forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma
  where
    !len' :: Int
len' = Int
len 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 = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len' forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len' a
x
    forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
sma Int
1 SmallArray a
arr Int
start Int
len
    forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma
  where
    !len' :: Int
len' = Int
len 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 = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len' forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray s a
sma <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len' forall a. a
uninitialized
    forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
sma Int
0 SmallArray a
arr1 Int
start1 Int
len1
    forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
copySmallArray SmallMutableArray s a
sma Int
len1 SmallArray a
arr2 Int
start2 Int
len2
    forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma
  where
    !len' :: Int
len' = Int
len1 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) = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray (PrimState (ST s)) b
sma <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len 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 = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j forall a. Ord a => a -> a -> Bool
< Int
len) forall a b. (a -> b) -> a -> b
$ do
            a
x <- forall (m :: * -> *) a. Applicative 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 (PrimState (ST s)) b
sma Int
j (a -> b
f a
x)
            Int -> Int -> ST s ()
loop (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Int
j forall a. Num a => a -> a -> a
+ Int
1)
    Int -> Int -> ST s ()
loop Int
start Int
0
    forall (f :: * -> *) a. Applicative f => a -> f a
pure 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) = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray (PrimState (ST s)) b
sma <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len 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 = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j forall a. Ord a => a -> a -> Bool
< Int
len) forall a b. (a -> b) -> a -> b
$ do
            a
x <- forall (m :: * -> *) a. Applicative 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 (PrimState (ST s)) b
sma Int
j forall a b. (a -> b) -> a -> b
$! a -> b
f a
x
            Int -> Int -> ST s ()
loop (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Int
j forall a. Num a => a -> a -> a
+ Int
1)
    Int -> Int -> ST s ()
loop Int
start Int
0
    forall (f :: * -> *) a. Applicative f => a -> f a
pure 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) = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray (PrimState (ST s)) b
sma <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len 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 = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j forall a. Ord a => a -> a -> Bool
< Int
len) forall a b. (a -> b) -> a -> b
$ do
            a
x <- forall (m :: * -> *) a. Applicative 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 (PrimState (ST s)) b
sma Int
j (Int -> a -> b
f Int
k a
x)
            Int -> Int -> Int -> ST s ()
loop (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Int
j forall a. Num a => a -> a -> a
+ Int
1) (Int
k forall a. Num a => a -> a -> a
+ a -> Int
step a
x)
    Int -> Int -> Int -> ST s ()
loop Int
start Int
0 Int
i0
    forall (f :: * -> *) a. Applicative f => a -> f a
pure 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) = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len forall a b. (a -> b) -> a -> b
$ forall a. (forall s. ST s (SmallMutableArray s a)) -> SmallArray a
runSmallArray forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray (PrimState (ST s)) b
sma <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len 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 = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j forall a. Ord a => a -> a -> Bool
< Int
len) forall a b. (a -> b) -> a -> b
$ do
            a
x <- forall (m :: * -> *) a. Applicative 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 (PrimState (ST s)) b
sma Int
j forall a b. (a -> b) -> a -> b
$! Int -> a -> b
f Int
k a
x
            Int -> Int -> Int -> ST s ()
loop (Int
i forall a. Num a => a -> a -> a
+ Int
1) (Int
j forall a. Num a => a -> a -> a
+ Int
1) (Int
k forall a. Num a => a -> a -> a
+ a -> Int
step a
x)
    Int -> Int -> Int -> ST s ()
loop Int
start Int
0 Int
i0
    forall (f :: * -> *) a. Applicative f => a -> f a
pure 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 a. (forall s. ST s a) -> a
runST forall a b. (a -> b) -> a -> b
$ do
    SmallMutableArray (PrimState (ST s)) b
sma1 <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len forall a. a
uninitialized
    SmallMutableArray (PrimState (ST s)) c
sma2 <- forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len 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 = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
j forall a. Ord a => a -> a -> Bool
< Int
len) forall a b. (a -> b) -> a -> b
$ do
            a
val <- 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
            forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState (ST s)) b
sma1 Int
j b
x
            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 forall a. Num a => a -> a -> a
+ Int
1) (Int
j forall a. Num a => a -> a -> a
+ Int
1)
    Int -> Int -> ST s ()
loop Int
start Int
0
    SmallArray b
arr1 <- forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray (PrimState (ST s)) b
sma1
    SmallArray c
arr2 <- forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray SmallMutableArray (PrimState (ST s)) c
sma2
    forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len SmallArray b
arr1, 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) = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len (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
len forall a. a
uninitialized forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= 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 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 (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray
            | (# 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
y (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
sma -> forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s b
sma Int
j b
y forall (f :: * -> *) a b. Applicative f => f a -> f b -> f 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 forall a. Num a => a -> a -> a
+ Int
1) (Int
j forall a. Num a => a -> a -> a
+ Int
1))
    in forall a. Int -> STA a -> Array a
runSTA Int
len 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 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 (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray
            | (# 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
y (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
sma -> forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s b
sma Int
j b
y forall (f :: * -> *) a b. Applicative f => f a -> f b -> f 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 forall a. Num a => a -> a -> a
+ Int
1) (Int
j forall a. Num a => a -> a -> a
+ Int
1))
    in forall a. Int -> STA a -> Array a
runSTA Int
len 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 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 (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray
            | (# 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
y (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
sma -> forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s b
sma Int
j b
y forall (f :: * -> *) a b. Applicative f => f a -> f b -> f 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 forall a. Num a => a -> a -> a
+ Int
1) (Int
j forall a. Num a => a -> a -> a
+ Int
1) (Int
k forall a. Num a => a -> a -> a
+ a -> Int
step a
x))
    in forall a. Int -> STA a -> Array a
runSTA Int
len 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 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 (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
unsafeFreezeSmallArray
            | (# 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
y (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
sma -> forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s b
sma Int
j b
y forall (f :: * -> *) a b. Applicative f => f a -> f b -> f 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 forall a. Num a => a -> a -> a
+ Int
1) (Int
j forall a. Num a => a -> a -> a
+ Int
1) (Int
k forall a. Num a => a -> a -> a
+ a -> Int
step a
x))
    in forall a. Int -> STA a -> Array a
runSTA Int
len 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 = forall s a. Int -> Int -> SmallMutableArray s a -> MutableArray s a
MutableArray Int
0 Int
len forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len 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 = forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> m a
readSmallArray SmallMutableArray s a
arr (Int
start 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 = forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray s a
arr (Int
start 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 = forall a. Int -> Int -> SmallArray a -> Array a
Array Int
0 Int
len forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a)
freezeSmallArray SmallMutableArray s a
arr (Int
start 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 = forall s a. Int -> Int -> SmallMutableArray s a -> MutableArray s a
MutableArray Int
0 Int
len forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
PrimMonad m =>
SmallArray a -> Int -> Int -> m (SmallMutableArray (PrimState m) a)
thawSmallArray SmallArray a
arr (Int
start forall a. Num a => a -> a -> a
+ Int
idx) Int
len