{-# LANGUAGE TypeFamilies #-}
-- | Use 1-length mutable unboxed vectors for mutable references.
--
-- Motivated by: <http://stackoverflow.com/questions/27261813/why-is-my-little-stref-int-require-allocating-gigabytes> and ArrayRef.
module Data.Mutable.URef
    ( -- * Types
      URef
    , IOURef
      -- * Functions
    , asURef
    , MutableRef (..)
    ) where

import           Control.Monad               (liftM)
import           Data.Mutable.Class
import qualified Data.Vector.Generic.Mutable as V
import qualified Data.Vector.Unboxed.Mutable as VU

-- | An unboxed vector reference, supporting any monad.
--
-- Since 0.2.0
newtype URef s a = URef (VU.MVector s a)

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

-- | An unboxed IO vector reference.
type IOURef = URef (PrimState IO)

instance MutableContainer (URef s a) where
    type MCState (URef s a) = s
instance VU.Unbox a => MutableRef (URef s a) where
    type RefElement (URef s a) = a

    newRef :: forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (URef s a)) =>
RefElement (URef s a) -> m (URef s a)
newRef = (MVector s a -> URef s a) -> m (MVector s a) -> m (URef s a)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM MVector s a -> URef s a
forall s a. MVector s a -> URef s a
URef (m (MVector s a) -> m (URef s a))
-> (a -> m (MVector s a)) -> a -> m (URef s a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a -> m (MVector (PrimState m) a)
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
Int -> a -> m (v (PrimState m) a)
V.replicate Int
1
    {-# INLINE newRef#-}

    readRef :: forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (URef s a)) =>
URef s a -> m (RefElement (URef s a))
readRef (URef MVector s a
v) = MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
V.unsafeRead MVector s a
MVector (PrimState m) a
v Int
0
    {-# INLINE readRef #-}

    writeRef :: forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (URef s a)) =>
URef s a -> RefElement (URef s a) -> m ()
writeRef (URef MVector s a
v) = MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
V.unsafeWrite MVector s a
MVector (PrimState m) a
v Int
0
    {-# INLINE writeRef #-}

    modifyRef :: forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (URef s a)) =>
URef s a
-> (RefElement (URef s a) -> RefElement (URef s a)) -> m ()
modifyRef (URef MVector s a
v) RefElement (URef s a) -> RefElement (URef s a)
f = MVector (PrimState m) a -> Int -> m a
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
V.unsafeRead MVector s a
MVector (PrimState m) a
v Int
0 m a -> (a -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= MVector (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
V.unsafeWrite MVector s a
MVector (PrimState m) a
v Int
0 (a -> m ()) -> (a -> a) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
RefElement (URef s a) -> RefElement (URef s a)
f
    {-# INLINE modifyRef #-}

    modifyRef' :: forall (m :: * -> *).
(PrimMonad m, PrimState m ~ MCState (URef s a)) =>
URef s a
-> (RefElement (URef s a) -> RefElement (URef s a)) -> m ()
modifyRef' = URef s a
-> (RefElement (URef s a) -> RefElement (URef 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 (URef s a)) =>
URef s a
-> (RefElement (URef s a) -> RefElement (URef s a)) -> m ()
modifyRef
    {-# INLINE modifyRef' #-}