{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# 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-2021
-- 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.Monad
import Control.DeepSeq (NFData(..), deepseq)
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.Operations
import Data.Massiv.Core.List
import Data.Massiv.Vector.Stream as S (steps, isteps)
import Data.Maybe (fromMaybe)
import Data.Primitive (sizeOf, Prim)
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 Prelude hiding (mapM)
import System.IO.Unsafe (unsafePerformIO)

#include "massiv.h"

-- | Representation for `Prim`itive elements
data P = P deriving Int -> P -> ShowS
[P] -> ShowS
P -> String
(Int -> P -> ShowS) -> (P -> String) -> ([P] -> ShowS) -> Show P
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 { Array P ix e -> Comp
pComp   :: !Comp
                                    , Array P ix e -> Sz ix
pSize   :: !(Sz ix)
                                    , Array P ix e -> Int
pOffset :: {-# UNPACK #-} !Int
                                    , 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 = (Array P ix e -> Array P ix e) -> Int -> Array P ix e -> ShowS
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 Array P ix e -> Array P ix e
forall a. a -> a
id
  showList :: [Array P ix e] -> ShowS
showList = [Array P ix e] -> ShowS
forall arr. Show arr => [arr] -> ShowS
showArrayList

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

instance NFData ix => NFData (MArray s P ix e) where
  rnf :: MArray s P ix e -> ()
rnf (MPArray sz _o _mb) = Sz ix
sz Sz ix -> () -> ()
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
(==) = (e -> e -> Bool) -> 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 e -> e -> Bool
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 = (e -> e -> Ordering) -> Array P ix e -> Array P ix e -> Ordering
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 e -> e -> Ordering
forall a. Ord a => a -> a -> Ordering
compare
  {-# INLINE compare #-}

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


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

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

instance Prim e => Source P e where
  unsafeLinearIndex :: Array P ix e -> Int -> e
unsafeLinearIndex _arr :: Array P ix e
_arr@(PArray _ _ o a) Int
i =
    INDEX_CHECK("(Source P ix e).unsafeLinearIndex",
                SafeSz . elemsBA _arr, indexByteArray) a (i + o)
  {-# INLINE unsafeLinearIndex #-}

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

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


instance Prim e => Manifest P e where

  unsafeLinearIndexM :: Array P ix e -> Int -> e
unsafeLinearIndexM _pa :: Array P ix e
_pa@(PArray _ _sz o a) Int
i =
    INDEX_CHECK("(Manifest P ix e).unsafeLinearIndexM",
                const (Sz (totalElem _sz)), indexByteArray) a (i + o)
  {-# INLINE unsafeLinearIndexM #-}

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

  unsafeResizeMArray :: Sz ix' -> MArray s P ix e -> MArray s P ix' e
unsafeResizeMArray Sz ix'
sz (MPArray _ off marr) = Sz ix' -> Int -> MutableByteArray s -> MArray s P ix' e
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 :: Int -> Sz1 -> MArray s P ix e -> MVector s P e
unsafeLinearSliceMArray Int
i Sz1
k (MPArray _ o a) = Sz1 -> Int -> MutableByteArray s -> MVector s P e
forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray Sz1
k (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
o) MutableByteArray s
a
  {-# INLINE unsafeLinearSliceMArray #-}

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

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

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

  initialize :: MArray (PrimState m) P ix e -> m ()
initialize (MPArray sz o mba) =
    let k :: Int
k = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
* e -> Int
forall a. Prim a => a -> Int
sizeOf (e
forall a. HasCallStack => a
undefined :: e)
    in Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ MutableByteArray (PrimState m) -> Int -> Int -> Word8 -> m ()
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 :: MArray (PrimState m) P ix e -> Int -> m e
unsafeLinearRead _mpa :: MArray (PrimState m) P ix e
_mpa@(MPArray _sz o ma) Int
i =
    INDEX_CHECK("(Manifest P ix e).unsafeLinearRead",
                const (Sz (totalElem _sz)), readByteArray) ma (i + o)
  {-# INLINE unsafeLinearRead #-}

  unsafeLinearWrite :: MArray (PrimState m) P ix e -> Int -> e -> m ()
unsafeLinearWrite _mpa :: MArray (PrimState m) P ix e
_mpa@(MPArray _sz o ma) Int
i =
    INDEX_CHECK("(Manifest P ix e).unsafeLinearWrite",
                const (Sz (totalElem _sz)), writeByteArray) ma (i + o)
  {-# INLINE unsafeLinearWrite #-}

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

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

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

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


instance (Prim e, Index ix) => Load P ix e where
  makeArrayLinear :: Comp -> Sz ix -> (Int -> e) -> Array P ix e
makeArrayLinear !Comp
comp !Sz ix
sz Int -> e
f = IO (Array P ix e) -> Array P ix e
forall a. IO a -> a
unsafePerformIO (IO (Array P ix e) -> Array P ix e)
-> IO (Array P ix e) -> Array P ix e
forall a b. (a -> b) -> a -> b
$ Comp -> Sz ix -> (Int -> IO e) -> IO (Array P ix e)
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 (e -> IO e
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> IO e) -> (Int -> e) -> Int -> IO e
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 s. ST s (Array P ix e)) -> Array P ix e
forall a. (forall s. ST s a) -> a
runST (Sz ix -> e -> ST s (MArray (PrimState (ST s)) P ix e)
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 ST s (MArray s P ix e)
-> (MArray s P ix e -> ST s (Array P ix e)) -> ST s (Array P ix e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Comp -> MArray (PrimState (ST s)) P ix e -> ST s (Array P ix e)
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_ :: Scheduler s () -> Array P ix e -> (Int -> e -> ST s ()) -> ST s ()
iterArrayLinearST_ !Scheduler s ()
scheduler !Array P ix e
arr =
    Scheduler s ()
-> Int -> (Int -> e) -> (Int -> e -> ST s ()) -> ST s ()
forall s (m :: * -> *) b.
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> b) -> (Int -> b -> m ()) -> m ()
splitLinearlyWith_ Scheduler s ()
scheduler (Array P ix e -> Int
forall ix r e. (Index ix, Size r) => Array r ix e -> Int
elemsCount Array P ix e
arr) (Array P ix e -> Int -> e
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 = Array P ix e -> Steps Id e
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 = Array P ix e -> Steps Id (ix, e)
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 :: Array P ix e -> Array P ix e -> e
unsafeDotProduct = Array P ix e -> Array P ix e -> e
forall e ix r.
(Num e, Index ix, Source r e) =>
Array r ix e -> Array r ix e -> e
defaultUnsafeDotProduct
  {-# INLINE unsafeDotProduct #-}
  powerSumArray :: Array P ix e -> Int -> e
powerSumArray = Array P ix e -> Int -> e
forall ix r e.
(Index ix, Source r e, Num e) =>
Array r ix e -> Int -> e
defaultPowerSumArray
  {-# INLINE powerSumArray #-}
  foldArray :: (e -> e -> e) -> e -> Array P ix e -> e
foldArray = (e -> e -> e) -> e -> Array P ix e -> e
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 :: (e -> e) -> Array P ix e -> Array P ix e
unsafeLiftArray = (e -> e) -> Array P ix e -> Array P ix e
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 :: (e -> e -> e) -> Array P ix e -> Array P ix e -> Array P ix e
unsafeLiftArray2 = (e -> e -> e) -> Array P ix e -> Array P ix e -> Array P ix e
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 = Comp -> [ListItem ix e] -> Array P ix e
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 = Array L ix e -> [ListItem ix e]
forall l. IsList l => l -> [Item l]
GHC.toList (Array L ix e -> [ListItem ix e])
-> (Array P ix e -> Array L ix e)
-> Array P ix e
-> [ListItem ix e]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array P ix e -> Array L ix e
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 :: proxy e -> ByteArray -> Int
elemsBA proxy e
_ ByteArray
a = ByteArray -> Int
sizeofByteArray ByteArray
a Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` e -> Int
forall a. Prim a => a -> Int
sizeOf (e
forall a. HasCallStack => a
undefined :: e)
{-# INLINE elemsBA #-}


elemsMBA :: forall proxy e s . Prim e => proxy e -> MutableByteArray s -> Int
elemsMBA :: proxy e -> MutableByteArray s -> Int
elemsMBA proxy e
_ MutableByteArray s
a = MutableByteArray s -> Int
forall s. MutableByteArray s -> Int
sizeofMutableByteArray MutableByteArray s
a Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` e -> Int
forall a. Prim a => a -> Int
sizeOf (e
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 :: Array P ix e -> ByteArray
toByteArray Array P ix e
arr = ByteArray -> Maybe ByteArray -> ByteArray
forall a. a -> Maybe a -> a
fromMaybe (Array P ix e -> ByteArray
forall ix e. Array P ix e -> ByteArray
unwrapByteArray (Array P ix e -> ByteArray) -> Array P ix e -> ByteArray
forall a b. (a -> b) -> a -> b
$ Array P ix e -> Array P ix e
forall r e ix.
(Manifest r e, Index ix) =>
Array r ix e -> Array r ix e
clone Array P ix e
arr) (Maybe ByteArray -> ByteArray) -> Maybe ByteArray -> ByteArray
forall a b. (a -> b) -> a -> b
$ Array P ix e -> Maybe ByteArray
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 :: Array P ix e -> ByteArray
unwrapByteArray = Array P ix e -> ByteArray
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 :: Array P ix e -> Int
unwrapByteArrayOffset = Array P ix e -> Int
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 :: Array P ix e -> m ByteArray
toByteArrayM arr :: Array P ix e
arr@PArray {pSize, pData} = do
  ByteArray
pData ByteArray -> m () -> m ByteArray
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Sz ix -> Sz1 -> m ()
forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements Sz ix
pSize (Int -> Sz1
forall ix. Index ix => ix -> Sz ix
Sz (Array P ix e -> ByteArray -> Int
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 :: Comp -> Sz ix -> ByteArray -> m (Array P ix e)
fromByteArrayM Comp
comp Sz ix
sz = Comp -> Sz ix -> Int -> ByteArray -> m (Array P ix e)
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 :: 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 Array P ix e -> m () -> m (Array P ix e)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Sz ix -> Sz1 -> m ()
forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements Sz ix
sz (Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Array P ix e -> ByteArray -> Int
forall (proxy :: * -> *) e. Prim e => proxy e -> ByteArray -> Int
elemsBA Array P ix e
arr ByteArray
ba Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
off))
  where
    arr :: Array P ix e
arr = Comp -> Sz ix -> Int -> ByteArray -> Array P ix e
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 :: Comp -> ByteArray -> Array P Int e
fromByteArray Comp
comp ByteArray
ba = Comp -> Sz1 -> Int -> ByteArray -> Array P Int e
forall ix e. Comp -> Sz ix -> Int -> ByteArray -> Array P ix e
PArray Comp
comp (Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Proxy e -> ByteArray -> Int
forall (proxy :: * -> *) e. Prim e => proxy e -> ByteArray -> Int
elemsBA (Proxy e
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 :: MArray s P ix e -> MutableByteArray s
unwrapMutableByteArray (MPArray _ _ 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 :: MArray s P ix e -> Int
unwrapMutableByteArrayOffset (MPArray _ off _) = 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 :: MArray (PrimState m) P ix e
-> m (Bool, MutableByteArray (PrimState m))
toMutableByteArray marr :: MArray (PrimState m) P ix e
marr@(MPArray sz offset mbas) =
  case MArray (PrimState m) P ix e
-> Maybe (MutableByteArray (PrimState m))
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 -> (Bool, MutableByteArray (PrimState m))
-> m (Bool, MutableByteArray (PrimState m))
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 = e -> Int
forall a. Prim a => a -> Int
sizeOf (e
forall a. HasCallStack => a
undefined :: e)
          szBytes :: Int
szBytes = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem Sz ix
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
eSize
      MutableByteArray (PrimState m)
mbad <- Int -> m (MutableByteArray (PrimState m))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newPinnedByteArray Int
szBytes
      MutableByteArray (PrimState m)
-> Int -> MutableByteArray (PrimState m) -> Int -> Int -> m ()
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
eSize) Int
szBytes
      (Bool, MutableByteArray (PrimState m))
-> m (Bool, MutableByteArray (PrimState m))
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 :: MArray s P ix e -> m (MutableByteArray s)
toMutableByteArrayM marr :: MArray s P ix e
marr@(MPArray sz _ mba) =
  MutableByteArray s
mba MutableByteArray s -> m () -> m (MutableByteArray s)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Sz ix -> Sz1 -> m ()
forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements Sz ix
sz (Int -> Sz1
forall ix. Index ix => ix -> Sz ix
Sz (MArray s P ix e -> MutableByteArray s -> Int
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 :: Sz ix -> MutableByteArray s -> m (MArray s P ix e)
fromMutableByteArrayM Sz ix
sz = Sz ix -> Int -> MutableByteArray s -> m (MArray s P ix e)
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 :: 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 MArray s P ix e -> m () -> m (MArray s P ix e)
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Sz ix -> Sz1 -> m ()
forall (m :: * -> *) ix ix'.
(MonadThrow m, Index ix, Index ix') =>
Sz ix -> Sz ix' -> m ()
guardNumberOfElements Sz ix
sz (Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (MArray s P ix e -> MutableByteArray s -> Int
forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA MArray s P ix e
marr MutableByteArray s
mba Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
off))
  where
    marr :: MArray s P ix e
marr = Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
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 :: MutableByteArray s -> MArray s P Int e
fromMutableByteArray MutableByteArray s
mba = Sz1 -> Int -> MutableByteArray s -> MArray s P Int e
forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray (Int -> Sz1
forall ix. ix -> Sz ix
SafeSz (Proxy e -> MutableByteArray s -> Int
forall (proxy :: * -> *) e s.
Prim e =>
proxy e -> MutableByteArray s -> Int
elemsMBA (Proxy e
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 :: Array P ix e -> Vector e
toPrimitiveVector PArray {pSize, pOffset, pData} = Int -> Int -> ByteArray -> Vector e
forall a. Int -> Int -> ByteArray -> Vector a
VP.Vector Int
pOffset (Sz ix -> Int
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 :: MArray s P ix e -> MVector s e
toPrimitiveMVector (MPArray sz offset mba) = Int -> Int -> MutableByteArray s -> MVector s e
forall s a. Int -> Int -> MutableByteArray s -> MVector s a
MVP.MVector Int
offset (Sz ix -> Int
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 :: Vector e -> Array P Int e
fromPrimitiveVector (VP.Vector Int
offset Int
len ByteArray
ba) =
  PArray :: forall ix e. Comp -> Sz ix -> Int -> ByteArray -> Array P ix e
PArray {pComp :: Comp
pComp = Comp
Seq, pSize :: Sz1
pSize = Int -> Sz1
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 :: MVector s e -> MArray s P Int e
fromPrimitiveMVector (MVP.MVector Int
offset Int
len MutableByteArray s
mba) = Sz1 -> Int -> MutableByteArray s -> MArray s P Int e
forall s ix e.
Sz ix -> Int -> MutableByteArray s -> MArray s P ix e
MPArray (Int -> Sz1
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 :: MArray (PrimState m) P ix Int -> ix -> m Int
unsafeAtomicReadIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray sz o mba) ix
ix =
  INDEX_CHECK( "unsafeAtomicReadIntArray"
             , SafeSz . elemsMBA _mpa
             , \(MutableByteArray mba#) (I# i#) ->
                 primitive $ \s# ->
                 case atomicReadIntArray# mba# i# s# of
                   (# s'#, e# #) -> (# s'#, I# e# #))
  MutableByteArray (PrimState m)
mba
  (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
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 :: MArray (PrimState m) P ix Int -> ix -> Int -> m ()
unsafeAtomicWriteIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray sz o mba) ix
ix (I# Int#
e#) =
  INDEX_CHECK( "unsafeAtomicWriteIntArray"
             , SafeSz . elemsMBA _mpa
             , \(MutableByteArray mba#) (I# i#) ->
                 primitive_ (atomicWriteIntArray# mba# i# e#))
  MutableByteArray (PrimState m)
mba
  (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
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 :: MArray (PrimState m) P ix Int -> ix -> Int -> Int -> m Int
unsafeCasIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray sz o mba) ix
ix (I# Int#
e#) (I# Int#
n#) =
  INDEX_CHECK( "unsafeCasIntArray"
             , SafeSz . elemsMBA _mpa
             , \(MutableByteArray mba#) (I# i#) ->
                 primitive $ \s# ->
                 case casIntArray# mba# i# e# n# s# of
                   (# s'#, o# #) -> (# s'#, I# o# #))
  MutableByteArray (PrimState m)
mba
  (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
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 :: MArray (PrimState m) P ix Int -> ix -> (Int -> Int) -> m Int
unsafeAtomicModifyIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray sz o mba) ix
ix Int -> Int
f =
  INDEX_CHECK("unsafeAtomicModifyIntArray", SafeSz . elemsMBA _mpa, atomicModify)
  MutableByteArray (PrimState m)
mba
  (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
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 MutableByteArray# (PrimState m)
-> Int#
-> Int#
-> Int#
-> State# (PrimState m)
-> (# State# (PrimState m), Int# #)
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 (State# (PrimState m) -> (# State# (PrimState m), Int #)) -> m Int
forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive ((State# (PrimState m) -> (# State# (PrimState m), Int #))
 -> m Int)
-> (State# (PrimState m) -> (# State# (PrimState m), Int #))
-> m Int
forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
            case MutableByteArray# (PrimState m)
-> Int# -> State# (PrimState m) -> (# State# (PrimState m), Int# #)
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 :: MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicAddIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray sz o mba) ix
ix (I# Int#
e#) =
  INDEX_CHECK( "unsafeAtomicAddIntArray"
             , SafeSz . elemsMBA _mpa
             , \(MutableByteArray mba#) (I# i#) ->
                 primitive $ \s# ->
                 case fetchAddIntArray# mba# i# e# s# of
                   (# s'#, p# #) -> (# s'#, I# p# #))
  MutableByteArray (PrimState m)
mba
  (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
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 :: MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicSubIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray sz o mba) ix
ix (I# Int#
e#) =
  INDEX_CHECK( "unsafeAtomicSubIntArray"
             , SafeSz . elemsMBA _mpa
             , \(MutableByteArray mba#) (I# i#) ->
                 primitive $ \s# ->
                 case fetchSubIntArray# mba# i# e# s# of
                   (# s'#, p# #) -> (# s'#, I# p# #))
  MutableByteArray (PrimState m)
mba
  (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
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 :: MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicAndIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray sz o mba) ix
ix (I# Int#
e#) =
  INDEX_CHECK( "unsafeAtomicAndIntArray"
             , SafeSz . elemsMBA _mpa
             , \(MutableByteArray mba#) (I# i#) ->
                 primitive $ \s# ->
                 case fetchAndIntArray# mba# i# e# s# of
                   (# s'#, p# #) -> (# s'#, I# p# #))
  MutableByteArray (PrimState m)
mba
  (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
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 :: MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicNandIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray sz o mba) ix
ix (I# Int#
e#) =
  INDEX_CHECK( "unsafeAtomicNandIntArray"
             , SafeSz . elemsMBA _mpa
             , \(MutableByteArray mba#) (I# i#) ->
                 primitive $ \s# ->
                 case fetchNandIntArray# mba# i# e# s# of
                   (# s'#, p# #) -> (# s'#, I# p# #))
  MutableByteArray (PrimState m)
mba
  (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
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 :: MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicOrIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray sz o mba) ix
ix (I# Int#
e#) =
  INDEX_CHECK( "unsafeAtomicOrIntArray"
             , SafeSz . elemsMBA _mpa
             , \(MutableByteArray mba#) (I# i#) ->
                 primitive $ \s# ->
                 case fetchOrIntArray# mba# i# e# s# of
                   (# s'#, p# #) -> (# s'#, I# p# #))
  MutableByteArray (PrimState m)
mba
  (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
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 :: MArray (PrimState m) P ix Int -> ix -> Int -> m Int
unsafeAtomicXorIntArray _mpa :: MArray (PrimState m) P ix Int
_mpa@(MPArray sz o mba) ix
ix (I# Int#
e#) =
  INDEX_CHECK( "unsafeAtomicXorIntArray"
             , SafeSz . elemsMBA _mpa
             , \(MutableByteArray mba#) (I# i#) ->
                 primitive $ \s# ->
                 case fetchXorIntArray# mba# i# e# s# of
                   (# s'#, p# #) -> (# s'#, I# p# #))
  MutableByteArray (PrimState m)
mba
  (Int
o Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Sz ix -> ix -> Int
forall ix. Index ix => Sz ix -> ix -> Int
toLinearIndex Sz ix
sz ix
ix)
{-# INLINE unsafeAtomicXorIntArray #-}