9710c7b · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Change newVar signature, add modifyVar'

9710c7bHarendra Kumar7 years ago

src/Streamly/Internal/Mutable/Prim/Var.hs+23−10

compare with master

unifiedsplit

@@ -18,8 +18,8 @@

18-- Stability : experimental

19-- Portability : GHC

20--

21−-- A mutable variable in a mutable monad (IO/ST) holding a 'Prim' value. This

21+-- A mutable variable in a mutation capable monad (IO/ST) holding a 'Prim'

22−-- allows fast modification of a value.

22+-- value. This allows fast modification because of unboxed storage.

24module Streamly.Internal.Mutable.Prim.Var

25 (

@@ -28,13 +28,13 @@ module Streamly.Internal.Mutable.Prim.Var

28 , Prim

30 -- * Construction

31−

32 , newVar

31 , newVar

34− -- * modification

33+ -- * Write

35 , writeVar

34 , writeVar

35+ , modifyVar'

37− -- * modification

37+ -- * Read

38 , readVar

39 )

40where

@@ -43,19 +43,23 @@ import Control.Monad.Primitive (PrimMonad(..), primitive_)

43import Data.Primitive.Types (Prim, sizeOf#, readByteArray#, writeByteArray#)

44import GHC.Exts (MutableByteArray#, newByteArray#)

46−-- XXX use unlifted newtype?

47−--

48-- | A 'Var' holds a single 'Prim' value.

46-- | A 'Var' holds a single 'Prim' value.

49data Var m a = Var (MutableByteArray# (PrimState m))

47data Var m a = Var (MutableByteArray# (PrimState m))

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.

51type MonadMut = PrimMonad

53type MonadMut = PrimMonad

53-- | Create a new mutable variable.

55-- | Create a new mutable variable.

54{-# INLINE newVar #-}

56{-# INLINE newVar #-}

55−newVar :: forall m a. (MonadMut m, Prim a) => m (Var m a)

57+newVar :: forall m a. (MonadMut m, Prim a) => a -> m (Var m a)

56−newVar = primitive (\s# ->

58+newVar x = primitive (\s# ->

57 case newByteArray# (sizeOf# (undefined :: a)) s# of

59 case newByteArray# (sizeOf# (undefined :: a)) s# of

60+ (# s1#, arr# #) ->

61+ case writeByteArray# arr# 0# x s1# of

58− (# s'#, arr# #) -> (# s'#, Var arr# #)

62+ s2# -> (# s2#, Var arr# #)

59 )

63 )

61-- | Write a value to a mutable variable.

65-- | Write a value to a mutable variable.

@@ -63,6 +67,15 @@ newVar = primitive (\s# ->

63writeVar :: (MonadMut m, Prim a) => Var m a -> a -> m ()

67writeVar :: (MonadMut m, Prim a) => Var m a -> a -> m ()

64writeVar (Var arr#) x = primitive_ (writeByteArray# arr# 0# x)

68writeVar (Var arr#) x = primitive_ (writeByteArray# arr# 0# x)

70+-- | Read a value from a variable.

66{-# INLINE readVar #-}

71{-# INLINE readVar #-}

67readVar :: (MonadMut m, Prim a) => Var m a -> m a

72readVar :: (MonadMut m, Prim a) => Var m a -> m a

68readVar (Var arr#) = primitive (readByteArray# arr# 0#)

73readVar (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'#