@@ -18,8 +18,8 @@
1818-- Stability : experimental
1919-- Portability : GHC
2020--
21−-- A mutable variable in a mutable monad (IO/ST) holding a 'Prim' value. This
22−-- allows fast modification of a value.
21+-- A mutable variable in a mutation capable monad (IO/ST) holding a 'Prim'
22+-- value. This allows fast modification because of unboxed storage.
2323
2424module Streamly.Internal.Mutable.Prim.Var
2525 (
@@ -28,13 +28,13 @@ module Streamly.Internal.Mutable.Prim.Var
2828 , Prim
2929
3030 -- * Construction
31−
3231 , newVar
3332
34− -- * modification
33+ -- * Write
3534 , writeVar
35+ , modifyVar'
3636
37− -- * modification
37+ -- * Read
3838 , readVar
3939 )
4040where
@@ -43,19 +43,23 @@ import Control.Monad.Primitive (PrimMonad(..), primitive_)
4343import Data.Primitive.Types (Prim, sizeOf#, readByteArray#, writeByteArray#)
4444import GHC.Exts (MutableByteArray#, newByteArray#)
4545
46−-- XXX use unlifted newtype?
47−--
4846-- | A 'Var' holds a single 'Prim' value.
4947data Var m a = Var (MutableByteArray# (PrimState m))
5048
49+-- The name PrimMonad does not give a clue what it means, an explicit "Mut"
50+-- suffix provides a better hint. MonadMut is just a generalization of MonadIO.
51+--
52+-- | A monad that allows mutable operations using a state token.
5153type MonadMut = PrimMonad
5254
5355-- | Create a new mutable variable.
5456{-# INLINE newVar #-}
55−newVar :: forall m a. (MonadMut m, Prim a) => m (Var m a)
56−newVar = primitive (\s# ->
57+newVar :: forall m a. (MonadMut m, Prim a) => a -> m (Var m a)
58+newVar x = primitive (\s# ->
5759 case newByteArray# (sizeOf# (undefined :: a)) s# of
58− (# s'#, arr# #) -> (# s'#, Var arr# #)
60+ (# s1#, arr# #) ->
61+ case writeByteArray# arr# 0# x s1# of
62+ s2# -> (# s2#, Var arr# #)
5963 )
6064
6165-- | Write a value to a mutable variable.
@@ -63,6 +67,15 @@ newVar = primitive (\s# ->
6367writeVar :: (MonadMut m, Prim a) => Var m a -> a -> m ()
6468writeVar (Var arr#) x = primitive_ (writeByteArray# arr# 0# x)
6569
70+-- | Read a value from a variable.
6671{-# INLINE readVar #-}
6772readVar :: (MonadMut m, Prim a) => Var m a -> m a
6873readVar (Var arr#) = primitive (readByteArray# arr# 0#)
74+
75+-- | Modify the value of a mutable variable using a function with strict
76+-- application.
77+{-# INLINE modifyVar' #-}
78+modifyVar' :: (MonadMut m, Prim a) => Var m a -> (a -> a) -> m ()
79+modifyVar' (Var arr#) g = primitive_ $ \s# ->
80+ case readByteArray# arr# 0# s# of
81+ (# s'#, a #) -> let a' = g a in a' `seq` writeByteArray# arr# 0# a' s'#