{-# LANGUAGE MagicHash    #-}
{-# LANGUAGE TypeFamilies #-}
-- | Use @ByteArray@s containing one element for mutable references.
--
-- This is similar to @URef@s, but avoids the overhead of storing the length of
-- the @Vector@, which we statically know will always be 1. This allows it to
-- be a bit faster.
--
-- Motivated by: <http://stackoverflow.com/questions/27261813/why-is-my-little-stref-int-require-allocating-gigabytes> and ArrayRef.
module Data.Mutable.PRef
    ( -- * Types
      PRef
    , IOPRef
      -- * Functions
    , asPRef
    , MutableRef (..)
    ) where

import Data.Mutable.Class
import Data.Primitive           (sizeOf)
import Data.Primitive.ByteArray (MutableByteArray, newByteArray, readByteArray,
                                 writeByteArray)
import Data.Primitive.Types     (Prim)

-- | A primitive ByteArray reference, supporting any monad.
--
-- Since 0.2.0
newtype PRef s a = PRef (MutableByteArray s)

-- |
-- Since 0.2.0
asPRef :: PRef s a -> PRef s a
asPRef :: forall s a. PRef s a -> PRef s a
asPRef PRef s a
x = PRef s a
x
{-# INLINE asPRef #-}

-- | A primitive ByteArray IO reference.
type IOPRef = PRef (PrimState IO)

instance MutableContainer (PRef s a) where
    type MCState (PRef s a) = s
instance Prim a => MutableRef (PRef s a) where
    type RefElement (PRef s a) = a

    newRef :: forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (PRef s a)) =>
RefElement (PRef s a) -> m (PRef s a)
newRef RefElement (PRef s a)
x = do
        MutableByteArray s
ba <- Int -> m (MutableByteArray (PrimState m))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (a -> Int
forall a. Prim a => a -> Int
sizeOf (a -> Int) -> a -> Int
forall a b. (a -> b) -> a -> b
$! a
RefElement (PRef s a)
x)
        MutableByteArray (PrimState m) -> Int -> a -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState m)
ba Int
0 a
RefElement (PRef s a)
x
        PRef s a -> m (PRef s a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (PRef s a -> m (PRef s a)) -> PRef s a -> m (PRef s a)
forall a b. (a -> b) -> a -> b
$! MutableByteArray s -> PRef s a
forall s a. MutableByteArray s -> PRef s a
PRef MutableByteArray s
ba
    {-# INLINE newRef #-}

    readRef :: forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (PRef s a)) =>
PRef s a -> m (RefElement (PRef s a))
readRef (PRef MutableByteArray s
ba) = MutableByteArray (PrimState m) -> Int -> m a
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState m)
ba Int
0
    {-# INLINE readRef #-}

    writeRef :: forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (PRef s a)) =>
PRef s a -> RefElement (PRef s a) -> m ()
writeRef (PRef MutableByteArray s
ba) = MutableByteArray (PrimState m) -> Int -> a -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState m)
ba Int
0
    {-# INLINE writeRef #-}

    modifyRef :: forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (PRef s a)) =>
PRef s a
-> (RefElement (PRef s a) -> RefElement (PRef s a)) -> m ()
modifyRef (PRef MutableByteArray s
ba) RefElement (PRef s a) -> RefElement (PRef s a)
f = do
        a
x <- MutableByteArray (PrimState m) -> Int -> m a
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> m a
readByteArray MutableByteArray s
MutableByteArray (PrimState m)
ba Int
0
        MutableByteArray (PrimState m) -> Int -> a -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState m)
ba Int
0 (a -> m ()) -> a -> m ()
forall a b. (a -> b) -> a -> b
$! RefElement (PRef s a) -> RefElement (PRef s a)
f a
RefElement (PRef s a)
x
    {-# INLINE modifyRef #-}

    modifyRef' :: forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (PRef s a)) =>
PRef s a
-> (RefElement (PRef s a) -> RefElement (PRef s a)) -> m ()
modifyRef' = PRef s a
-> (RefElement (PRef s a) -> RefElement (PRef s a)) -> m ()
forall c (m :: * -> *).
(MutableRef c, PrimMonad m, PrimState m ~ MCState c) =>
c -> (RefElement c -> RefElement c) -> m ()
forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (PRef s a)) =>
PRef s a
-> (RefElement (PRef s a) -> RefElement (PRef s a)) -> m ()
modifyRef
    {-# INLINE modifyRef' #-}