{-# 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 = forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall s a. MVector s a -> URef s a
URef forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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) = forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
V.unsafeRead MVector s 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) = forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
V.unsafeWrite MVector s 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 = forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> m a
V.unsafeRead MVector s a
v Int
0 forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) (v :: * -> * -> *) a.
(PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
V.unsafeWrite MVector s a
v Int
0 forall b c a. (b -> c) -> (a -> b) -> a -> c
. 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' = forall c (m :: * -> *).
(MutableRef c, PrimMonad m, PrimState m ~ MCState c) =>
c -> (RefElement c -> RefElement c) -> m ()
modifyRef
    {-# INLINE modifyRef' #-}