{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UnliftedFFITypes #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UnboxedTuples #-}
-- |
-- Module      : Data.Primitive.PVar.Internal
-- Copyright   : (c) Alexey Kuleshevich 2020
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Primitive.PVar.Internal
  ( PVar(..)
  , newPVar
  , newPinnedPVar
  , newAlignedPinnedPVar
  , rawPVar
  , rawPinnedPVar
  , rawAlignedPinnedPVar
  , rawStorablePVar
  , rawAlignedStorablePVar
  , unsafeToPtrPVar
  , runWithPokedPtr
  , peekPrim
  , pokePrim
  , readPVar
  , writePVar
  , isPinnedPVar
  , sizeOfPVar
  , sizeOfPVar#
  , alignmentPVar
  , alignmentPVar#
  , unI#
  -- * Atomic operations
  , atomicModifyIntArray#
  , atomicModifyIntPVar
  , atomicModifyIntArray_#
  , atomicModifyIntPVar_
  -- * Re-exports
  , isByteArrayPinned#
  , isMutableByteArrayPinned#
  , sizeOf
  , alignment
  )
  where

import Control.DeepSeq
import Control.Monad.Primitive (MonadPrim, primitive, primitive_,
                                touch, unsafePrimToPrim)
import Data.Primitive.Types
import qualified Foreign.Storable as S
import GHC.Exts

#if !MIN_VERSION_primitive(0,6,3)
import Data.Primitive (sizeOf, alignment)
#endif

-- | Mutable variable with primitive value.
--
-- @since 0.1.0
data PVar a s = PVar (MutableByteArray# s)

-- | @`S.poke`+`S.peek`@ will result in a new copy of a `PVar`
instance Prim a => S.Storable (PVar a RealWorld) where
  sizeOf :: PVar a RealWorld -> Int
sizeOf = PVar a RealWorld -> Int
forall a s. Prim a => PVar a s -> Int
sizeOfPVar
  {-# INLINE sizeOf #-}
  alignment :: PVar a RealWorld -> Int
alignment = PVar a RealWorld -> Int
forall a s. Prim a => PVar a s -> Int
alignmentPVar
  {-# INLINE alignment #-}
  peekElemOff :: Ptr (PVar a RealWorld) -> Int -> IO (PVar a RealWorld)
peekElemOff (Ptr Addr#
addr#) (I# Int#
i#) = do
    a
a <- (State# (PrimState IO) -> (# State# (PrimState IO), a #)) -> IO a
forall a.
(State# (PrimState IO) -> (# State# (PrimState IO), a #)) -> IO a
forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (Addr# -> Int# -> State# RealWorld -> (# State# RealWorld, a #)
forall s. Addr# -> Int# -> State# s -> (# State# s, a #)
forall a s.
Prim a =>
Addr# -> Int# -> State# s -> (# State# s, a #)
readOffAddr# Addr#
addr# Int#
i#)
    a -> IO (PVar a RealWorld)
forall s (m :: * -> *) a.
(MonadPrim s m, Prim a) =>
a -> m (PVar a s)
newAlignedPinnedPVar a
a
  {-# INLINE peekElemOff #-}
  pokeElemOff :: Ptr (PVar a RealWorld) -> Int -> PVar a RealWorld -> IO ()
pokeElemOff (Ptr Addr#
addr#) (I# Int#
i#) PVar a RealWorld
pvar = do
    a
a <- PVar a RealWorld -> IO a
forall s (m :: * -> *) a.
(MonadPrim s m, Prim a) =>
PVar a s -> m a
readPVar PVar a RealWorld
pvar
    (State# (PrimState IO) -> State# (PrimState IO)) -> IO ()
forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (Addr# -> Int# -> a -> State# RealWorld -> State# RealWorld
forall s. Addr# -> Int# -> a -> State# s -> State# s
forall a s. Prim a => Addr# -> Int# -> a -> State# s -> State# s
writeOffAddr# Addr#
addr# Int#
i# a
a)
  {-# INLINE pokeElemOff #-}

-- | Values are already written into `PVar` in NF, this instance is trivial.
instance NFData (PVar a s) where
  rnf :: PVar a s -> ()
rnf (PVar MutableByteArray# s
_) = ()

-- | Create a mutable variable in unpinned memory (i.e. GC can move it) with an initial
-- value. This is a prefered way to create a mutable variable, since it will not
-- contribute to memory fragmentation. For pinned memory versions see `newPinnedPVar` and
-- `newAlignedPinnedPVar`
--
-- @since 0.1.0
newPVar :: (MonadPrim s m, Prim a) => a -> m (PVar a s)
newPVar :: forall s (m :: * -> *) a.
(MonadPrim s m, Prim a) =>
a -> m (PVar a s)
newPVar a
v = do
  PVar a s
pvar <- m (PVar a s)
forall a (m :: * -> *) s. (MonadPrim s m, Prim a) => m (PVar a s)
rawPVar
  PVar a s -> a -> m ()
forall s (m :: * -> *) a.
(MonadPrim s m, Prim a) =>
PVar a s -> a -> m ()
writePVar PVar a s
pvar a
v
  PVar a s -> m (PVar a s)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return PVar a s
pvar
{-# INLINE newPVar #-}

-- | Create a mutable variable in unpinned and unititialized memory
--
-- @since 0.1.0
rawPVar ::
     forall a m s. (MonadPrim s m, Prim a)
  => m (PVar a s)
rawPVar :: forall a (m :: * -> *) s. (MonadPrim s m, Prim a) => m (PVar a s)
rawPVar =
  (State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
-> m (PVar a s)
forall a.
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive ((State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
 -> m (PVar a s))
-> (State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
-> m (PVar a s)
forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
    case Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# (a -> Int#
forall a. Prim a => a -> Int#
sizeOf# (a
forall a. HasCallStack => a
undefined :: a)) State# s
State# (PrimState m)
s# of
      (# State# s
s'#, MutableByteArray# s
mba# #) -> (# State# s
State# (PrimState m)
s'#, MutableByteArray# s -> PVar a s
forall a s. MutableByteArray# s -> PVar a s
PVar MutableByteArray# s
mba# #)
{-# INLINE rawPVar #-}


-- | Create a mutable variable in pinned memory with an initial value.
--
-- @since 0.1.0
newPinnedPVar :: (MonadPrim s m, Prim a) => a -> m (PVar a s)
newPinnedPVar :: forall s (m :: * -> *) a.
(MonadPrim s m, Prim a) =>
a -> m (PVar a s)
newPinnedPVar a
v = do
  PVar a s
pvar <- m (PVar a s)
forall a (m :: * -> *) s. (MonadPrim s m, Prim a) => m (PVar a s)
rawPinnedPVar
  PVar a s -> a -> m ()
forall s (m :: * -> *) a.
(MonadPrim s m, Prim a) =>
PVar a s -> a -> m ()
writePVar PVar a s
pvar a
v
  PVar a s -> m (PVar a s)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return PVar a s
pvar
{-# INLINE newPinnedPVar #-}

-- | Create a mutable variable in pinned memory with uninitialized memory.
--
-- @since 0.1.0
rawPinnedPVar ::
     forall a m s. (MonadPrim s m, Prim a)
  => m (PVar a s)
rawPinnedPVar :: forall a (m :: * -> *) s. (MonadPrim s m, Prim a) => m (PVar a s)
rawPinnedPVar =
  (State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
-> m (PVar a s)
forall a.
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive ((State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
 -> m (PVar a s))
-> (State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
-> m (PVar a s)
forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
    case Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newPinnedByteArray# (a -> Int#
forall a. Prim a => a -> Int#
sizeOf# (a
forall a. HasCallStack => a
undefined :: a)) State# s
State# (PrimState m)
s# of
      (# State# s
s'#, MutableByteArray# s
mba# #) -> (# State# s
State# (PrimState m)
s'#, MutableByteArray# s -> PVar a s
forall a s. MutableByteArray# s -> PVar a s
PVar MutableByteArray# s
mba# #)
{-# INLINE rawPinnedPVar #-}


-- | Create a mutable variable in pinned memory with an initial value and aligned
-- according to its `Data.Primitive.Types.alignment`
--
-- @since 0.1.0
newAlignedPinnedPVar :: (MonadPrim s m, Prim a) => a -> m (PVar a s)
newAlignedPinnedPVar :: forall s (m :: * -> *) a.
(MonadPrim s m, Prim a) =>
a -> m (PVar a s)
newAlignedPinnedPVar a
v = do
  PVar a s
pvar <- m (PVar a s)
forall a (m :: * -> *) s. (MonadPrim s m, Prim a) => m (PVar a s)
rawAlignedPinnedPVar
  PVar a s -> a -> m ()
forall s (m :: * -> *) a.
(MonadPrim s m, Prim a) =>
PVar a s -> a -> m ()
writePVar PVar a s
pvar a
v
  PVar a s -> m (PVar a s)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return PVar a s
pvar
{-# INLINE newAlignedPinnedPVar #-}


-- | Create a mutable variable in pinned uninitialized memory.
--
-- @since 0.1.0
rawAlignedPinnedPVar ::
     forall a m s. (MonadPrim s m, Prim a)
  => m (PVar a s)
rawAlignedPinnedPVar :: forall a (m :: * -> *) s. (MonadPrim s m, Prim a) => m (PVar a s)
rawAlignedPinnedPVar =
  let dummy :: a
dummy = a
forall a. HasCallStack => a
undefined :: a
   in (State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
-> m (PVar a s)
forall a.
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive ((State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
 -> m (PVar a s))
-> (State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
-> m (PVar a s)
forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
        case Int# -> Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# (a -> Int#
forall a. Prim a => a -> Int#
sizeOf# a
dummy) (a -> Int#
forall a. Prim a => a -> Int#
alignment# a
dummy) State# s
State# (PrimState m)
s# of
          (# State# s
s'#, MutableByteArray# s
mba# #) -> (# State# s
State# (PrimState m)
s'#, MutableByteArray# s -> PVar a s
forall a s. MutableByteArray# s -> PVar a s
PVar MutableByteArray# s
mba# #)
{-# INLINE rawAlignedPinnedPVar #-}

-- | Create a mutable variable in pinned uninitialized memory using Storable interface for
-- getting the number of bytes for memory allocation.
--
-- @since 0.1.0
rawStorablePVar ::
     forall a m s. (MonadPrim s m, S.Storable a)
  => m (PVar a s)
rawStorablePVar :: forall a (m :: * -> *) s.
(MonadPrim s m, Storable a) =>
m (PVar a s)
rawStorablePVar =
  case a -> Int
forall a. Storable a => a -> Int
S.sizeOf (a
forall a. HasCallStack => a
undefined :: a) of
    I# Int#
size# ->
      (State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
-> m (PVar a s)
forall a.
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive ((State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
 -> m (PVar a s))
-> (State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
-> m (PVar a s)
forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
        case Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newPinnedByteArray# Int#
size# State# s
State# (PrimState m)
s# of
          (# State# s
s'#, MutableByteArray# s
mba# #) -> (# State# s
State# (PrimState m)
s'#, MutableByteArray# s -> PVar a s
forall a s. MutableByteArray# s -> PVar a s
PVar MutableByteArray# s
mba# #)
{-# INLINE rawStorablePVar #-}

-- | Create a mutable variable in pinned uninitialized memory using Storable interface for
-- getting the number of bytes for memory allocation and alignement.
--
-- @since 0.1.0
rawAlignedStorablePVar ::
     forall a m s. (MonadPrim s m, S.Storable a)
  => m (PVar a s)
rawAlignedStorablePVar :: forall a (m :: * -> *) s.
(MonadPrim s m, Storable a) =>
m (PVar a s)
rawAlignedStorablePVar =
  let dummy :: a
dummy = a
forall a. HasCallStack => a
undefined :: a
   in case a -> Int
forall a. Storable a => a -> Int
S.sizeOf a
dummy of
        I# Int#
size# ->
          case a -> Int
forall a. Storable a => a -> Int
S.alignment a
dummy of
            I# Int#
align# ->
              (State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
-> m (PVar a s)
forall a.
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive ((State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
 -> m (PVar a s))
-> (State# (PrimState m) -> (# State# (PrimState m), PVar a s #))
-> m (PVar a s)
forall a b. (a -> b) -> a -> b
$ \State# (PrimState m)
s# ->
                case Int# -> Int# -> State# s -> (# State# s, MutableByteArray# s #)
forall d.
Int# -> Int# -> State# d -> (# State# d, MutableByteArray# d #)
newAlignedPinnedByteArray# Int#
size# Int#
align# State# s
State# (PrimState m)
s# of
                  (# State# s
s'#, MutableByteArray# s
mba# #) -> (# State# s
State# (PrimState m)
s'#, MutableByteArray# s -> PVar a s
forall a s. MutableByteArray# s -> PVar a s
PVar MutableByteArray# s
mba# #)
{-# INLINE rawAlignedStorablePVar #-}


-- | Get the address to the contents. This is highly unsafe, espcially if memory is not pinned
--
-- @since 0.1.0
unsafeToPtrPVar :: PVar a s -> Ptr a
unsafeToPtrPVar :: forall a s. PVar a s -> Ptr a
unsafeToPtrPVar (PVar MutableByteArray# s
mba#) = Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr (ByteArray# -> Addr#
byteArrayContents# (MutableByteArray# s -> ByteArray#
forall a b. a -> b
unsafeCoerce# MutableByteArray# s
mba#))
{-# INLINE unsafeToPtrPVar #-}

-- helper that filles the PVar before running the action
runWithPokedPtr ::
     (S.Storable a, MonadPrim s m)
  => PVar a s
  -> a
  -> (PVar a s -> Ptr a -> m b)
  -> m b
runWithPokedPtr :: forall a s (m :: * -> *) b.
(Storable a, MonadPrim s m) =>
PVar a s -> a -> (PVar a s -> Ptr a -> m b) -> m b
runWithPokedPtr PVar a s
pvar a
a PVar a s -> Ptr a -> m b
f = do
  let ptr :: Ptr a
ptr = PVar a s -> Ptr a
forall a s. PVar a s -> Ptr a
unsafeToPtrPVar PVar a s
pvar
  Ptr a -> a -> m ()
forall a s (m :: * -> *).
(Storable a, MonadPrim s m) =>
Ptr a -> a -> m ()
pokePrim Ptr a
ptr a
a
  b
r <- PVar a s -> Ptr a -> m b
f PVar a s
pvar Ptr a
ptr
  PVar a s -> m ()
forall (m :: * -> *) a. PrimMonad m => a -> m ()
touch PVar a s
pvar
  b -> m b
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return b
r
{-# INLINE runWithPokedPtr #-}


-- | Use `S.Storable` reading functionality inside the `PrimMonad`.
--
-- @since 0.1.0
peekPrim :: (S.Storable a, MonadPrim s m) => Ptr a -> m a
peekPrim :: forall a s (m :: * -> *).
(Storable a, MonadPrim s m) =>
Ptr a -> m a
peekPrim = IO a -> m a
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim (IO a -> m a) -> (Ptr a -> IO a) -> Ptr a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr a -> IO a
forall a. Storable a => Ptr a -> IO a
S.peek
{-# INLINE peekPrim #-}

-- | Use `S.Storable` wrting functionality inside the `PrimMonad`.
--
-- @since 0.1.0
pokePrim :: (S.Storable a, MonadPrim s m) => Ptr a -> a -> m ()
pokePrim :: forall a s (m :: * -> *).
(Storable a, MonadPrim s m) =>
Ptr a -> a -> m ()
pokePrim Ptr a
ptr = IO () -> m ()
forall (m1 :: * -> *) (m2 :: * -> *) a.
(PrimBase m1, PrimMonad m2) =>
m1 a -> m2 a
unsafePrimToPrim (IO () -> m ()) -> (a -> IO ()) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr a -> a -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
S.poke Ptr a
ptr
{-# INLINE pokePrim #-}

-- | Read a value from a mutable variable
--
-- @since 0.1.0
readPVar :: (MonadPrim s m, Prim a) => PVar a s -> m a
readPVar :: forall s (m :: * -> *) a.
(MonadPrim s m, Prim a) =>
PVar a s -> m a
readPVar (PVar MutableByteArray# s
mba#) = (State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall a.
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (MutableByteArray# s -> Int# -> State# s -> (# State# s, a #)
forall s.
MutableByteArray# s -> Int# -> State# s -> (# State# s, a #)
forall a s.
Prim a =>
MutableByteArray# s -> Int# -> State# s -> (# State# s, a #)
readByteArray# MutableByteArray# s
mba# Int#
0#)
{-# INLINE readPVar #-}

-- | Write a value into a mutable variable
--
-- @since 0.1.0
writePVar :: (MonadPrim s m, Prim a) => PVar a s -> a -> m ()
writePVar :: forall s (m :: * -> *) a.
(MonadPrim s m, Prim a) =>
PVar a s -> a -> m ()
writePVar (PVar MutableByteArray# s
mba#) a
v = (State# (PrimState m) -> State# (PrimState m)) -> m ()
forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (MutableByteArray# s -> Int# -> a -> State# s -> State# s
forall s. MutableByteArray# s -> Int# -> a -> State# s -> State# s
forall a s.
Prim a =>
MutableByteArray# s -> Int# -> a -> State# s -> State# s
writeByteArray# MutableByteArray# s
mba# Int#
0# a
v)
{-# INLINE writePVar #-}

-- | Get the size of the mutable variable in bytes as an unpacked integer
--
-- @since 0.1.0
sizeOfPVar# :: forall a s. Prim a => PVar a s -> Int#
sizeOfPVar# :: forall a s. Prim a => PVar a s -> Int#
sizeOfPVar# PVar a s
_ = a -> Int#
forall a. Prim a => a -> Int#
sizeOf# (a
forall a. HasCallStack => a
undefined :: a)
{-# INLINE sizeOfPVar# #-}

-- | Get the alignment of the mutable variable in bytes as an unpacked integer
--
-- @since 0.1.0
alignmentPVar# :: forall a s. Prim a => PVar a s -> Int#
alignmentPVar# :: forall a s. Prim a => PVar a s -> Int#
alignmentPVar# PVar a s
_ = a -> Int#
forall a. Prim a => a -> Int#
alignment# (a
forall a. HasCallStack => a
undefined :: a)
{-# INLINE alignmentPVar# #-}


-- | Size in bytes of a value stored inside the mutable variable. `PVar` itself is neither
-- accessed nor evaluated.
--
-- @since 0.1.0
sizeOfPVar :: Prim a => PVar a s -> Int
sizeOfPVar :: forall a s. Prim a => PVar a s -> Int
sizeOfPVar PVar a s
pvar = Int# -> Int
I# (PVar a s -> Int#
forall a s. Prim a => PVar a s -> Int#
sizeOfPVar# PVar a s
pvar)
{-# INLINE sizeOfPVar #-}

-- | Alignment in bytes of the value stored inside of the mutable variable. `PVar` itself is
-- neither accessed nor evaluated.
--
-- @since 0.1.0
alignmentPVar :: Prim a => PVar a s -> Int
alignmentPVar :: forall a s. Prim a => PVar a s -> Int
alignmentPVar PVar a s
pvar = Int# -> Int
I# (PVar a s -> Int#
forall a s. Prim a => PVar a s -> Int#
alignmentPVar# PVar a s
pvar)
{-# INLINE alignmentPVar #-}

-- | Unwrap the primitive `Int`
--
-- @since 0.1.0
unI# :: Int -> Int#
unI# :: Int -> Int#
unI# (I# Int#
i#) = Int#
i#
{-# INLINE unI# #-}



-- | Check if `PVar` is backed by pinned memory or not
--
-- @since 0.1.0
isPinnedPVar :: PVar a s -> Bool
isPinnedPVar :: forall a s. PVar a s -> Bool
isPinnedPVar (PVar MutableByteArray# s
mba#) = Int# -> Bool
isTrue# (MutableByteArray# s -> Int#
forall d. MutableByteArray# d -> Int#
isMutableByteArrayPinned# MutableByteArray# s
mba#)
{-# INLINE isPinnedPVar #-}


-- | Using `casIntArray#` perform atomic modification of an integer element in a
-- `MutableByteArray#`. Implies a full memory barrier.
--
-- @since 0.1.0
atomicModifyIntArray# ::
     MutableByteArray# d -- ^ Array to be mutated
  -> Int# -- ^ Index in number of `Int#` elements into the `MutableByteArray#`
  -> (Int# -> (# Int#, b #)) -- ^ Function to be applied atomically to the element
  -> State# d -- ^ Starting state
  -> (# State# d, b #)
atomicModifyIntArray# :: forall d b.
MutableByteArray# d
-> Int# -> (Int# -> (# Int#, b #)) -> State# d -> (# State# d, b #)
atomicModifyIntArray# MutableByteArray# d
mba# Int#
i# Int# -> (# Int#, b #)
f State# d
s0# =
  let go :: State# d -> Int# -> (# State# d, b #)
go State# d
s# Int#
o# =
        case Int# -> (# Int#, b #)
f Int#
o# of
          (# Int#
n#, b
artifact #) ->
            case MutableByteArray# d
-> Int# -> Int# -> Int# -> State# d -> (# State# d, Int# #)
forall d.
MutableByteArray# d
-> Int# -> Int# -> Int# -> State# d -> (# State# d, Int# #)
casIntArray# MutableByteArray# d
mba# Int#
i# Int#
o# Int#
n# State# d
s# of
              (# State# d
s'#, Int#
o'# #) ->
                case Int#
o# Int# -> Int# -> Int#
==# Int#
o'# of
                  Int#
0# -> State# d -> Int# -> (# State# d, b #)
go State# d
s# Int#
o'#
                  Int#
_  -> b -> State# d -> (# State# d, b #)
forall a d. a -> State# d -> (# State# d, a #)
seq# b
artifact State# d
s'#
   in case MutableByteArray# d -> Int# -> State# d -> (# State# d, Int# #)
forall d.
MutableByteArray# d -> Int# -> State# d -> (# State# d, Int# #)
atomicReadIntArray# MutableByteArray# d
mba# Int#
i# State# d
s0# of
        (# State# d
s'#, Int#
o# #) -> State# d -> Int# -> (# State# d, b #)
go State# d
s'# Int#
o#
{-# INLINE atomicModifyIntArray# #-}



-- | Apply a function to an integer element of a `PVar` atomically. Implies a full memory
-- barrier.
--
-- @since 0.1.0
atomicModifyIntPVar ::
     MonadPrim s m => PVar Int s -> (Int -> (Int, a)) -> m a
atomicModifyIntPVar :: forall s (m :: * -> *) a.
MonadPrim s m =>
PVar Int s -> (Int -> (Int, a)) -> m a
atomicModifyIntPVar (PVar MutableByteArray# s
mba#) Int -> (Int, a)
f = (State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall a.
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive (MutableByteArray# s
-> Int# -> (Int# -> (# Int#, a #)) -> State# s -> (# State# s, a #)
forall d b.
MutableByteArray# d
-> Int# -> (Int# -> (# Int#, b #)) -> State# d -> (# State# d, b #)
atomicModifyIntArray# MutableByteArray# s
mba# Int#
0# Int# -> (# Int#, a #)
g)
  where
    g :: Int# -> (# Int#, a #)
g Int#
i# =
      case Int -> (Int, a)
f (Int# -> Int
I# Int#
i#) of
        (I# Int#
o#, a
a) -> (# Int#
o#, a
a #)
    {-# INLINE g #-}
{-# INLINE atomicModifyIntPVar #-}


-- | Uses `casIntArray#` to perform atomic modification of an integer element in a
-- `MutableByteArray#`. Implies a full memory barrier.
--
-- @since 0.1.0
atomicModifyIntArray_# ::
     MutableByteArray# d -- ^ Array to be mutated
  -> Int# -- ^ Index in number of `Int#` elements into the `MutableByteArray#`
  -> (Int# -> Int#) -- ^ Function to be applied atomically to the element
  -> State# d -- ^ Starting state
  -> State# d
atomicModifyIntArray_# :: forall d.
MutableByteArray# d
-> Int# -> (Int# -> Int#) -> State# d -> State# d
atomicModifyIntArray_# MutableByteArray# d
mba# Int#
i# Int# -> Int#
f State# d
s0# =
  let go :: State# d -> Int# -> State# d
go State# d
s# Int#
o# =
        case MutableByteArray# d
-> Int# -> Int# -> Int# -> State# d -> (# State# d, Int# #)
forall d.
MutableByteArray# d
-> Int# -> Int# -> Int# -> State# d -> (# State# d, Int# #)
casIntArray# MutableByteArray# d
mba# Int#
i# Int#
o# (Int# -> Int#
f Int#
o#) State# d
s# of
          (# State# d
s'#, Int#
o'# #) ->
            case Int#
o# Int# -> Int# -> Int#
==# Int#
o'# of
              Int#
0# -> State# d -> Int# -> State# d
go State# d
s# Int#
o'#
              Int#
_  -> State# d
s'#
   in case MutableByteArray# d -> Int# -> State# d -> (# State# d, Int# #)
forall d.
MutableByteArray# d -> Int# -> State# d -> (# State# d, Int# #)
atomicReadIntArray# MutableByteArray# d
mba# Int#
i# State# d
s0# of
        (# State# d
s'#, Int#
o# #) -> State# d -> Int# -> State# d
go State# d
s'# Int#
o#
{-# INLINE atomicModifyIntArray_# #-}


-- | Apply a function to an integer element of a `PVar` atomically. Returns the old
-- value. Implies a full memory barrier.
--
-- @since 0.1.0
atomicModifyIntPVar_ ::
     MonadPrim s m => PVar Int s -> (Int -> Int) -> m ()
atomicModifyIntPVar_ :: forall s (m :: * -> *).
MonadPrim s m =>
PVar Int s -> (Int -> Int) -> m ()
atomicModifyIntPVar_ (PVar MutableByteArray# s
mba#) Int -> Int
f =
  (State# (PrimState m) -> State# (PrimState m)) -> m ()
forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (MutableByteArray# s
-> Int# -> (Int# -> Int#) -> State# s -> State# s
forall d.
MutableByteArray# d
-> Int# -> (Int# -> Int#) -> State# d -> State# d
atomicModifyIntArray_# MutableByteArray# s
mba# Int#
0# (\Int#
i# -> Int -> Int#
unI# (Int -> Int
f (Int# -> Int
I# Int#
i#))))
{-# INLINE atomicModifyIntPVar_ #-}


-- ghc-8.2 (i.e. 802 version) introduced these two functions, for versions before those
-- use their reimplementations in C:
# if __GLASGOW_HASKELL__ < 802
foreign import ccall unsafe "pvar.c pvar_is_byte_array_pinned"
  isByteArrayPinned# :: ByteArray# -> Int#
foreign import ccall unsafe "pvar.c pvar_is_byte_array_pinned"
  isMutableByteArrayPinned# :: MutableByteArray# s -> Int#
#endif