Replace Prim with Unboxed in Data.IORef

core/src/Streamly/Internal/Data/IORef/Prim.hs

@@ -9,7 +9,7 @@
 -- Stability   : experimental
 -- Portability : GHC
 --
--- A mutable variable in a mutation capable monad (IO) holding a 'Prim'
+-- A mutable variable in a mutation capable monad (IO) holding a 'Unboxed'
 -- value. This allows fast modification because of unboxed storage.
 --
 -- = Multithread Consistency Notes
@@ -24,7 +24,6 @@
 module Streamly.Internal.Data.IORef.Prim
     (
       IORef
-    , Prim
 
     -- * Construction
     , newIORef
@@ -42,56 +41,57 @@ 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 Streamly.Internal.Data.Unboxed
+    ( ArrayContents(..)
+    , Unboxed(..)
+    , peekWith
+    , pokeWith
+    , newUnpinnedArrayContents
+    )
 
 import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
 
--- | An 'IORef' holds a single 'Prim' value.
-data IORef a = IORef (MutableByteArray# RealWorld)
+-- | An 'IORef' holds a single 'Unboxed' value.
+newtype IORef a = IORef (ArrayContents a)
 
 -- | Create a new 'IORef'.
 --
 -- /Pre-release/
 {-# INLINE newIORef #-}
-newIORef :: forall a. Prim a => a -> IO (IORef a)
-newIORef x = IO (\s# ->
-      case newByteArray# (sizeOf# (undefined :: a)) s# of
-        (# s1#, arr# #) ->
-            case writeByteArray# arr# 0# x s1# of
-                s2# -> (# s2#, IORef arr# #)
-    )
+newIORef :: forall a. Unboxed a => a -> IO (IORef a)
+newIORef x = do
+    var <- newUnpinnedArrayContents (sizeOf (undefined :: a))
+    pokeWith var 0 x
+    return $ IORef var
 
 -- | Write a value to an 'IORef'.
 --
 -- /Pre-release/
 {-# INLINE writeIORef #-}
-writeIORef :: Prim a => IORef a -> a -> IO ()
-writeIORef (IORef arr#) x = primitive_ (writeByteArray# arr# 0# x)
+writeIORef :: Unboxed a => IORef a -> a -> IO ()
+writeIORef (IORef var) x = pokeWith var 0 x
 
 -- | Read a value from an 'IORef'.
 --
 -- /Pre-release/
 {-# INLINE readIORef #-}
-readIORef :: Prim a => IORef a -> IO a
-readIORef (IORef arr#) = IO (readByteArray# arr# 0#)
+readIORef :: Unboxed a => IORef a -> IO a
+readIORef (IORef var) = peekWith var 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 arr#) g = primitive_ $ \s# ->
-  case readByteArray# arr# 0# s# of
-    (# s'#, a #) -> let a' = g a in a' `seq` writeByteArray# arr# 0# a' s'#
+modifyIORef' :: Unboxed a => IORef a -> (a -> a) -> IO ()
+modifyIORef' var g = do
+  x <- readIORef var
+  writeIORef var (g x)
 
 -- | 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 :: (MonadIO m, Unboxed a) => IORef a -> D.Stream m a
 toStreamD var = D.Stream step ()
 
     where

core/src/Streamly/Internal/Data/Time/Units.hs

@@ -51,7 +51,7 @@ where
 import Text.Printf (printf)
 
 import Data.Int
-import Data.Primitive.Types (Prim(..))
+import Streamly.Internal.Data.Unboxed (Unboxed)
 import Streamly.Internal.Data.Time.TimeSpec
 
 -------------------------------------------------------------------------------
@@ -106,7 +106,7 @@ newtype NanoSecond64 = NanoSecond64 Int64
              , Real
              , Integral
              , Ord
-             , Prim
+             , Unboxed
              )
 
 -- | An 'Int64' time representation with a microsecond resolution.
@@ -121,7 +121,7 @@ newtype MicroSecond64 = MicroSecond64 Int64
              , Real
              , Integral
              , Ord
-             , Prim
+             , Unboxed
              )
 
 -- | An 'Int64' time representation with a millisecond resolution.
@@ -136,7 +136,7 @@ newtype MilliSecond64 = MilliSecond64 Int64
              , Real
              , Integral
              , Ord
-             , Prim
+             , Unboxed
              )
 
 -------------------------------------------------------------------------------

core/src/Streamly/Internal/Data/Unboxed.hs

@@ -14,6 +14,7 @@ module Streamly.Internal.Data.Unboxed
     , getMutableByteArray#
     , pin
     , unpin
+    , newUnpinnedArrayContents
     ) where
 
 #include "MachDeps.h"
@@ -22,7 +23,7 @@ module Streamly.Internal.Data.Unboxed
 
 import Data.Complex (Complex((:+)), realPart)
 import GHC.Base (IO(..))
-import GHC.Int (Int32(..))
+import GHC.Int (Int32(..), Int64(..))
 import GHC.Word (Word8(..), Word64(..))
 
 import GHC.Exts
@@ -47,6 +48,20 @@ touch :: ArrayContents a -> IO ()
 touch (ArrayContents contents) =
     IO $ \s -> case touch# contents s of s' -> (# s', () #)
 
+--------------------------------------------------------------------------------
+-- Creation
+--------------------------------------------------------------------------------
+
+{-# INLINE newUnpinnedArrayContents #-}
+newUnpinnedArrayContents :: Int -> IO (ArrayContents a)
+newUnpinnedArrayContents nbytes | nbytes < 0 =
+  errorWithoutStackTrace "newUnpinnedArrayContents: size must be >= 0"
+newUnpinnedArrayContents (I# nbytes) = IO $ \s ->
+    case newByteArray# nbytes s of
+        (# s', mbarr# #) ->
+           let c = ArrayContents mbarr#
+            in (# s', c #)
+
 -------------------------------------------------------------------------------
 -- Pinning & Unpinning
 -------------------------------------------------------------------------------
@@ -106,6 +121,7 @@ unpin arr@(ArrayContents marr#) =
 #define SIZEOF_WORD64_PRIMITIVE 8#
 #define SIZEOF_HSDOUBLE_PRIMITIVE 8#
 #define SIZEOF_INT32_PRIMITIVE 4#
+#define SIZEOF_INT64_PRIMITIVE 8#
 
 #ifdef __GHCJS__
 #define WORD64TYP Word64#
@@ -113,6 +129,12 @@ unpin arr@(ArrayContents marr#) =
 #define WORD64TYP Word#
 #endif
 
+#ifdef __GHCJS__
+#define INT64TYP Int64#
+#else
+#define INT64TYP Int#
+#endif
+
 {-# INLINE readWord8ArrayAsWideChar# #-}
 readWord8ArrayAsWideChar# ::
        MutableByteArray# d -> Int# -> State# d -> (# State# d, Char# #)
@@ -149,6 +171,18 @@ writeWord8ArrayAsInt32# ::
 writeWord8ArrayAsInt32# arr# i# a# s# =
     writeInt32Array# arr# (quotInt# i# SIZEOF_INT32_PRIMITIVE) a# s#
 
+{-# INLINE readWord8ArrayAsInt64# #-}
+readWord8ArrayAsInt64# ::
+       MutableByteArray# d -> Int# -> State# d -> (# State# d, INT64TYP #)
+readWord8ArrayAsInt64# arr# i# s# =
+    readInt64Array# arr# (quotInt# i# SIZEOF_INT64_PRIMITIVE) s#
+
+{-# INLINE writeWord8ArrayAsInt64# #-}
+writeWord8ArrayAsInt64# ::
+       MutableByteArray# d -> Int# -> INT64TYP -> State# d -> State# d
+writeWord8ArrayAsInt64# arr# i# a# s# =
+    writeInt64Array# arr# (quotInt# i# SIZEOF_INT64_PRIMITIVE) a# s#
+
 {-# INLINE readWord8ArrayAsWord# #-}
 readWord8ArrayAsWord# ::
        MutableByteArray# d -> Int# -> State# d -> (# State# d, Word# #)
@@ -269,6 +303,13 @@ DERIVE_UNBOXED( Int
               , readWord8ArrayAsInt#
               , writeWord8ArrayAsInt#)
 
+DERIVE_UNBOXED( Int64
+              , I64#
+              , SIZEOF_INT64
+              , ALIGNMENT_INT64
+              , readWord8ArrayAsInt64#
+              , writeWord8ArrayAsInt64#)
+
 DERIVE_UNBOXED( Word
               , W#
               , SIZEOF_HSWORD

src/Streamly/Internal/Data/Stream/IsStream/Generate.hs

@@ -89,7 +89,6 @@ where
 #include "inline.hs"
 
 import Control.Monad.IO.Class (MonadIO(..))
-import Data.Primitive.Types (Prim)
 import Data.Void (Void)
 import Streamly.Internal.Control.Concurrent (MonadAsync)
 import Streamly.Internal.Data.Unfold.Type (Unfold)
@@ -105,6 +104,7 @@ import Streamly.Internal.Data.Stream.Serial (SerialT)
 import Streamly.Internal.Data.Stream.WSerial (WSerialT)
 import Streamly.Internal.Data.Stream.Zip (ZipSerialM)
 import Streamly.Internal.Data.Time.Units (AbsTime , RelTime64, addToAbsTime64)
+import Streamly.Internal.Data.Unboxed (Unboxed)
 
 import qualified Streamly.Internal.Data.IORef.Prim as Prim
 import qualified Streamly.Internal.Data.Stream.IsStream.Type as IsStream
@@ -660,5 +660,5 @@ fromCallback setCallback = concatM $ do
 -- /Pre-release/
 --
 {-# INLINE fromPrimIORef #-}
-fromPrimIORef :: (IsStream t, MonadIO m, Prim a) => Prim.IORef a -> t m a
+fromPrimIORef :: (IsStream t, MonadIO m, Unboxed a) => Prim.IORef a -> t m a
 fromPrimIORef = fromStreamD . Prim.toStreamD