PrimVar.hs

-- SPDX-License-Identifier: CC0-1.0
{-
Each contributor licenses you to do everything with this work that
would otherwise infringe any patent claims they can license or become
able to license.
-}
{-|
`PrimVar` is internally a `PrimArray` of one element,
but with a convenient `MutVar`-like API.
-}
module Data.Primitive.PrimVar
  ( -- * Types
    PrimVar(..)
    -- * Allocation
  , newPrimVar
  , newPinnedPrimVar
  , newAlignedPinnedPrimVar
    -- * Access
  , readPrimVar
  , writePrimVar
  , modifyPrimVar
    -- * Information
  , samePrimVar
  , primVarContents
  , isPrimVarPinned
  ) where

import Data.Functor ( ($>) )
import Foreign (Ptr)

import Control.Monad.Primitive ( PrimMonad(PrimState) )
import Data.Primitive
  ( MutablePrimArray
  , Prim
  , newPrimArray
  , newPinnedPrimArray
  , newAlignedPinnedPrimArray
  , readPrimArray
  , writePrimArray
  , sameMutablePrimArray
  , mutablePrimArrayContents
  , isMutablePrimArrayPinned
  )

newtype PrimVar s a = PrimVar (MutablePrimArray s a)

newPrimVar
  :: (Prim a, PrimMonad m)
  => a -> m (PrimVar (PrimState m) a)
-- ^ Create a new `PrimVar` with the specified initial value.
newPrimVar a =
  newPrimArray 1 >>= \arr ->
    writePrimArray arr 0 a $> PrimVar arr

newPinnedPrimVar
  :: (Prim a, PrimMonad m)
  => a -> m (PrimVar (PrimState m) a)
-- ^ Create a new /pinned/ `PrimVar` with the specified initial value.
-- The garbage collector is guaranteed not to move it.
newPinnedPrimVar a =
  newPinnedPrimArray 1 >>= \arr ->
    writePrimArray arr 0 a $> PrimVar arr

newAlignedPinnedPrimVar
  :: (Prim a, PrimMonad m)
  => a -> m (PrimVar (PrimState m) a)
-- ^ Create a new /pinned/ `PrimVar` with the specified initial value
-- and with the alignment given by its Prim instance.
-- The garbage collector is guaranteed not to move it.
newAlignedPinnedPrimVar a =
  newAlignedPinnedPrimArray 1 >>= \arr ->
    writePrimArray arr 0 a $> PrimVar arr

readPrimVar
  :: (Prim a, PrimMonad m)
  => PrimVar (PrimState m) a -> m a
-- ^ Read the value of a `PrimVar`.
readPrimVar (PrimVar v) = readPrimArray v 0

writePrimVar
  :: (Prim a, PrimMonad m)
  => PrimVar (PrimState m) a -> a -> m ()
-- ^ Write a new value into a `PrimVar`.
writePrimVar (PrimVar v) = writePrimArray v 0

modifyPrimVar
  :: (Prim a, PrimMonad m)
  => PrimVar (PrimState m) a -> (a -> a) -> m ()
-- ^ Mutate the contents of a `PrimVar`.
modifyPrimVar v f =
  readPrimVar v >>= \a ->
    writePrimVar v (f a)

samePrimVar :: PrimVar s a -> PrimVar s a -> Bool
-- ^ Check if two `PrimVar`s refer to the same memory block.
samePrimVar (PrimVar v0) (PrimVar v1) =
  sameMutablePrimArray v0 v1

primVarContents :: PrimVar s a -> Ptr a
-- ^ Yield a pointer to a `PrimVar`'s data.
-- This operation is only safe on pinned `PrimVar`s
-- allocated by `newPinnedPrimVar` or `newAlignedPinnedPrimVar`.
primVarContents (PrimVar v) =
  mutablePrimArrayContents v

isPrimVarPinned :: PrimVar s a -> Bool
-- ^ Check whether or not a `PrimVar` is pinned.
isPrimVarPinned (PrimVar v) =
  isMutablePrimArrayPinned v