{-|
Module      :  Z.Data.PrimRef.PrimSTRef
Description :  Primitive ST Reference
Copyright   :  (c) Dong Han 2017~2019
License     :  BSD-style

Maintainer  :  winterland1989@gmail.com
Stability   :  experimental
Portability :  portable

This package provide fast unboxed references for ST monad. Unboxed reference is implemented using single cell MutableByteArray s to eliminate indirection overhead which MutVar# s a carry, on the otherhand unboxed reference only support limited type(instances of 'Prim' class).
-}


module Z.Data.PrimRef.PrimSTRef
  ( -- * Unboxed ST references
    PrimSTRef(..)
  , newPrimSTRef
  , readPrimSTRef
  , writePrimSTRef
  , modifyPrimSTRef
  ) where

import Data.Primitive.Types
import Data.Primitive.ByteArray
import GHC.ST
import GHC.Exts

-- | A mutable variable in the ST monad which can hold an instance of 'Prim'.
--
newtype PrimSTRef s a = PrimSTRef (MutableByteArray s)

-- | Build a new 'PrimSTRef'
--
newPrimSTRef :: Prim a => a -> ST s (PrimSTRef s a)
newPrimSTRef :: a -> ST s (PrimSTRef s a)
newPrimSTRef a
x = do
     MutableByteArray s
mba <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (Int# -> Int
I# (a -> Int#
forall a. Prim a => a -> Int#
sizeOf# a
x))
     MutableByteArray (PrimState (ST s)) -> Int -> a -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
mba Int
0 a
x
     PrimSTRef s a -> ST s (PrimSTRef s a)
forall (m :: * -> *) a. Monad m => a -> m a
return (MutableByteArray s -> PrimSTRef s a
forall k s (a :: k). MutableByteArray s -> PrimSTRef s a
PrimSTRef MutableByteArray s
mba)
{-# INLINE newPrimSTRef #-}

-- | Read the value of an 'PrimSTRef'
--
readPrimSTRef :: Prim a => PrimSTRef s a -> ST s a
readPrimSTRef :: PrimSTRef s a -> ST s a
readPrimSTRef (PrimSTRef MutableByteArray s
mba) = MutableByteArray (PrimState (ST s)) -> Int -> ST s a
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
mba Int
0
{-# INLINE readPrimSTRef #-}

-- | Write a new value into an 'PrimSTRef'
--
writePrimSTRef :: Prim a => PrimSTRef s a -> a -> ST s ()
writePrimSTRef :: PrimSTRef s a -> a -> ST s ()
writePrimSTRef (PrimSTRef MutableByteArray s
mba) a
x = MutableByteArray (PrimState (ST s)) -> Int -> a -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
mba Int
0 a
x
{-# INLINE writePrimSTRef #-}

-- | Mutate the contents of an 'PrimSTRef'.
--
--  Unboxed reference is always strict on the value it hold.
--
modifyPrimSTRef :: Prim a => PrimSTRef s a -> (a -> a) -> ST s ()
modifyPrimSTRef :: PrimSTRef s a -> (a -> a) -> ST s ()
modifyPrimSTRef PrimSTRef s a
ref a -> a
f = PrimSTRef s a -> ST s a
forall a s. Prim a => PrimSTRef s a -> ST s a
readPrimSTRef PrimSTRef s a
ref ST s a -> (a -> ST s ()) -> ST s ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= PrimSTRef s a -> a -> ST s ()
forall a s. Prim a => PrimSTRef s a -> a -> ST s ()
writePrimSTRef PrimSTRef s a
ref (a -> ST s ()) -> (a -> a) -> a -> ST s ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f
{-# INLINE modifyPrimSTRef #-}