{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : Data.Massiv.Array.Manifest.Primitive
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Array.Manifest.Primitive (
  P (..),
  Array (..),
  MArray (..),
  Prim,
  toPrimitiveVector,
  toPrimitiveMVector,
  fromPrimitiveVector,
  fromPrimitiveMVector,
  toByteArray,
  toByteArrayM,
  unwrapByteArray,
  unwrapByteArrayOffset,
  unwrapMutableByteArray,
  unwrapMutableByteArrayOffset,
  fromByteArray,
  fromByteArrayM,
  fromByteArrayOffsetM,
  toMutableByteArray,
  toMutableByteArrayM,
  fromMutableByteArrayM,
  fromMutableByteArrayOffsetM,
  fromMutableByteArray,
  shrinkMutableByteArray,
  unsafeAtomicReadIntArray,
  unsafeAtomicWriteIntArray,
  unsafeCasIntArray,
  unsafeAtomicModifyIntArray,
  unsafeAtomicAddIntArray,
  unsafeAtomicSubIntArray,
  unsafeAtomicAndIntArray,
  unsafeAtomicNandIntArray,
  unsafeAtomicOrIntArray,
  unsafeAtomicXorIntArray,
) where

import Control.DeepSeq (NFData (..), deepseq)
import Control.Monad
import Control.Monad.Primitive (PrimMonad (..), primitive_)
import Data.Massiv.Array.Delayed.Pull -- (eq, ord)
import Data.Massiv.Array.Manifest.Internal
import Data.Massiv.Array.Manifest.List as A
import Data.Massiv.Array.Mutable
import Data.Massiv.Core.Common
import Data.Massiv.Core.List
import Data.Massiv.Core.Operations
import Data.Massiv.Vector.Stream as S (isteps, steps)
import Data.Maybe (fromMaybe)
import Data.Primitive (Prim, sizeOf)
import Data.Primitive.ByteArray
import qualified Data.Vector.Primitive as VP
import qualified Data.Vector.Primitive.Mutable as MVP
import GHC.Exts as GHC
import System.IO.Unsafe (unsafePerformIO)
import Prelude hiding (mapM)

-- | Representation for `Prim`itive elements
data P = P deriving (Int -> P -> ShowS
[P] -> ShowS
P -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [P] -> ShowS
$cshowList :: [P] -> ShowS
show :: P -> String
$cshow :: P -> String
showsPrec :: Int -> P -> ShowS
$cshowsPrec :: Int -> P -> ShowS
Show)

data instance Array P ix e = PArray
  { forall ix e. Array P ix e -> Comp
pComp :: !Comp
  , forall ix e. Array P ix e -> Sz ix
pSize :: !(Sz ix)
  , forall ix e. Array P ix e -> Int
pOffset :: {-# UNPACK #-} !Int
  , forall ix e. Array P ix e -> ByteArray
pData :: {-# UNPACK #-} !ByteArray
  }

data instance MArray s P ix e
  = MPArray !(Sz ix) {-# UNPACK #-} !Int {-# UNPACK #-} !(MutableByteArray s)

instance (Ragged L ix e, Show e, Prim e) => Show (Array P ix e) where
  showsPrec :: Int -> Array P ix e -> ShowS
showsPrec = forall r r' ix e.
(Ragged L ix e, Load r ix e, Load r' ix e, Source r' e, Show e) =>
(Array r ix e -> Array r' ix e) -> Int -> Array r ix e -> ShowS
showsArrayPrec forall a. a -> a
id
  showList :: [Array P ix e] -> ShowS
showList = forall arr. Show arr => [arr] -> ShowS
showArrayList

instance Index ix => NFData (Array P ix e) where
  rnf :: Array P ix e -> ()
rnf (PArray Comp
c Sz ix
sz Int
o ByteArray
a) = Comp
c forall a b. NFData a => a -> b -> b
`deepseq` Sz ix
sz forall a b. NFData a => a -> b -> b
`deepseq` Int
o seq :: forall a b. a -> b -> b
`seq` ByteArray
a seq :: forall a b. a -> b -> b
`seq` ()
  {-# INLINE rnf #-}

instance NFData ix => NFData (MArray s P ix e) where
  rnf :: MArray s P ix e -> ()
rnf (MPArray Sz ix
sz Int
_o MutableByteArray s
_mb) = Sz ix
sz forall a b. NFData a => a -> b -> b
`deepseq` ()
  {-# INLINE rnf #-}

instance (Prim e, Eq e, Index ix) => Eq (Array P ix e) where
  == :: Array P ix e -> Array P ix e -> Bool
(==) = forall ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> Bool) -> Array r1 ix e1 -> Array r2 ix e2 -> Bool
eqArrays forall a. Eq a => a -> a -> Bool
(==)
  {-# INLINE (==) #-}

instance (Prim e, Ord e, Index ix) => Ord (Array P ix e) where
  compare :: Array P ix e -> Array P ix e -> Ordering
compare = forall ix r1 e1 r2 e2.
(Index ix, Source r1 e1, Source r2 e2) =>
(e1 -> e2 -> Ordering)
-> Array r1 ix e1 -> Array r2 ix e2 -> Ordering
compareArrays forall a. Ord a => a -> a -> Ordering
compare
  {-# INLINE compare #-}

instance Strategy P where
  getComp :: forall ix e. Array P ix e -> Comp
getComp = forall ix e. Array P ix e -> Comp
pComp
  {-# INLINE getComp #-}
  setComp :: forall ix e. Comp -> Array P ix e -> Array P ix e
setComp Comp
c Array P ix e
arr = Array P ix e
arr{pComp :: Comp
pComp = Comp
c}
  {-# INLINE setComp #-}
  repr :: P
repr = P
P

instance Index ix => Shape P ix where
  maxLinearSize :: forall e. Array P ix e -> Maybe Sz1
maxLinearSize = forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount
  {-# INLINE maxLinearSize #-}

instance Size P where
  size :: forall ix e. Array P ix e -> Sz ix
size = forall ix e. Array P ix e -> Sz ix
pSize
  {-# INLINE size #-}
  unsafeResize :: forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array P ix e -> Array P ix' e
unsafeResize !Sz ix'
sz !Array P ix e
arr = Array P ix e
arr{pSize :: Sz ix'
pSize = Sz ix'
sz}
  {-# INLINE unsafeResize #-}

instance Prim e => Source P e where
  unsafeLinearIndex :: forall ix. Index ix => Array P ix e -> Int -> e
unsafeLinearIndex _arr :: Array P ix e
_arr@(PArray Comp
_ Sz ix
_ Int
o ByteArray
a) Int
i =
    forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
"P.unsafeLinearIndex" (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e. Prim e => proxy e -> ByteArray -> Int
elemsBA Array P ix e
_arr) forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
a (Int
i forall a. Num a => a -> a -> a
+ Int
o)
  {-# INLINE unsafeLinearIndex #-}

  unsafeOuterSlice :: forall ix.
(Index ix, Index (Lower ix)) =>
Array P ix e -> Sz (Lower ix) -> Int -> Array P (Lower ix) e
unsafeOuterSlice (PArray Comp
c Sz ix
_ Int
o ByteArray
a) Sz (Lower ix)
szL Int
i =
    forall ix e. Comp -> Sz ix -> Int -> ByteArray -> Array P ix e
PArray Comp
c Sz (Lower ix)
szL (Int
i forall a. Num a => a -> a -> a
* forall ix. Index ix => Sz ix -> Int
totalElem Sz (Lower ix)
szL forall a. Num a => a -> a -> a
+ Int
o) ByteArray
a
  {-# INLINE unsafeOuterSlice #-}

  unsafeLinearSlice :: forall ix. Index ix => Int -> Sz1 -> Array P ix e -> Array P Int e
unsafeLinearSlice Int
i Sz1
k (PArray Comp
c Sz ix
_ Int
o ByteArray
a) = forall ix e. Comp -> Sz ix -> Int -> ByteArray -> Array P ix e
PArray Comp
c Sz1
k (Int
i forall a. Num a => a -> a -> a
+ Int
o) ByteArray
a
  {-# INLINE unsafeLinearSlice #-}

instance Prim e => Manifest P e where
  unsafeLinearIndexM :: forall ix. Index ix => Array P ix e -> Int -> e
unsafeLinearIndexM _pa :: Array P ix e
_pa@(PArray Comp
_ Sz ix
_sz Int
o ByteArray
a) Int
i =
    forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
"P.unsafeLinearIndexM" (forall a b. a -> b -> a
const (forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
_sz))) forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
a (Int
i forall a. Num a => a -> a -> a
+ Int
o)
  {-# INLINE unsafeLinearIndexM #-}

  sizeOfMArray :: forall ix s. Index ix => MArray s P ix e -> Sz ix
sizeOfMArray (MPArray Sz ix
sz Int
_ MutableByteArray s
_) = Sz ix
sz
  {-# INLINE sizeOfMArray #-}

  unsafeResizeMArray :: forall ix' ix s.
(Index ix', Index ix) =>
Sz ix' -> MArray s P ix e -> MArray s P ix' e
unsafeResizeMArray Sz ix'
sz (MPArray Sz ix
_ Int
off MutableByteArray s
marr) = forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray Sz ix'
sz Int
off MutableByteArray s
marr
  {-# INLINE unsafeResizeMArray #-}

  unsafeLinearSliceMArray :: forall ix s.
Index ix =>
Int -> Sz1 -> MArray s P ix e -> MVector s P e
unsafeLinearSliceMArray Int
i Sz1
k (MPArray Sz ix
_ Int
o MutableByteArray s
a) = forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray Sz1
k (Int
i forall a. Num a => a -> a -> a
+ Int
o) MutableByteArray s
a
  {-# INLINE unsafeLinearSliceMArray #-}

  unsafeThaw :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Array P ix e -> m (MArray (PrimState m) P ix e)
unsafeThaw (PArray Comp
_ Sz ix
sz Int
o ByteArray
a) = forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray Sz ix
sz Int
o forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
PrimMonad m =>
ByteArray -> m (MutableByteArray (PrimState m))
unsafeThawByteArray ByteArray
a
  {-# INLINE unsafeThaw #-}

  unsafeFreeze :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) P ix e -> m (Array P ix e)
unsafeFreeze Comp
comp (MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
a) = forall ix e. Comp -> Sz ix -> Int -> ByteArray -> Array P ix e
PArray Comp
comp Sz ix
sz Int
o forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray (PrimState m)
a
  {-# INLINE unsafeFreeze #-}

  unsafeNew :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) P ix e)
unsafeNew Sz ix
sz
    | Int
n forall a. Ord a => a -> a -> Bool
<= (forall a. Bounded a => a
maxBound :: Int) forall a. Integral a => a -> a -> a
`div` Int
eSize = forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray Sz ix
sz Int
0 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (Int
n forall a. Num a => a -> a -> a
* Int
eSize)
    | Bool
otherwise = forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"Array size is too big: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Sz ix
sz
    where
      !n :: Int
n = forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz
      !eSize :: Int
eSize = forall a. Prim a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e)
  {-# INLINE unsafeNew #-}

  initialize :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix e -> m ()
initialize (MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) =
    let k :: Int
k = forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz forall a. Num a => a -> a -> a
* forall a. Prim a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e)
     in forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
k forall a. Ord a => a -> a -> Bool
> Int
0) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> Int -> Word8 -> m ()
fillByteArray MutableByteArray (PrimState m)
mba Int
o Int
k Word8
0
  {-# INLINE initialize #-}

  unsafeLinearRead :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix e -> Int -> m e
unsafeLinearRead _mpa :: MArray (PrimState m) P ix e
_mpa@(MPArray Sz ix
_sz Int
o MutableByteArray (PrimState m)
ma) Int
i =
    forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
"P.unsafeLinearRead" (forall a b. a -> b -> a
const (forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
_sz))) forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray (PrimState m)
ma (Int
i forall a. Num a => a -> a -> a
+ Int
o)
  {-# INLINE unsafeLinearRead #-}

  unsafeLinearWrite :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix e -> Int -> e -> m ()
unsafeLinearWrite _mpa :: MArray (PrimState m) P ix e
_mpa@(MPArray Sz ix
_sz Int
o MutableByteArray (PrimState m)
ma) Int
i =
    forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert String
"P.unsafeLinearWrite" (forall a b. a -> b -> a
const (forall ix. Index ix => ix -> Sz ix
Sz (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
_sz))) forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray (PrimState m)
ma (Int
i forall a. Num a => a -> a -> a
+ Int
o)
  {-# INLINE unsafeLinearWrite #-}

  unsafeLinearSet :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix e -> Int -> Sz1 -> e -> m ()
unsafeLinearSet (MPArray Sz ix
_ Int
o MutableByteArray (PrimState m)
ma) Int
offset (SafeSz Int
sz) = forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> Int -> a -> m ()
setByteArray MutableByteArray (PrimState m)
ma (Int
offset forall a. Num a => a -> a -> a
+ Int
o) Int
sz
  {-# INLINE unsafeLinearSet #-}

  unsafeLinearCopy :: forall ix' ix (m :: * -> *).
(Index ix', Index ix, PrimMonad m) =>
MArray (PrimState m) P ix' e
-> Int -> MArray (PrimState m) P ix e -> Int -> Sz1 -> m ()
unsafeLinearCopy (MPArray Sz ix'
_ Int
oFrom MutableByteArray (PrimState m)
maFrom) Int
iFrom (MPArray Sz ix
_ Int
oTo MutableByteArray (PrimState m)
maTo) Int
iTo (Sz Int
k) =
    forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> MutableByteArray (PrimState m) -> Int -> Int -> m ()
copyMutableByteArray MutableByteArray (PrimState m)
maTo ((Int
oTo forall a. Num a => a -> a -> a
+ Int
iTo) forall a. Num a => a -> a -> a
* Int
esz) MutableByteArray (PrimState m)
maFrom ((Int
oFrom forall a. Num a => a -> a -> a
+ Int
iFrom) forall a. Num a => a -> a -> a
* Int
esz) (Int
k forall a. Num a => a -> a -> a
* Int
esz)
    where
      esz :: Int
esz = forall a. Prim a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e)
  {-# INLINE unsafeLinearCopy #-}

  unsafeArrayLinearCopy :: forall ix' ix (m :: * -> *).
(Index ix', Index ix, PrimMonad m) =>
Array P ix' e
-> Int -> MArray (PrimState m) P ix e -> Int -> Sz1 -> m ()
unsafeArrayLinearCopy (PArray Comp
_ Sz ix'
_ Int
oFrom ByteArray
aFrom) Int
iFrom (MPArray Sz ix
_ Int
oTo MutableByteArray (PrimState m)
maTo) Int
iTo (Sz Int
k) =
    forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
copyByteArray MutableByteArray (PrimState m)
maTo ((Int
oTo forall a. Num a => a -> a -> a
+ Int
iTo) forall a. Num a => a -> a -> a
* Int
esz) ByteArray
aFrom ((Int
oFrom forall a. Num a => a -> a -> a
+ Int
iFrom) forall a. Num a => a -> a -> a
* Int
esz) (Int
k forall a. Num a => a -> a -> a
* Int
esz)
    where
      esz :: Int
esz = forall a. Prim a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e)
  {-# INLINE unsafeArrayLinearCopy #-}

  unsafeLinearShrink :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix e
-> Sz ix -> m (MArray (PrimState m) P ix e)
unsafeLinearShrink (MPArray Sz ix
_ Int
o MutableByteArray (PrimState m)
ma) Sz ix
sz = do
    forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> Int -> m ()
shrinkMutableByteArray MutableByteArray (PrimState m)
ma ((Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) forall a. Num a => a -> a -> a
* forall a. Prim a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e))
    forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
ma
  {-# INLINE unsafeLinearShrink #-}

  unsafeLinearGrow :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix e
-> Sz ix -> m (MArray (PrimState m) P ix e)
unsafeLinearGrow (MPArray Sz ix
_ Int
o MutableByteArray (PrimState m)
ma) Sz ix
sz =
    forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray Sz ix
sz Int
o forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> m (MutableByteArray (PrimState m))
resizeMutableByteArray MutableByteArray (PrimState m)
ma ((Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) forall a. Num a => a -> a -> a
* forall a. Prim a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e))
  {-# INLINE unsafeLinearGrow #-}

instance (Prim e, Index ix) => Load P ix e where
  makeArray :: Comp -> Sz ix -> (ix -> e) -> Array P ix e
makeArray Comp
comp Sz ix
sz ix -> e
f = forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute (forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp Sz ix
sz ix -> e
f :: Array D ix e)
  {-# INLINE makeArray #-}
  makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array P ix e
makeArrayLinear !Comp
comp !Sz ix
sz Int -> e
f = forall a. IO a -> a
unsafePerformIO forall a b. (a -> b) -> a -> b
$ forall r ix e (m :: * -> *).
(MonadUnliftIO m, Manifest r e, Index ix) =>
Comp -> Sz ix -> (Int -> m e) -> m (Array r ix e)
generateArrayLinear Comp
comp Sz ix
sz (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> e
f)
  {-# INLINE makeArrayLinear #-}

  replicate :: Comp -> Sz ix -> e -> Array P ix e
replicate Comp
comp !Sz ix
sz !e
e = forall a. (forall s. ST s a) -> a
runST (forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> e -> m (MArray (PrimState m) r ix e)
newMArray Sz ix
sz e
e forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp)
  {-# INLINE replicate #-}

  iterArrayLinearST_ :: forall s.
Scheduler s () -> Array P ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ !Scheduler s ()
scheduler !Array P ix e
arr =
    forall s (m :: * -> *) b.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler s ()
scheduler (forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount Array P ix e
arr) (forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Array P ix e
arr)
  {-# INLINE iterArrayLinearST_ #-}

instance (Prim e, Index ix) => StrideLoad P ix e

instance (Prim e, Index ix) => Stream P ix e where
  toStream :: Array P ix e -> Steps Id e
toStream = forall r ix e (m :: * -> *).
(Monad m, Index ix, Source r e) =>
Array r ix e -> Steps m e
S.steps
  {-# INLINE toStream #-}
  toStreamIx :: Array P ix e -> Steps Id (ix, e)
toStreamIx = forall r ix e (m :: * -> *).
(Monad m, Index ix, Source r e) =>
Array r ix e -> Steps m (ix, e)
S.isteps
  {-# INLINE toStreamIx #-}

instance (Prim e, Num e) => FoldNumeric P e where
  unsafeDotProduct :: forall ix. Index ix => Array P ix e -> Array P ix e -> e
unsafeDotProduct = forall e ix r.
(Num e, Index ix, Source r e) =>
Array r ix e -> Array r ix e -> e
defaultUnsafeDotProduct
  {-# INLINE unsafeDotProduct #-}
  powerSumArray :: forall ix. Index ix => Array P ix e -> Int -> e
powerSumArray = forall ix r e.
(Index ix, Source r e, Num e) =>
Array r ix e -> Int -> e
defaultPowerSumArray
  {-# INLINE powerSumArray #-}
  foldArray :: forall ix. Index ix => (e -> e -> e) -> e -> Array P ix e -> e
foldArray = forall ix r e.
(Index ix, Source r e) =>
(e -> e -> e) -> e -> Array r ix e -> e
defaultFoldArray
  {-# INLINE foldArray #-}

instance (Prim e, Num e) => Numeric P e where
  unsafeLiftArray :: forall ix. Index ix => (e -> e) -> Array P ix e -> Array P ix e
unsafeLiftArray = forall r ix e.
(Load r ix e, Source r e) =>
(e -> e) -> Array r ix e -> Array r ix e
defaultUnsafeLiftArray
  {-# INLINE unsafeLiftArray #-}
  unsafeLiftArray2 :: forall ix.
Index ix =>
(e -> e -> e) -> Array P ix e -> Array P ix e -> Array P ix e
unsafeLiftArray2 = forall r ix e.
(Load r ix e, Source r e) =>
(e -> e -> e) -> Array r ix e -> Array r ix e -> Array r ix e
defaultUnsafeLiftArray2
  {-# INLINE unsafeLiftArray2 #-}

instance (Prim e, Floating e) => NumericFloat P e

instance (Prim e, IsList (Array L ix e), Ragged L ix e) => IsList (Array P ix e) where
  type Item (Array P ix e) = Item (Array L ix e)
  fromList :: [Item (Array P ix e)] -> Array P ix e
fromList = forall r ix e.
(HasCallStack, Ragged L ix e, Manifest r e) =>
Comp -> [ListItem ix e] -> Array r ix e
A.fromLists' Comp
Seq
  {-# INLINE fromList #-}
  toList :: Array P ix e -> [Item (Array P ix e)]
toList = forall l. IsList l => l -> [Item l]
GHC.toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ix e r.
(Ragged L ix e, Shape r ix, Source r e) =>
Array r ix e -> Array L ix e
toListArray
  {-# INLINE toList #-}

elemsBA :: forall proxy e. Prim e => proxy e -> ByteArray -> Int
elemsBA :: forall (proxy :: * -> *) e. Prim e => proxy e -> ByteArray -> Int
elemsBA proxy e
_ ByteArray
a = ByteArray -> Int
sizeofByteArray ByteArray
a forall a. Integral a => a -> a -> a
`div` forall a. Prim a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e)
{-# INLINE elemsBA #-}

elemsMBA :: forall proxy e s. Prim e => proxy e -> MutableByteArray s -> Int
elemsMBA :: forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA proxy e
_ MutableByteArray s
a = forall s. MutableByteArray s -> Int
sizeofMutableByteArray MutableByteArray s
a forall a. Integral a => a -> a -> a
`div` forall a. Prim a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e)
{-# INLINE elemsMBA #-}

-- | /O(n)/ - Ensure that the size matches the internal `ByteArray`. If not make a copy of
-- the slice and return it as `ByteArray`
--
-- @since 0.2.1
toByteArray :: (Index ix, Prim e) => Array P ix e -> ByteArray
toByteArray :: forall ix e. (Index ix, Prim e) => Array P ix e -> ByteArray
toByteArray Array P ix e
arr = forall a. a -> Maybe a -> a
fromMaybe (forall ix e. Array P ix e -> ByteArray
unwrapByteArray forall a b. (a -> b) -> a -> b
$ forall r e ix.
(Manifest r e, Index ix) =>
Array r ix e -> Array r ix e
clone Array P ix e
arr) forall a b. (a -> b) -> a -> b
$ forall e ix (m :: * -> *).
(Prim e, Index ix, MonadThrow m) =>
Array P ix e -> m ByteArray
toByteArrayM Array P ix e
arr
{-# INLINE toByteArray #-}

-- | /O(1)/ - Extract the internal `ByteArray`. This will ignore any possible slicing that
-- has been applied to the array. Use `toByteArray` in order to preserve slicing or
-- `unwrapByteArrayOffset` to get ahold of the offset
--
-- @since 0.5.0
unwrapByteArray :: Array P ix e -> ByteArray
unwrapByteArray :: forall ix e. Array P ix e -> ByteArray
unwrapByteArray = forall ix e. Array P ix e -> ByteArray
pData
{-# INLINE unwrapByteArray #-}

-- | /O(1)/ - Extract potential linear offset into the underlying `ByteArray`, which can
-- also be extracted with `unwrapByteArray`.
--
-- @since 0.5.9
unwrapByteArrayOffset :: Array P ix e -> Int
unwrapByteArrayOffset :: forall ix e. Array P ix e -> Int
unwrapByteArrayOffset = forall ix e. Array P ix e -> Int
pOffset
{-# INLINE unwrapByteArrayOffset #-}

-- | /O(1)/ - Unwrap Ensure that the size matches the internal `ByteArray`.
--
-- @since 0.5.0
toByteArrayM :: (Prim e, Index ix, MonadThrow m) => Array P ix e -> m ByteArray
toByteArrayM :: forall e ix (m :: * -> *).
(Prim e, Index ix, MonadThrow m) =>
Array P ix e -> m ByteArray
toByteArrayM arr :: Array P ix e
arr@PArray{Sz ix
pSize :: Sz ix
pSize :: forall ix e. Array P ix e -> Sz ix
pSize, ByteArray
pData :: ByteArray
pData :: forall ix e. Array P ix e -> ByteArray
pData} = do
  ByteArray
pData forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements Sz ix
pSize (forall ix. Index ix => ix -> Sz ix
Sz (forall (proxy :: * -> *) e. Prim e => proxy e -> ByteArray -> Int
elemsBA Array P ix e
arr ByteArray
pData))
{-# INLINE toByteArrayM #-}

-- | /O(1)/ - Construct a primitive array from the `ByteArray`. Will return `Nothing` if
-- number of elements doesn't match.
--
-- @since 0.3.0
fromByteArrayM :: (MonadThrow m, Index ix, Prim e) => Comp -> Sz ix -> ByteArray -> m (Array P ix e)
fromByteArrayM :: forall (m :: * -> *) ix e.
(MonadThrow m, Index ix, Prim e) =>
Comp -> Sz ix -> ByteArray -> m (Array P ix e)
fromByteArrayM Comp
comp Sz ix
sz = forall (m :: * -> *) ix e.
(MonadThrow m, Index ix, Prim e) =>
Comp -> Sz ix -> Int -> ByteArray -> m (Array P ix e)
fromByteArrayOffsetM Comp
comp Sz ix
sz Int
0
{-# INLINE fromByteArrayM #-}

-- | /O(1)/ - Construct a primitive array from the `ByteArray`. Will return `Nothing` if
-- number of elements doesn't match.
--
-- @since 0.5.9
fromByteArrayOffsetM
  :: (MonadThrow m, Index ix, Prim e) => Comp -> Sz ix -> Int -> ByteArray -> m (Array P ix e)
fromByteArrayOffsetM :: forall (m :: * -> *) ix e.
(MonadThrow m, Index ix, Prim e) =>
Comp -> Sz ix -> Int -> ByteArray -> m (Array P ix e)
fromByteArrayOffsetM Comp
comp Sz ix
sz Int
off ByteArray
ba =
  Array P ix e
arr forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements Sz ix
sz (forall ix. ix -> Sz ix
SafeSz (forall (proxy :: * -> *) e. Prim e => proxy e -> ByteArray -> Int
elemsBA Array P ix e
arr ByteArray
ba forall a. Num a => a -> a -> a
- Int
off))
  where
    arr :: Array P ix e
arr = forall ix e. Comp -> Sz ix -> Int -> ByteArray -> Array P ix e
PArray Comp
comp Sz ix
sz Int
off ByteArray
ba
{-# INLINE fromByteArrayOffsetM #-}

-- | /O(1)/ - Construct a flat Array from `ByteArray`
--
-- @since 0.4.0
fromByteArray :: forall e. Prim e => Comp -> ByteArray -> Array P Ix1 e
fromByteArray :: forall e. Prim e => Comp -> ByteArray -> Array P Int e
fromByteArray Comp
comp ByteArray
ba = forall ix e. Comp -> Sz ix -> Int -> ByteArray -> Array P ix e
PArray Comp
comp (forall ix. ix -> Sz ix
SafeSz (forall (proxy :: * -> *) e. Prim e => proxy e -> ByteArray -> Int
elemsBA (forall {k} (t :: k). Proxy t
Proxy :: Proxy e) ByteArray
ba)) Int
0 ByteArray
ba
{-# INLINE fromByteArray #-}

-- | /O(1)/ - Extract the internal `MutableByteArray`. This will discard any possible
-- slicing that has been applied to the array.
--
-- @since 0.5.0
unwrapMutableByteArray :: MArray s P ix e -> MutableByteArray s
unwrapMutableByteArray :: forall s ix e. MArray s P ix e -> MutableByteArray s
unwrapMutableByteArray (MPArray Sz ix
_ Int
_ MutableByteArray s
mba) = MutableByteArray s
mba
{-# INLINE unwrapMutableByteArray #-}

-- | /O(1)/ - Extract the linear offset into underlying `MutableByteArray`, which can aslo
-- be extracted with `unwrapMutableByteArray`.
--
-- @since 0.5.9
unwrapMutableByteArrayOffset :: MArray s P ix e -> Int
unwrapMutableByteArrayOffset :: forall s ix e. MArray s P ix e -> Int
unwrapMutableByteArrayOffset (MPArray Sz ix
_ Int
off MutableByteArray s
_) = Int
off
{-# INLINE unwrapMutableByteArrayOffset #-}

-- | /O(n)/ - Try to cast a mutable array to `MutableByteArray`, if sizes do not match make
-- a copy. Returns `True` if an array was converted without a copy, in which case it means
-- that the source at the resulting array are still pointing to the same location in memory.
--
-- @since 0.5.0
toMutableByteArray
  :: forall ix e m
   . (Prim e, Index ix, PrimMonad m)
  => MArray (PrimState m) P ix e
  -> m (Bool, MutableByteArray (PrimState m))
toMutableByteArray :: forall ix e (m :: * -> *).
(Prim e, Index ix, PrimMonad m) =>
MArray (PrimState m) P ix e
-> m (Bool, MutableByteArray (PrimState m))
toMutableByteArray marr :: MArray (PrimState m) P ix e
marr@(MPArray Sz ix
sz Int
offset MutableByteArray (PrimState m)
mbas) =
  case forall ix e (m :: * -> *) s.
(Index ix, Prim e, MonadThrow m) =>
MArray s P ix e -> m (MutableByteArray s)
toMutableByteArrayM MArray (PrimState m) P ix e
marr of
    Just MutableByteArray (PrimState m)
mba -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
True, MutableByteArray (PrimState m)
mba)
    Maybe (MutableByteArray (PrimState m))
Nothing -> do
      let eSize :: Int
eSize = forall a. Prim a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: e)
          szBytes :: Int
szBytes = forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz forall a. Num a => a -> a -> a
* Int
eSize
      MutableByteArray (PrimState m)
mbad <- forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newPinnedByteArray Int
szBytes
      forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> MutableByteArray (PrimState m) -> Int -> Int -> m ()
copyMutableByteArray MutableByteArray (PrimState m)
mbad Int
0 MutableByteArray (PrimState m)
mbas (Int
offset forall a. Num a => a -> a -> a
* Int
eSize) Int
szBytes
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool
False, MutableByteArray (PrimState m)
mbad)
{-# INLINE toMutableByteArray #-}

-- | /O(1)/ - Extract the internal `MutableByteArray`.
--
-- @since 0.2.1
toMutableByteArrayM :: (Index ix, Prim e, MonadThrow m) => MArray s P ix e -> m (MutableByteArray s)
toMutableByteArrayM :: forall ix e (m :: * -> *) s.
(Index ix, Prim e, MonadThrow m) =>
MArray s P ix e -> m (MutableByteArray s)
toMutableByteArrayM marr :: MArray s P ix e
marr@(MPArray Sz ix
sz Int
_ MutableByteArray s
mba) =
  MutableByteArray s
mba forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements Sz ix
sz (forall ix. Index ix => ix -> Sz ix
Sz (forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray s P ix e
marr MutableByteArray s
mba))
{-# INLINE toMutableByteArrayM #-}

-- | /O(1)/ - Construct a primitive mutable array from the `MutableByteArray`. Will throw
-- `SizeElementsMismatchException` if number of elements doesn't match.
--
-- @since 0.3.0
fromMutableByteArrayM
  :: (MonadThrow m, Index ix, Prim e) => Sz ix -> MutableByteArray s -> m (MArray s P ix e)
fromMutableByteArrayM :: forall (m :: * -> *) ix e s.
(MonadThrow m, Index ix, Prim e) =>
Sz ix -> MutableByteArray s -> m (MArray s P ix e)
fromMutableByteArrayM Sz ix
sz = forall (m :: * -> *) ix e s.
(MonadThrow m, Index ix, Prim e) =>
Sz ix -> Int -> MutableByteArray s -> m (MArray s P ix e)
fromMutableByteArrayOffsetM Sz ix
sz Int
0
{-# INLINE fromMutableByteArrayM #-}

-- | /O(1)/ - Construct a primitive mutable array from the `MutableByteArray`. Will throw
-- `SizeElementsMismatchException` if number of elements doesn't match.
--
-- @since 0.5.9
fromMutableByteArrayOffsetM
  :: (MonadThrow m, Index ix, Prim e) => Sz ix -> Ix1 -> MutableByteArray s -> m (MArray s P ix e)
fromMutableByteArrayOffsetM :: forall (m :: * -> *) ix e s.
(MonadThrow m, Index ix, Prim e) =>
Sz ix -> Int -> MutableByteArray s -> m (MArray s P ix e)
fromMutableByteArrayOffsetM Sz ix
sz Int
off MutableByteArray s
mba =
  MArray s P ix e
marr forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements Sz ix
sz (forall ix. ix -> Sz ix
SafeSz (forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray s P ix e
marr MutableByteArray s
mba forall a. Num a => a -> a -> a
- Int
off))
  where
    marr :: MArray s P ix e
marr = forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray Sz ix
sz Int
off MutableByteArray s
mba
{-# INLINE fromMutableByteArrayOffsetM #-}

-- | /O(1)/ - Construct a flat Array from `MutableByteArray`
--
-- @since 0.4.0
fromMutableByteArray :: forall e s. Prim e => MutableByteArray s -> MArray s P Ix1 e
fromMutableByteArray :: forall e s. Prim e => MutableByteArray s -> MArray s P Int e
fromMutableByteArray MutableByteArray s
mba = forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray (forall ix. ix -> Sz ix
SafeSz (forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA (forall {k} (t :: k). Proxy t
Proxy :: Proxy e) MutableByteArray s
mba)) Int
0 MutableByteArray s
mba
{-# INLINE fromMutableByteArray #-}

-- | /O(1)/ - Cast a primitive array to a primitive vector.
--
-- @since 0.5.0
toPrimitiveVector :: Index ix => Array P ix e -> VP.Vector e
toPrimitiveVector :: forall ix e. Index ix => Array P ix e -> Vector e
toPrimitiveVector PArray{Sz ix
pSize :: Sz ix
pSize :: forall ix e. Array P ix e -> Sz ix
pSize, Int
pOffset :: Int
pOffset :: forall ix e. Array P ix e -> Int
pOffset, ByteArray
pData :: ByteArray
pData :: forall ix e. Array P ix e -> ByteArray
pData} = forall a. Int -> Int -> ByteArray -> Vector a
VP.Vector Int
pOffset (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
pSize) ByteArray
pData
{-# INLINE toPrimitiveVector #-}

-- | /O(1)/ - Cast a mutable primitive array to a mutable primitive vector.
--
-- @since 0.5.0
toPrimitiveMVector :: Index ix => MArray s P ix e -> MVP.MVector s e
toPrimitiveMVector :: forall ix s e. Index ix => MArray s P ix e -> MVector s e
toPrimitiveMVector (MPArray Sz ix
sz Int
offset MutableByteArray s
mba) = forall s a. Int -> Int -> MutableByteArray s -> MVector s a
MVP.MVector Int
offset (forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz) MutableByteArray s
mba
{-# INLINE toPrimitiveMVector #-}

-- | /O(1)/ - Cast a primitive vector to a primitive array.
--
-- @since 0.5.0
fromPrimitiveVector :: VP.Vector e -> Array P Ix1 e
fromPrimitiveVector :: forall e. Vector e -> Array P Int e
fromPrimitiveVector (VP.Vector Int
offset Int
len ByteArray
ba) =
  PArray{pComp :: Comp
pComp = Comp
Seq, pSize :: Sz1
pSize = forall ix. ix -> Sz ix
SafeSz Int
len, pOffset :: Int
pOffset = Int
offset, pData :: ByteArray
pData = ByteArray
ba}
{-# INLINE fromPrimitiveVector #-}

-- | /O(1)/ - Cast a mutable primitive vector to a mutable primitive array.
--
-- @since 0.5.0
fromPrimitiveMVector :: MVP.MVector s e -> MArray s P Ix1 e
fromPrimitiveMVector :: forall s e. MVector s e -> MArray s P Int e
fromPrimitiveMVector (MVP.MVector Int
offset Int
len MutableByteArray s
mba) = forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray (forall ix. ix -> Sz ix
SafeSz Int
len) Int
offset MutableByteArray s
mba
{-# INLINE fromPrimitiveMVector #-}

-- | Atomically read an `Int` element from the array
--
-- @since 0.3.0
unsafeAtomicReadIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> m Int
unsafeAtomicReadIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> m Int
unsafeAtomicReadIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) ix
ix =
  forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
    String
"P.unsafeAtomicReadIntArray"
    (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray (PrimState m) P ix Int
_mpa)
    ( \(MutableByteArray MutableByteArray# (PrimState m)
mba#) (I# Int#
i#) ->
        forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
          case forall d.
MutableByteArray# d -> Int# -> State# d -> (# State# d, Int# #)
atomicReadIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# State# (PrimState m)
s# of
            (# State# (PrimState m)
s'#, Int#
e# #) -> (# State# (PrimState m)
s'#, Int# -> Int
I# Int#
e# #)
    )
    MutableByteArray (PrimState m)
mba
    (Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
{-# INLINE unsafeAtomicReadIntArray #-}

-- | Atomically write an `Int` element int the array
--
-- @since 0.3.0
unsafeAtomicWriteIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m ()
unsafeAtomicWriteIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m ()
unsafeAtomicWriteIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) ix
ix (I# Int#
e#) =
  forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
    String
"P.unsafeAtomicWriteIntArray"
    (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray (PrimState m) P ix Int
_mpa)
    ( \(MutableByteArray MutableByteArray# (PrimState m)
mba#) (I# Int#
i#) ->
        forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (forall d.
MutableByteArray# d -> Int# -> Int# -> State# d -> State# d
atomicWriteIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# Int#
e#)
    )
    MutableByteArray (PrimState m)
mba
    (Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
{-# INLINE unsafeAtomicWriteIntArray #-}

-- | Atomically CAS an `Int` in the array. Returns the old value.
--
-- @since 0.3.0
unsafeCasIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m Int
unsafeCasIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m Int
unsafeCasIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) ix
ix (I# Int#
e#) (I# Int#
n#) =
  forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
    String
"P.unsafeCasIntArray"
    (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray (PrimState m) P ix Int
_mpa)
    ( \(MutableByteArray MutableByteArray# (PrimState m)
mba#) (I# Int#
i#) ->
        forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
          case forall d.
MutableByteArray# d
-> Int# -> Int# -> Int# -> State# d -> (# State# d, Int# #)
casIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# Int#
e# Int#
n# State# (PrimState m)
s# of
            (# State# (PrimState m)
s'#, Int#
o# #) -> (# State# (PrimState m)
s'#, Int# -> Int
I# Int#
o# #)
    )
    MutableByteArray (PrimState m)
mba
    (Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
{-# INLINE unsafeCasIntArray #-}

-- | Atomically modify an `Int` element of the array. Returns the old value.
--
-- @since 0.3.0
unsafeAtomicModifyIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m Int
unsafeAtomicModifyIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m Int
unsafeAtomicModifyIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) ix
ix Int -> Int
f =
  forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
    String
"P.unsafeAtomicModifyIntArray"
    (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray (PrimState m) P ix Int
_mpa)
    MutableByteArray (PrimState m) -> Int -> m Int
atomicModify
    MutableByteArray (PrimState m)
mba
    (Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
  where
    atomicModify :: MutableByteArray (PrimState m) -> Int -> m Int
atomicModify (MutableByteArray MutableByteArray# (PrimState m)
mba#) (I# Int#
i#) =
      let go :: State# (PrimState m) -> Int# -> (# State# (PrimState m), Int #)
go State# (PrimState m)
s# Int#
o# =
            let !(I# Int#
n#) = Int -> Int
f (Int# -> Int
I# Int#
o#)
             in case forall d.
MutableByteArray# d
-> Int# -> Int# -> Int# -> State# d -> (# State# d, Int# #)
casIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# Int#
o# Int#
n# State# (PrimState m)
s# of
                  (# State# (PrimState m)
s'#, Int#
o'# #) ->
                    case Int#
o# Int# -> Int# -> Int#
==# Int#
o'# of
                      Int#
0# -> State# (PrimState m) -> Int# -> (# State# (PrimState m), Int #)
go State# (PrimState m)
s# Int#
o'#
                      Int#
_ -> (# State# (PrimState m)
s'#, Int# -> Int
I# Int#
o# #)
       in forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
            case forall d.
MutableByteArray# d -> Int# -> State# d -> (# State# d, Int# #)
atomicReadIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# State# (PrimState m)
s# of
              (# State# (PrimState m)
s'#, Int#
o# #) -> State# (PrimState m) -> Int# -> (# State# (PrimState m), Int #)
go State# (PrimState m)
s'# Int#
o#
    {-# INLINE atomicModify #-}
{-# INLINE unsafeAtomicModifyIntArray #-}

-- | Atomically add to an `Int` element in the array. Returns the old value.
--
-- @since 0.3.0
unsafeAtomicAddIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicAddIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicAddIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) ix
ix (I# Int#
e#) =
  forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
    String
"P.unsafeAtomicAddIntArray"
    (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray (PrimState m) P ix Int
_mpa)
    ( \(MutableByteArray MutableByteArray# (PrimState m)
mba#) (I# Int#
i#) ->
        forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
          case forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchAddIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# Int#
e# State# (PrimState m)
s# of
            (# State# (PrimState m)
s'#, Int#
p# #) -> (# State# (PrimState m)
s'#, Int# -> Int
I# Int#
p# #)
    )
    MutableByteArray (PrimState m)
mba
    (Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
{-# INLINE unsafeAtomicAddIntArray #-}

-- | Atomically subtract from an `Int` element in the array. Returns the old value.
--
-- @since 0.3.0
unsafeAtomicSubIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicSubIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicSubIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) ix
ix (I# Int#
e#) =
  forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
    String
"P.unsafeAtomicSubIntArray"
    (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray (PrimState m) P ix Int
_mpa)
    ( \(MutableByteArray MutableByteArray# (PrimState m)
mba#) (I# Int#
i#) ->
        forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
          case forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchSubIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# Int#
e# State# (PrimState m)
s# of
            (# State# (PrimState m)
s'#, Int#
p# #) -> (# State# (PrimState m)
s'#, Int# -> Int
I# Int#
p# #)
    )
    MutableByteArray (PrimState m)
mba
    (Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
{-# INLINE unsafeAtomicSubIntArray #-}

-- | Atomically AND an `Int` element in the array. Returns the old value.
--
-- @since 0.3.0
unsafeAtomicAndIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicAndIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicAndIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) ix
ix (I# Int#
e#) =
  forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
    String
"P.unsafeAtomicAndIntArray"
    (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray (PrimState m) P ix Int
_mpa)
    ( \(MutableByteArray MutableByteArray# (PrimState m)
mba#) (I# Int#
i#) ->
        forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
          case forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchAndIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# Int#
e# State# (PrimState m)
s# of
            (# State# (PrimState m)
s'#, Int#
p# #) -> (# State# (PrimState m)
s'#, Int# -> Int
I# Int#
p# #)
    )
    MutableByteArray (PrimState m)
mba
    (Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
{-# INLINE unsafeAtomicAndIntArray #-}

-- | Atomically NAND an `Int` element in the array. Returns the old value.
--
-- @since 0.3.0
unsafeAtomicNandIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicNandIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicNandIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) ix
ix (I# Int#
e#) =
  forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
    String
"P.unsafeAtomicNandIntArray"
    (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray (PrimState m) P ix Int
_mpa)
    ( \(MutableByteArray MutableByteArray# (PrimState m)
mba#) (I# Int#
i#) ->
        forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
          case forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchNandIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# Int#
e# State# (PrimState m)
s# of
            (# State# (PrimState m)
s'#, Int#
p# #) -> (# State# (PrimState m)
s'#, Int# -> Int
I# Int#
p# #)
    )
    MutableByteArray (PrimState m)
mba
    (Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
{-# INLINE unsafeAtomicNandIntArray #-}

-- | Atomically OR an `Int` element in the array. Returns the old value.
--
-- @since 0.3.0
unsafeAtomicOrIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicOrIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicOrIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) ix
ix (I# Int#
e#) =
  forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
    String
"P.unsafeAtomicOrIntArray"
    (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray (PrimState m) P ix Int
_mpa)
    ( \(MutableByteArray MutableByteArray# (PrimState m)
mba#) (I# Int#
i#) ->
        forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
          case forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchOrIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# Int#
e# State# (PrimState m)
s# of
            (# State# (PrimState m)
s'#, Int#
p# #) -> (# State# (PrimState m)
s'#, Int# -> Int
I# Int#
p# #)
    )
    MutableByteArray (PrimState m)
mba
    (Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
{-# INLINE unsafeAtomicOrIntArray #-}

-- | Atomically XOR an `Int` element in the array. Returns the old value.
--
-- @since 0.3.0
unsafeAtomicXorIntArray
  :: (Index ix, PrimMonad m) => MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicXorIntArray :: forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicXorIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray Sz ix
sz Int
o MutableByteArray (PrimState m)
mba) ix
ix (I# Int#
e#) =
  forall a ix e.
String -> (a -> Sz ix) -> (a -> ix -> e) -> a -> ix -> e
indexAssert
    String
"P.unsafeAtomicXorIntArray"
    (forall ix. ix -> Sz ix
SafeSz forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray (PrimState m) P ix Int
_mpa)
    ( \(MutableByteArray MutableByteArray# (PrimState m)
mba#) (I# Int#
i#) ->
        forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
          case forall d.
MutableByteArray# d
-> Int# -> Int# -> State# d -> (# State# d, Int# #)
fetchXorIntArray# MutableByteArray# (PrimState m)
mba# Int#
i# Int#
e# State# (PrimState m)
s# of
            (# State# (PrimState m)
s'#, Int#
p# #) -> (# State# (PrimState m)
s'#, Int# -> Int
I# Int#
p# #)
    )
    MutableByteArray (PrimState m)
mba
    (Int
o forall a. Num a => a -> a -> a
+ forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
{-# INLINE unsafeAtomicXorIntArray #-}