{-# LANGUAGE UnboxedTuples #-}

-- |
-- Module      : Streamly.Internal.Data.IORef.Prim
-- Copyright   : (c) 2019 Composewell Technologies
--
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- A mutable variable in a mutation capable monad (IO) holding a 'Prim'
-- value. This allows fast modification because of unboxed storage.
--
-- = Multithread Consistency Notes
--
-- In general, any value that straddles a machine word cannot be guaranteed to
-- be consistently read from another thread without a lock.  GHC heap objects
-- are always machine word aligned, therefore, a 'IORef' is also word aligned.
-- On a 64-bit platform, writing a 64-bit aligned type from one thread and
-- reading it from another thread should give consistent old or new value. The
-- same holds true for 32-bit values on a 32-bit platform.

module Streamly.Internal.Data.IORef.Prim
    (
      IORef
    , Prim

    -- * Construction
    , newIORef

    -- * Write
    , writeIORef
    , modifyIORef'

    -- * Read
    , readIORef
    , toStream
    , toStreamD
    )
where

#include "inline.hs"

import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Primitive (primitive_)
import Data.Primitive.Types (Prim, sizeOf#, readByteArray#, writeByteArray#)
import GHC.Exts (MutableByteArray#, newByteArray#, RealWorld)
import GHC.IO (IO(..))

import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
import qualified Streamly.Internal.Data.Stream.StreamD.Type as D

-- | An 'IORef' holds a single 'Prim' value.
data IORef a = IORef (MutableByteArray# RealWorld)

-- | Create a new 'IORef'.
--
-- /Pre-release/
{-# INLINE newIORef #-}
newIORef :: forall a. Prim a => a -> IO (IORef a)
newIORef :: a -> IO (IORef a)
newIORef a
x = (State# RealWorld -> (# State# RealWorld, IORef a #))
-> IO (IORef a)
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (\State# RealWorld
s# ->
      case Int#
-> State# RealWorld
-> (# State# RealWorld, MutableByteArray# RealWorld #)
forall d. Int# -> State# d -> (# State# d, MutableByteArray# d #)
newByteArray# (a -> Int#
forall a. Prim a => a -> Int#
sizeOf# (a
forall a. HasCallStack => a
undefined :: a)) State# RealWorld
s# of
        (# State# RealWorld
s1#, MutableByteArray# RealWorld
arr# #) ->
            case MutableByteArray# RealWorld
-> Int# -> a -> State# RealWorld -> State# RealWorld
forall a s.
Prim a =>
MutableByteArray# s -> Int# -> a -> State# s -> State# s
writeByteArray# MutableByteArray# RealWorld
arr# Int#
0# a
x State# RealWorld
s1# of
                State# RealWorld
s2# -> (# State# RealWorld
s2#, MutableByteArray# RealWorld -> IORef a
forall a. MutableByteArray# RealWorld -> IORef a
IORef MutableByteArray# RealWorld
arr# #)
    )

-- | Write a value to an 'IORef'.
--
-- /Pre-release/
{-# INLINE writeIORef #-}
writeIORef :: Prim a => IORef a -> a -> IO ()
writeIORef :: IORef a -> a -> IO ()
writeIORef (IORef MutableByteArray# RealWorld
arr#) a
x = (State# (PrimState IO) -> State# (PrimState IO)) -> IO ()
forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ (MutableByteArray# RealWorld
-> Int# -> a -> State# RealWorld -> State# RealWorld
forall a s.
Prim a =>
MutableByteArray# s -> Int# -> a -> State# s -> State# s
writeByteArray# MutableByteArray# RealWorld
arr# Int#
0# a
x)

-- | Read a value from an 'IORef'.
--
-- /Pre-release/
{-# INLINE readIORef #-}
readIORef :: Prim a => IORef a -> IO a
readIORef :: IORef a -> IO a
readIORef (IORef MutableByteArray# RealWorld
arr#) = (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
forall a. (State# RealWorld -> (# State# RealWorld, a #)) -> IO a
IO (MutableByteArray# RealWorld
-> Int# -> State# RealWorld -> (# State# RealWorld, a #)
forall a s.
Prim a =>
MutableByteArray# s -> Int# -> State# s -> (# State# s, a #)
readByteArray# MutableByteArray# RealWorld
arr# Int#
0#)

-- | Modify the value of an 'IORef' using a function with strict application.
--
-- /Pre-release/
{-# INLINE modifyIORef' #-}
modifyIORef' :: Prim a => IORef a -> (a -> a) -> IO ()
modifyIORef' :: IORef a -> (a -> a) -> IO ()
modifyIORef' (IORef MutableByteArray# RealWorld
arr#) a -> a
g = (State# (PrimState IO) -> State# (PrimState IO)) -> IO ()
forall (m :: * -> *).
PrimMonad m =>
(State# (PrimState m) -> State# (PrimState m)) -> m ()
primitive_ ((State# (PrimState IO) -> State# (PrimState IO)) -> IO ())
-> (State# (PrimState IO) -> State# (PrimState IO)) -> IO ()
forall a b. (a -> b) -> a -> b
$ \State# (PrimState IO)
s# ->
  case MutableByteArray# RealWorld
-> Int# -> State# RealWorld -> (# State# RealWorld, a #)
forall a s.
Prim a =>
MutableByteArray# s -> Int# -> State# s -> (# State# s, a #)
readByteArray# MutableByteArray# RealWorld
arr# Int#
0# State# RealWorld
State# (PrimState IO)
s# of
    (# State# RealWorld
s'#, a
a #) -> let a' :: a
a' = a -> a
g a
a in a
a' a -> State# RealWorld -> State# RealWorld
`seq` MutableByteArray# RealWorld
-> Int# -> a -> State# RealWorld -> State# RealWorld
forall a s.
Prim a =>
MutableByteArray# s -> Int# -> a -> State# s -> State# s
writeByteArray# MutableByteArray# RealWorld
arr# Int#
0# a
a' State# RealWorld
s'#

-- | Generate a stream by continuously reading the IORef.
--
-- /Pre-release/
{-# INLINE_NORMAL toStreamD #-}
toStreamD :: (MonadIO m, Prim a) => IORef a -> D.Stream m a
toStreamD :: IORef a -> Stream m a
toStreamD IORef a
var = (State Stream m a -> () -> m (Step () a)) -> () -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State Stream m a -> () -> m (Step () a)
forall (m :: * -> *) p. MonadIO m => p -> () -> m (Step () a)
step ()

    where

    {-# INLINE_LATE step #-}
    step :: p -> () -> m (Step () a)
step p
_ () = IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IORef a -> IO a
forall a. Prim a => IORef a -> IO a
readIORef IORef a
var) m a -> (a -> m (Step () a)) -> m (Step () a)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
x -> Step () a -> m (Step () a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step () a -> m (Step () a)) -> Step () a -> m (Step () a)
forall a b. (a -> b) -> a -> b
$ a -> () -> Step () a
forall s a. a -> s -> Step s a
D.Yield a
x ()

-- | Construct a stream by reading a 'Prim' 'IORef' repeatedly.
--
-- /Pre-release/
--
{-# INLINE toStream #-}
toStream :: (K.IsStream t, MonadIO m, Prim a) => IORef a -> t m a
toStream :: IORef a -> t m a
toStream = Stream m a -> t m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
Stream m a -> t m a
D.fromStreamD (Stream m a -> t m a)
-> (IORef a -> Stream m a) -> IORef a -> t m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IORef a -> Stream m a
forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
IORef a -> Stream m a
toStreamD