From 8d49b677fc478bfef84f36b3350adbc696e9d9e4 Mon Sep 17 00:00:00 2001
From: KingoftheHomeless <combiner8761@gmail.com>
Date: Thu, 13 Feb 2020 21:20:24 +0100
Subject: [PATCH] Add lazy interpretations for Output and Writer (#311)

* Add lazy interpretations for Output and Writer

* INLINE the helper function of interpretViaLazyWriter

* Added strictness tests for runLazyWriter
---
 src/Polysemy/Internal/Writer.hs | 34 +++++++++++++++++
 src/Polysemy/Output.hs          | 67 +++++++++++++++++++++++++++++++++
 src/Polysemy/Writer.hs          | 61 +++++++++++++++++++++++++++++-
 test/WriterSpec.hs              | 42 +++++++++++++++++++++
 4 files changed, 203 insertions(+), 1 deletion(-)

diff --git a/src/Polysemy/Internal/Writer.hs b/src/Polysemy/Internal/Writer.hs
index adc62f7..8b4b4b7 100644
--- a/src/Polysemy/Internal/Writer.hs
+++ b/src/Polysemy/Internal/Writer.hs
@@ -5,12 +5,17 @@ module Polysemy.Internal.Writer where
 import Control.Concurrent.STM
 import Control.Exception
 import Control.Monad
+import qualified Control.Monad.Trans.Writer.Lazy as Lazy
 
+import Data.Bifunctor (first)
 import Data.Semigroup
 
 import Polysemy
 import Polysemy.Final
 
+import Polysemy.Internal
+import Polysemy.Internal.Union
+
 
 ------------------------------------------------------------------------------
 -- | An effect capable of emitting and intercepting messages.
@@ -156,3 +161,32 @@ runWriterSTMAction write = interpretH $ \case
       writeTVar switch True
       return o'
 {-# INLINE runWriterSTMAction #-}
+
+
+-- TODO (KingoftheHomeless):
+-- Benchmark to see if switching to a more flexible variant
+-- would incur a performance loss
+interpretViaLazyWriter
+  :: forall o e r a
+   . Monoid o
+  => (forall m x. Monad m => Weaving e (Lazy.WriterT o m) x -> Lazy.WriterT o m x)
+  -> Sem (e ': r) a
+  -> Sem r (o, a)
+interpretViaLazyWriter f sem = Sem $ \(k :: forall x. Union r (Sem r) x -> m x) ->
+  let
+    go :: forall x. Sem (e ': r) x -> Lazy.WriterT o m x
+    go = usingSem $ \u -> case decomp u of
+      Right (Weaving e s wv ex ins) -> f $ Weaving e s (go . wv) ex ins
+      Left g -> Lazy.WriterT $ do
+        ~(o, a) <- k $
+          weave
+            (mempty, ())
+            (\ ~(o, m) -> (fmap . first) (o <>) (interpretViaLazyWriter f m))
+            (Just . snd)
+            g
+        return (a, o)
+    {-# INLINE go #-}
+  in do
+    ~(a,s) <- Lazy.runWriterT (go sem)
+    return (s, a)
+{-# INLINE interpretViaLazyWriter #-}
diff --git a/src/Polysemy/Output.hs b/src/Polysemy/Output.hs
index 193268e..ae64de2 100644
--- a/src/Polysemy/Output.hs
+++ b/src/Polysemy/Output.hs
@@ -9,8 +9,11 @@ module Polysemy.Output
 
     -- * Interpretations
   , runOutputList
+  , runLazyOutputList
   , runOutputMonoid
+  , runLazyOutputMonoid
   , runOutputMonoidAssocR
+  , runLazyOutputMonoidAssocR
   , runOutputMonoidIORef
   , runOutputMonoidTVar
   , outputToIOMonoid
@@ -22,6 +25,7 @@ module Polysemy.Output
 
 import Data.IORef
 import Control.Concurrent.STM
+import qualified Control.Monad.Trans.Writer.Lazy as Lazy
 
 import Data.Semigroup (Endo(..))
 import Data.Bifunctor (first)
@@ -29,6 +33,9 @@ import Polysemy
 import Polysemy.State
 import Control.Monad (when)
 
+import Polysemy.Internal.Union
+import Polysemy.Internal.Writer
+
 
 ------------------------------------------------------------------------------
 -- | An effect capable of sending messages. Useful for streaming output and for
@@ -53,6 +60,22 @@ runOutputList = fmap (first reverse) . runState [] . reinterpret
   )
 {-# INLINE runOutputList #-}
 
+
+------------------------------------------------------------------------------
+-- | Run an 'Output' effect by transforming it into a list of its values,
+-- lazily.
+--
+-- __Warning: This inherits the nasty space leak issue of__
+-- __'Lazy.WriterT'! Don't use this if you don't have to.__
+--
+-- @since 1.3.0.0
+runLazyOutputList
+    :: forall o r a
+     . Sem (Output o ': r) a
+    -> Sem r ([o], a)
+runLazyOutputList = runLazyOutputMonoidAssocR pure
+{-# INLINE runLazyOutputList #-}
+
 ------------------------------------------------------------------------------
 -- | Run an 'Output' effect by transforming it into a monoid.
 --
@@ -69,6 +92,25 @@ runOutputMonoid f = runState mempty . reinterpret
   )
 {-# INLINE runOutputMonoid #-}
 
+
+------------------------------------------------------------------------------
+-- | Run an 'Output' effect by transforming it into a monoid, and accumulate
+-- it lazily.
+--
+-- __Warning: This inherits the nasty space leak issue of__
+-- __'Lazy.WriterT'! Don't use this if you don't have to.__
+--
+-- @since 1.3.0.0
+runLazyOutputMonoid
+    :: forall o m r a
+     . Monoid m
+    => (o -> m)
+    -> Sem (Output o ': r) a
+    -> Sem r (m, a)
+runLazyOutputMonoid f = interpretViaLazyWriter $ \(Weaving e s _ ex _) ->
+  case e of
+    Output o -> ex s <$ Lazy.tell (f o)
+
 ------------------------------------------------------------------------------
 -- | Like 'runOutputMonoid', but right-associates uses of '<>'.
 --
@@ -90,6 +132,31 @@ runOutputMonoidAssocR f =
   . runOutputMonoid (\o -> let !o' = f o in Endo (o' <>))
 {-# INLINE runOutputMonoidAssocR #-}
 
+------------------------------------------------------------------------------
+-- | Like 'runLazyOutputMonoid', 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.
+--
+-- You should always use this instead of 'runLazyOutputMonoid' if the monoid
+-- is a list, such as 'String'.
+--
+-- __Warning: This inherits the nasty space leak issue of__
+-- __'Lazy.WriterT'! Don't use this if you don't have to.__
+--
+-- @since 1.3.0.0
+runLazyOutputMonoidAssocR
+    :: forall o m r a
+     . Monoid m
+    => (o -> m)
+    -> Sem (Output o ': r) a
+    -> Sem r (m, a)
+runLazyOutputMonoidAssocR f =
+    fmap (first (`appEndo` mempty))
+  . runLazyOutputMonoid (\o -> let o' = f o in Endo (o' <>))
+                              --   ^ N.B. No bang pattern
+{-# INLINE runLazyOutputMonoidAssocR #-}
+
 ------------------------------------------------------------------------------
 -- | Run an 'Output' effect by transforming it into atomic operations
 -- over an 'IORef'.
diff --git a/src/Polysemy/Writer.hs b/src/Polysemy/Writer.hs
index 430a541..c533217 100644
--- a/src/Polysemy/Writer.hs
+++ b/src/Polysemy/Writer.hs
@@ -12,7 +12,9 @@ module Polysemy.Writer
 
     -- * Interpretations
   , runWriter
+  , runLazyWriter
   , runWriterAssocR
+  , runLazyWriterAssocR
   , runWriterTVar
   , writerToIOFinal
   , writerToIOAssocRFinal
@@ -23,6 +25,7 @@ module Polysemy.Writer
   ) where
 
 import Control.Concurrent.STM
+import qualified Control.Monad.Trans.Writer.Lazy as Lazy
 
 import Data.Bifunctor (first)
 import Data.Semigroup
@@ -31,6 +34,7 @@ import Polysemy
 import Polysemy.Output
 import Polysemy.State
 
+import Polysemy.Internal.Union
 import Polysemy.Internal.Writer
 
 
@@ -55,7 +59,8 @@ outputToWriter = interpret $ \case
 
 
 ------------------------------------------------------------------------------
--- | Run a 'Writer' effect in the style of 'Control.Monad.Trans.Writer.WriterT'
+-- | Run a 'Writer' effect in the style of
+-- 'Control.Monad.Trans.Writer.Strict.WriterT'
 -- (but without the nasty space leak!)
 runWriter
     :: Monoid o
@@ -81,6 +86,35 @@ runWriter = runState mempty . reinterpretH
   )
 {-# INLINE runWriter #-}
 
+
+------------------------------------------------------------------------------
+-- | Run a 'Writer' effect in the style of 'Control.Monad.Trans.Writer.WriterT'
+-- lazily.
+--
+-- __Warning: This inherits the nasty space leak issue of__
+-- __'Lazy.WriterT'! Don't use this if you don't have to.__
+--
+-- @since 1.3.0.0
+runLazyWriter
+    :: forall o r a
+     . Monoid o
+    => Sem (Writer o ': r) a
+    -> Sem r (o, a)
+runLazyWriter = interpretViaLazyWriter $ \(Weaving e s wv ex ins) ->
+  case e of
+    Tell o   -> ex s <$ Lazy.tell o
+    Listen m -> do
+      let m' = wv (m <$ s)
+      ~(fa, o) <- Lazy.listen m'
+      return $ ex $ (,) o <$> fa
+    Pass m -> do
+      let m' = wv (m <$ s)
+      Lazy.pass $ do
+        ft <- m'
+        let f = maybe id fst (ins ft)
+        return (ex (fmap snd ft), f)
+{-# INLINE runLazyWriter #-}
+
 -----------------------------------------------------------------------------
 -- | Like 'runWriter', but right-associates uses of '<>'.
 --
@@ -102,6 +136,31 @@ runWriterAssocR =
   . raiseUnder
 {-# INLINE runWriterAssocR #-}
 
+
+-----------------------------------------------------------------------------
+-- | Like 'runLazyWriter', 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.
+--
+-- You should always use this instead of 'runLazyWriter' if the monoid
+-- is a list, such as 'String'.
+--
+-- __Warning: This inherits the nasty space leak issue of__
+-- __'Lazy.WriterT'! Don't use this if you don't have to.__
+--
+-- @since 1.3.0.0
+runLazyWriterAssocR
+    :: Monoid o
+    => Sem (Writer o ': r) a
+    -> Sem r (o, a)
+runLazyWriterAssocR =
+    (fmap . first) (`appEndo` mempty)
+  . runLazyWriter
+  . writerToEndoWriter
+  . raiseUnder
+{-# INLINE runLazyWriterAssocR #-}
+
 --------------------------------------------------------------------
 -- | Transform a 'Writer' effect into atomic operations
 -- over a 'TVar' through final 'IO'.
diff --git a/test/WriterSpec.hs b/test/WriterSpec.hs
index 16445a4..a72134c 100644
--- a/test/WriterSpec.hs
+++ b/test/WriterSpec.hs
@@ -11,6 +11,7 @@ import Control.Exception (evaluate)
 import Polysemy
 import Polysemy.Async
 import Polysemy.Error
+import Polysemy.Input
 import Polysemy.Writer
 
 censor' :: forall e s a r
@@ -156,3 +157,44 @@ spec = do
       Right end2 <- runFinal . runError $ test6
       end1 `shouldBe` "message has been received"
       end2 `shouldBe` "message has been received"
+
+  describe "runLazyWriter" $ do
+    let
+      runLazily     = run . runInputConst () . runLazyWriter @[Int]
+      runSemiLazily = runLazily . runError @()
+      runStrictly   = run . runError @() . runLazyWriter @[Int]
+      runStrictlyM  = runM . runLazyWriter @[Int]
+
+      act :: Member (Writer [Int]) r => Sem r ()
+      act = do
+        tell @[Int] [1]
+        tell @[Int] [2]
+        error "strict"
+
+    it "should build the final output lazily, if the interpreters after \
+       \runLazyWriter and the final monad are lazy" $ do
+      (take 2 . fst . runLazily) act `shouldBe` [1,2]
+      (take 2 . fst . runSemiLazily) act `shouldBe` [1,2]
+      evaluate (runStrictly act) `shouldThrow` errorCall "strict"
+      runStrictlyM act `shouldThrow` errorCall "strict"
+
+    it "should listen lazily if all interpreters and final monad are lazy" $ do
+      let
+        listenAct :: Member (Writer [Int]) r => Sem r [Int]
+        listenAct = do
+          (end,_) <- listen @[Int] act
+          return (take 2 end)
+      (snd . runLazily) listenAct `shouldBe` [1,2]
+
+      evaluate ((snd . runSemiLazily) listenAct) `shouldThrow` errorCall "strict"
+      evaluate (runStrictly listenAct) `shouldThrow` errorCall "strict"
+      runStrictlyM listenAct `shouldThrow` errorCall "strict"
+
+    it "should censor lazily if all interpreters and final monad are lazy" $ do
+      let
+        censorAct :: Member (Writer [Int]) r => Sem r ()
+        censorAct = censor @[Int] (\(_:y:_) -> [0,y]) act
+      (fst . runLazily) censorAct `shouldBe` [0,2]
+      evaluate ((fst . runSemiLazily) censorAct) `shouldThrow` errorCall "strict"
+      evaluate (runStrictly censorAct) `shouldThrow` errorCall "strict"
+      runStrictlyM censorAct `shouldThrow` errorCall "strict"