module RIO.Prelude.URef
  ( URef
  , IOURef
  , newURef
  , readURef
  , writeURef
  , modifyURef
  ) where

import RIO.Prelude.Reexports
import qualified Data.Vector.Unboxed.Mutable as MUVector

-- | An unboxed reference. This works like an 'IORef', but the data is
-- stored in a bytearray instead of a heap object, avoiding
-- significant allocation overhead in some cases. For a concrete
-- example, see this Stack Overflow question:
-- <https://stackoverflow.com/questions/27261813/why-is-my-little-stref-int-require-allocating-gigabytes>.
--
-- The first parameter is the state token type, the same as would be
-- used for the 'ST' monad. If you're using an 'IO'-based monad, you
-- can use the convenience 'IOURef' type synonym instead.
--
-- @since 0.0.2.0
newtype URef s a = URef (MUVector.MVector s a)

-- | Helpful type synonym for using a 'URef' from an 'IO'-based stack.
--
-- @since 0.0.2.0
type IOURef = URef (PrimState IO)

-- | Create a new 'URef'
--
-- @since 0.0.2.0
newURef :: (PrimMonad m, Unbox a) => a -> m (URef (PrimState m) a)
newURef :: a -> m (URef (PrimState m) a)
newURef a
a = (MVector (PrimState m) a -> URef (PrimState m) a)
-> m (MVector (PrimState m) a) -> m (URef (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap MVector (PrimState m) a -> URef (PrimState m) a
forall s a. MVector s a -> URef s a
URef (Int -> a -> m (MVector (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
MUVector.replicate Int
1 a
a)

-- | Read the value in a 'URef'
--
-- @since 0.0.2.0
readURef :: (PrimMonad m, Unbox a) => URef (PrimState m) a -> m a
readURef :: URef (PrimState m) a -> m a
readURef (URef MVector (PrimState m) a
v) = MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> m a
MUVector.unsafeRead MVector (PrimState m) a
v Int
0

-- | Write a value into a 'URef'. Note that this action is strict, and
-- will force evalution of the value.
--
-- @since 0.0.2.0
writeURef :: (PrimMonad m, Unbox a) => URef (PrimState m) a -> a -> m ()
writeURef :: URef (PrimState m) a -> a -> m ()
writeURef (URef MVector (PrimState m) a
v) = MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
MVector (PrimState m) a -> Int -> a -> m ()
MUVector.unsafeWrite MVector (PrimState m) a
v Int
0

-- | Modify a value in a 'URef'. Note that this action is strict, and
-- will force evaluation of the result value.
--
-- @since 0.0.2.0
modifyURef :: (PrimMonad m, Unbox a) => URef (PrimState m) a -> (a -> a) -> m ()
modifyURef :: URef (PrimState m) a -> (a -> a) -> m ()
modifyURef URef (PrimState m) a
u a -> a
f = URef (PrimState m) a -> m a
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
URef (PrimState m) a -> m a
readURef URef (PrimState m) a
u m a -> (a -> m ()) -> m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= URef (PrimState m) a -> a -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
URef (PrimState m) a -> a -> m ()
writeURef URef (PrimState m) a
u (a -> m ()) -> (a -> a) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
f