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

-- Warning: No bound checks are performed!

module Data.RRBVector.Internal.Array
    ( Array, MutableArray
    , empty, singleton, from2
    , index, head, last
    , update, adjust, adjust'
    , take, drop, splitAt
    , snoc, cons
    , map, map'
    , traverse, traverse'
    , new, read, write
    , freeze, thaw
    ) where

import Control.Applicative (Applicative(liftA2))
import Control.DeepSeq (NFData(..))
import Control.Monad (when)
import Control.Monad.ST
import Data.Foldable (Foldable(..))
import Data.Primitive.SmallArray
import Prelude hiding (take, drop, splitAt, head, last, map, traverse, read)

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

instance Semigroup (Array a) where
    Array Int
start1 Int
len1 SmallArray a
arr1 <> :: 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 (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

instance Foldable Array where
    foldr :: (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
    {-# INLINE foldr #-}

    foldl :: (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)
    {-# INLINE foldl #-}

    foldr' :: (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
    {-# INLINE foldr' #-}

    foldl' :: (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
    {-# INLINE foldl' #-}

    null :: Array a -> Bool
null Array a
arr = 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
    {-# INLINE null #-}

    length :: Array a -> Int
length (Array Int
_ Int
len SmallArray a
_) = Int
len
    {-# INLINE length #-}

instance (NFData a) => NFData (Array a) where
    rnf :: Array a -> ()
rnf = (() -> a -> ()) -> () -> Array a -> ()
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) ()

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

empty :: Array a
empty :: 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 :: 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 :: 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 (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma

index :: Array a -> Int -> a
index :: 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 :: 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 (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma

adjust :: Array a -> Int -> (a -> a) -> Array a
adjust :: 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. Monad 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 (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma

adjust' :: Array a -> Int -> (a -> a) -> Array a
adjust' :: 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. Monad 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 (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s a
sma

take :: Array a -> Int -> Array a
take :: 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 :: 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 :: 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 :: 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 :: 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 (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 :: Array a -> a -> Array a
snoc (Array Int
_ 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
0 Int
len
    SmallMutableArray s a -> ST s (SmallMutableArray 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 :: Array a -> a -> Array a
cons (Array Int
_ 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
0 Int
len
    SmallMutableArray s a -> ST s (SmallMutableArray 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

map :: (a -> b) -> Array a -> Array b
map :: (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 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. 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
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 (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s b
sma

map' :: (a -> b) -> Array a -> Array b
map' :: (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 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. 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
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 (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray s b
sma

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

runSTA :: Int -> STA a -> Array a
runSTA :: 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 (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 :: (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 (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)
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 (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 (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' :: (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 (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)
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 (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 (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

new :: Int -> ST s (MutableArray s a)
new :: 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 :: 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 :: 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 :: 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 :: 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