Add AtomicState and atomic interpreters for Output/Writer (#206)

* Add AtomicState and atomic interpreters for Output/Writer

* Fixed a word in docs, moved test functions around

* Remove runWriterIORef and runWriterTVar

package.yaml

@@ -20,6 +20,7 @@ dependencies:
 - containers >= 0.5 && < 0.7
 - mtl >= 2.2.2 && < 3
 - syb >= 0.7 && < 0.8
+- stm >= 2 && < 3
 - template-haskell >= 2.12.0.0 && < 3
 - th-abstraction >= 0.3.1.0 && < 0.4
 - transformers >= 0.5.2.0 && < 0.6

polysemy.cabal

@@ -4,7 +4,7 @@ cabal-version: 1.12
 --
 -- see: https://github.com/sol/hpack
 --
--- hash: ff7213d283da75830205760ef903cd73a3feb2c082b648ce653ab85d1291db8f
+-- hash: 7a803bf8b5358a99d104c182fbbb4cea4d7190e5da71481af76b6ca17c05bc27
 
 name:           polysemy
 version:        1.0.0.0
@@ -41,6 +41,7 @@ library
   exposed-modules:
       Polysemy
       Polysemy.Async
+      Polysemy.AtomicState
       Polysemy.Embed
       Polysemy.Embed.Type
       Polysemy.Error
@@ -80,6 +81,7 @@ library
     , containers >=0.5 && <0.7
     , first-class-families >=0.5.0.0 && <0.6
     , mtl >=2.2.2 && <3
+    , stm >=2 && <3
     , syb >=0.7 && <0.8
     , template-haskell >=2.12.0.0 && <3
     , th-abstraction >=0.3.1.0 && <0.4
@@ -142,6 +144,7 @@ test-suite polysemy-test
     , inspection-testing >=0.4.2 && <0.5
     , mtl >=2.2.2 && <3
     , polysemy
+    , stm >=2 && <3
     , syb >=0.7 && <0.8
     , template-haskell >=2.12.0.0 && <3
     , th-abstraction >=0.3.1.0 && <0.4
@@ -175,6 +178,7 @@ benchmark polysemy-bench
     , freer-simple
     , mtl
     , polysemy
+    , stm >=2 && <3
     , syb >=0.7 && <0.8
     , template-haskell >=2.12.0.0 && <3
     , th-abstraction >=0.3.1.0 && <0.4

src/Polysemy/AtomicState.hs

@@ -0,0 +1,126 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE TemplateHaskell #-}
+module Polysemy.AtomicState
+  ( -- * Effect
+    AtomicState (..)
+
+    -- * Actions
+  , atomicState
+  , atomicState'
+  , atomicGet
+  , atomicPut
+  , atomicModify
+  , atomicModify'
+
+    -- * Interpretations
+  , runAtomicStateIORef
+  , runAtomicStateTVar
+  , atomicStateToState
+  ) where
+
+
+import Control.Concurrent.STM
+
+import Polysemy
+import Polysemy.State
+
+import Data.IORef
+
+data AtomicState s m a where
+  AtomicState :: (s -> (s, a)) -> AtomicState s m a
+  AtomicGet   :: AtomicState s m s
+
+makeSem_ ''AtomicState
+
+-----------------------------------------------------------------------------
+-- | Atomically reads and modifies the state.
+atomicState :: forall s a r
+             . Member (AtomicState s) r
+            => (s -> (s, a))
+            -> Sem r a
+
+atomicGet :: forall s r
+           . Member (AtomicState s) r
+          => Sem r s
+
+-----------------------------------------------------------------------------
+-- | A variant of 'atomicState' in which the computation is strict in the new
+-- state and return value.
+atomicState' :: Member (AtomicState s) r
+              => (s -> (s, a))
+              -> Sem r a
+atomicState' f = do
+  -- KingoftheHomeless: return value needs to be forced due to how
+  -- 'atomicModifyIORef' is implemented: the computation
+  -- (and thus the new state) is forced only once the return value is.
+  !a <- atomicState $ \s ->
+    case f s of
+      v@(!_, _) -> v
+  return a
+{-# INLINE atomicState' #-}
+
+atomicPut :: Member (AtomicState s) r
+          => s
+          -> Sem r ()
+atomicPut s = do
+  !_ <- atomicState $ \_ -> (s, ()) -- strict put with atomicModifyIORef
+  return ()
+{-# INLINE atomicPut #-}
+
+atomicModify :: Member (AtomicState s) r
+             => (s -> s)
+             -> Sem r ()
+atomicModify f = atomicState $ \s -> (f s, ())
+{-# INLINE atomicModify #-}
+
+-----------------------------------------------------------------------------
+-- | A variant of 'atomicModify' in which the computation is strict in the
+-- new state.
+atomicModify' :: Member (AtomicState s) r
+              => (s -> s)
+              -> Sem r ()
+atomicModify' f = do
+  !_ <- atomicState $ \s -> let !s' = f s in (s', ())
+  return ()
+{-# INLINE atomicModify' #-}
+
+------------------------------------------------------------------------------
+-- | Run an 'AtomicState' effect by transforming it into atomic operations
+-- over an 'IORef'.
+runAtomicStateIORef :: Member (Embed IO) r
+                    => IORef s
+                    -> Sem (AtomicState s ': r) a
+                    -> Sem r a
+runAtomicStateIORef ref = interpret $ \case
+  AtomicState f -> embed $ atomicModifyIORef ref f
+  AtomicGet     -> embed $ readIORef ref
+{-# INLINE runAtomicStateIORef #-}
+
+------------------------------------------------------------------------------
+-- | Run an 'AtomicState' effect by transforming it into atomic operations
+-- over a 'TVar'.
+runAtomicStateTVar :: Member (Embed IO) r
+                   => TVar s
+                   -> Sem (AtomicState s ': r) a
+                   -> Sem r a
+runAtomicStateTVar tvar = interpret $ \case
+  AtomicState f -> embed $ atomically $ do
+    (s', a) <- f <$> readTVar tvar
+    writeTVar tvar s'
+    return a
+  AtomicGet -> embed $ readTVarIO tvar
+{-# INLINE runAtomicStateTVar #-}
+
+------------------------------------------------------------------------------
+-- | Transform an 'AtomicState' effect to a 'State' effect, discarding
+-- the notion of atomicity.
+atomicStateToState :: Member (State s) r
+                   => Sem (AtomicState s ': r) a
+                   -> Sem r a
+atomicStateToState = interpret $ \case
+  AtomicState f -> do
+    (s', a) <- f <$> get
+    put s'
+    return a
+  AtomicGet -> get
+{-# INLINE atomicStateToState #-}

src/Polysemy/Output.hs

@@ -11,11 +11,16 @@ module Polysemy.Output
   , runOutputList
   , runOutputMonoid
   , runOutputMonoidAssocR
+  , runOutputMonoidIORef
+  , runOutputMonoidTVar
   , ignoreOutput
   , runOutputBatched
   , runOutputSem
   ) where
 
+import Data.IORef
+import Control.Concurrent.STM
+
 import Data.Semigroup (Endo(..))
 import Data.Bifunctor (first)
 import Polysemy
@@ -81,6 +86,36 @@ runOutputMonoidAssocR f =
   . runOutputMonoid (\a -> Endo (f a <>))
 {-# INLINE runOutputMonoidAssocR #-}
 
+------------------------------------------------------------------------------
+-- | Run an 'Output' effect by transforming it into atomic operations
+-- over an 'IORef'.
+runOutputMonoidIORef
+    :: forall o m r a
+     . (Monoid m, Member (Embed IO) r)
+    => IORef m
+    -> (o -> m)
+    -> Sem (Output o ': r) a
+    -> Sem r a
+runOutputMonoidIORef ref f = interpret $ \case
+  Output o -> embed $ atomicModifyIORef' ref (\s -> (s <> f o, ()))
+{-# INLINE runOutputMonoidIORef #-}
+
+------------------------------------------------------------------------------
+-- | Run an 'Output' effect by transforming it into atomic operations
+-- over a 'TVar'.
+runOutputMonoidTVar
+    :: forall o m r a
+     . (Monoid m, Member (Embed IO) r)
+    => TVar m
+    -> (o -> m)
+    -> Sem (Output o ': r) a
+    -> Sem r a
+runOutputMonoidTVar tvar f = interpret $ \case
+  Output o -> embed $ atomically $ do
+    s <- readTVar tvar
+    writeTVar tvar $! s <> f o
+{-# INLINE runOutputMonoidTVar #-}
+
 ------------------------------------------------------------------------------
 -- | Run an 'Output' effect by ignoring it.
 --

src/Polysemy/State.hs

@@ -105,6 +105,11 @@ evalLazyState s = fmap snd . runLazyState s
 ------------------------------------------------------------------------------
 -- | Run a 'State' effect by transforming it into operations over an 'IORef'.
 --
+-- /Note/: This is not safe in a concurrent setting, as 'modify' isn't atomic.
+-- If you need operations over the state to be atomic,
+-- use 'Polysemy.AtomicState.runAtomicStateIORef' or
+-- 'Polysemy.AtomicState.runAtomicStateTVar' instead.
+--
 -- @since 1.0.0.0
 runStateIORef
     :: forall s r a

src/Polysemy/Writer.hs

@@ -81,11 +81,11 @@ runWriter = runState mempty . reinterpretH
   )
 {-# INLINE runWriter #-}
 
-------------------------------------------------------------------------------
+-----------------------------------------------------------------------------
 -- | Like 'runWriter', but right-associates uses of '<>'.
 --
--- This asymptotically improves performance if the time complexity of '<>' for
--- the 'Monoid' depends only on the size of the first argument.
+-- This asymptotically improves performance if the time complexity of '<>'
+-- for the 'Monoid' depends only on the size of the first argument.
 --
 -- You should always use this instead of 'runWriter' if the monoid
 -- is a list, such as 'String'.

test/OutputSpec.hs

@@ -1,12 +1,17 @@
 module OutputSpec where
 
-import Polysemy
-import Polysemy.Output
-import Data.Foldable
-import Test.Hspec
-
+import Control.Concurrent.STM
 import Control.Exception (evaluate)
 
+import Data.IORef
+import Data.Foldable
+
+import Polysemy
+import Polysemy.Async
+import Polysemy.Output
+
+import Test.Hspec
+
 spec :: Spec
 spec = parallel $ do
   describe "runOutputBatched" $ do
@@ -38,8 +43,39 @@ spec = parallel $ do
         runM t           `shouldThrow` errorCall "strict"
         evaluate (run t) `shouldThrow` errorCall "strict"
 
+  describe "runOutputMonoidIORef" $ do
+    it "should commit writes of asynced computations" $
+      let io = do
+            ref <- newIORef ""
+            (runM .@ lowerAsync) . runOutputMonoidIORef ref (show @Int) $
+              test1
+            readIORef ref
+      in do
+        res <- io
+        res `shouldBe` "12"
+
+  describe "runOutputMonoidTVar" $ do
+    it "should commit writes of asynced computations" $
+      let io = do
+            ref <- newTVarIO ""
+            (runM .@ lowerAsync) . runOutputMonoidTVar ref (show @Int) $
+              test1
+            readTVarIO ref
+      in do
+        res <- io
+        res `shouldBe` "12"
+
 runOutput :: Int -> Sem '[Output Int, Output [Int]] a -> ([[Int]], a)
 runOutput size = run . runOutputMonoid (:[]) . runOutputBatched size
 
 runOutputList' :: Sem '[Output Int] a -> ([Int], a)
 runOutputList' = run . runOutputList
+
+test1 :: Members '[Async, Output Int] r
+     => Sem r ()
+test1 = do
+  output @Int 1
+  a <- async $ do
+    output @Int 2
+  _ <- await a
+  return ()

test/WriterSpec.hs

@@ -3,9 +3,15 @@ module WriterSpec where
 
 import Test.Hspec
 
+import qualified Control.Concurrent.Async as A
+import Control.Concurrent.MVar
+import Control.Concurrent.STM
 import Control.Exception (evaluate)
 
+import Data.IORef
+
 import Polysemy
+import Polysemy.Async
 import Polysemy.Error
 import Polysemy.Writer
 
@@ -50,33 +56,34 @@ test3 :: (String, (String, ()))
 test3 = run . runWriter $ listen (tell "and hear")
 
 spec :: Spec
-spec = describe "writer" $ do
-  it "should not censor" $ do
-    test1 `shouldBe` ("censoring not applied", Right ())
+spec = do
+  describe "writer" $ do
+    it "should not censor" $ do
+      test1 `shouldBe` ("censoring not applied", Right ())
 
-  it "should censor" $ do
-    test2 `shouldBe` ("censoring applied", Right ())
+    it "should censor" $ do
+      test2 `shouldBe` ("censoring applied", Right ())
 
-  it "should have a proper listen" $ do
-    test3 `shouldBe` ("and hear", ("and hear", ()))
+    it "should have a proper listen" $ do
+      test3 `shouldBe` ("and hear", ("and hear", ()))
 
-  it "should be strict in the output" $
-    let
-      t1 = runWriter @String $ do
-        tell @String (error "strict")
-        return ()
+    it "should be strict in the output" $
+      let
+        t1 = runWriter @String $ do
+          tell @String (error "strict")
+          return ()
 
-      t2 = runWriter @String $ do
-        _ <- listen @String (tell @String (error "strict"))
-        return ()
+        t2 = runWriter @String $ do
+          _ <- listen @String (tell @String (error "strict"))
+          return ()
 
-      t3 = runWriter @String $ do
-        pass @String $ pure (\_ -> error "strict", ())
-        return ()
-    in do
-      runM t1           `shouldThrow` errorCall "strict"
-      evaluate (run t1) `shouldThrow` errorCall "strict"
-      runM t2           `shouldThrow` errorCall "strict"
-      evaluate (run t2) `shouldThrow` errorCall "strict"
-      runM t3           `shouldThrow` errorCall "strict"
-      evaluate (run t3) `shouldThrow` errorCall "strict"
+        t3 = runWriter @String $ do
+          pass @String $ pure (\_ -> error "strict", ())
+          return ()
+      in do
+        runM t1           `shouldThrow` errorCall "strict"
+        evaluate (run t1) `shouldThrow` errorCall "strict"
+        runM t2           `shouldThrow` errorCall "strict"
+        evaluate (run t2) `shouldThrow` errorCall "strict"
+        runM t3           `shouldThrow` errorCall "strict"
+        evaluate (run t3) `shouldThrow` errorCall "strict"