{-# LANGUAGE BangPatterns, TemplateHaskell, TupleSections #-}
{-# OPTIONS_HADDOCK not-home, prune #-}

-- | Description: The 'Writer' effect
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.
data Writer o m a where
  -- | Write a message to the log.
  Tell   :: o -> Writer o m ()
  -- | Return the log produced by the higher-order action.
  Listen :: ∀ o m a. m a -> Writer o m (o, a)
  -- | Run the given action and apply the function it returns to the log.
  Pass   :: m (o -> o, a) -> Writer o m a

makeSem ''Writer

-- TODO(KingoftheHomeless): Research if this is more or less efficient than
-- using 'reinterpretH' + 'subsume'

-----------------------------------------------------------------------------
-- | Transform a @'Writer' o@ effect into a  @'Writer' ('Endo' o)@ effect,
-- right-associating all uses of '<>' for @o@.
--
-- This can be used together with 'raiseUnder' in order to create
-- @-AssocR@ variants out of regular 'Writer' interpreters.
--
-- @since 1.2.0.0
writerToEndoWriter
    :: (Monoid o, Member (Writer (Endo o)) r)
    => Sem (Writer o ': r) a
    -> Sem r a
writerToEndoWriter :: forall o (r :: EffectRow) a.
(Monoid o, Member (Writer (Endo o)) r) =>
Sem (Writer o : r) a -> Sem r a
writerToEndoWriter = forall (e :: Effect) (r :: EffectRow) a.
(forall (rInitial :: EffectRow) x.
 e (Sem rInitial) x -> Tactical e (Sem rInitial) r x)
-> Sem (e : r) a -> Sem r a
interpretH forall a b. (a -> b) -> a -> b
$ \case
      Tell o
o   -> forall o (r :: EffectRow). Member (Writer o) r => o -> Sem r ()
tell (forall a. (a -> a) -> Endo a
Endo (o
o forall a. Semigroup a => a -> a -> a
<>)) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (f :: * -> *) a (e :: Effect) (m :: * -> *)
       (r :: EffectRow).
Functor f =>
a -> Sem (WithTactics e f m r) (f a)
pureT
      Listen Sem rInitial a
m -> do
        Sem r (f a)
m' <- forall o (r :: EffectRow) a.
(Monoid o, Member (Writer (Endo o)) r) =>
Sem (Writer o : r) a -> Sem r a
writerToEndoWriter forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a (e :: Effect) (f :: * -> *)
       (r :: EffectRow).
m a -> Sem (WithTactics e f m r) (Sem (e : r) (f a))
runT Sem rInitial a
m
        forall (e :: Effect) (r :: EffectRow) a. Sem r a -> Sem (e : r) a
raise forall a b. (a -> b) -> a -> b
$ do
          (Endo o
o, f a
fa) <- forall o (r :: EffectRow) a.
Member (Writer o) r =>
Sem r a -> Sem r (o, a)
listen Sem r (f a)
m'
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ (,) (forall a. Endo a -> a -> a
appEndo Endo o
o forall a. Monoid a => a
mempty) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a
fa
      Pass Sem rInitial (o -> o, x)
m -> do
        Inspector f
ins <- forall (e :: Effect) (f :: * -> *) (m :: * -> *) (r :: EffectRow).
Sem (WithTactics e f m r) (Inspector f)
getInspectorT
        Sem r (f (o -> o, x))
m'  <- forall o (r :: EffectRow) a.
(Monoid o, Member (Writer (Endo o)) r) =>
Sem (Writer o : r) a -> Sem r a
writerToEndoWriter forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a (e :: Effect) (f :: * -> *)
       (r :: EffectRow).
m a -> Sem (WithTactics e f m r) (Sem (e : r) (f a))
runT Sem rInitial (o -> o, x)
m
        forall (e :: Effect) (r :: EffectRow) a. Sem r a -> Sem (e : r) a
raise forall a b. (a -> b) -> a -> b
$ forall o (r :: EffectRow) a.
Member (Writer o) r =>
Sem r (o -> o, a) -> Sem r a
pass forall a b. (a -> b) -> a -> b
$ do
          f (o -> o, x)
t <- Sem r (f (o -> o, x))
m'
          let
            f' :: Endo o -> Endo o
f' =
              forall b a. b -> (a -> b) -> Maybe a -> b
maybe
                forall a. a -> a
id
                (\(o -> o
f, x
_) (Endo o -> o
oo) -> let !o' :: o
o' = o -> o
f (o -> o
oo forall a. Monoid a => a
mempty) in forall a. (a -> a) -> Endo a
Endo (o
o' forall a. Semigroup a => a -> a -> a
<>))
                (forall (f :: * -> *). Inspector f -> forall x. f x -> Maybe x
inspect Inspector f
ins f (o -> o, x)
t)
          forall (m :: * -> *) a. Monad m => a -> m a
return (Endo o -> Endo o
f', forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f (o -> o, x)
t)
{-# INLINE writerToEndoWriter #-}


-- TODO(KingoftheHomeless): Make this mess more palatable
--
-- 'interpretFinal' is too weak for our purposes, so we
-- use 'interpretH' + 'withWeavingToFinal'.

------------------------------------------------------------------------------
-- | A variant of 'Polysemy.Writer.runWriterTVar' where an 'STM' action is
-- used instead of a 'TVar' to commit 'tell's.
runWriterSTMAction :: forall o r a
                          . (Member (Final IO) r, Monoid o)
                         => (o -> STM ())
                         -> Sem (Writer o ': r) a
                         -> Sem r a
runWriterSTMAction :: forall o (r :: EffectRow) a.
(Member (Final IO) r, Monoid o) =>
(o -> STM ()) -> Sem (Writer o : r) a -> Sem r a
runWriterSTMAction o -> STM ()
write = forall (e :: Effect) (r :: EffectRow) a.
(forall (rInitial :: EffectRow) x.
 e (Sem rInitial) x -> Tactical e (Sem rInitial) r x)
-> Sem (e : r) a -> Sem r a
interpretH forall a b. (a -> b) -> a -> b
$ \case
  Tell o
o -> do
    ()
t <- forall (m :: * -> *) (r :: EffectRow) a.
(Member (Final m) r, Functor m) =>
m a -> Sem r a
embedFinal forall a b. (a -> b) -> a -> b
$ forall a. STM a -> IO a
atomically (o -> STM ()
write o
o)
    forall (f :: * -> *) a (e :: Effect) (m :: * -> *)
       (r :: EffectRow).
Functor f =>
a -> Sem (WithTactics e f m r) (f a)
pureT ()
t
  Listen Sem rInitial a
m -> do
    Sem (Writer o : r) (f a)
m' <- forall (m :: * -> *) a (e :: Effect) (f :: * -> *)
       (r :: EffectRow).
m a -> Sem (WithTactics e f m r) (Sem (e : r) (f a))
runT Sem rInitial a
m
    -- Using 'withWeavingToFinal' instead of 'withStrategicToFinal'
    -- here allows us to avoid using two additional 'embedFinal's in
    -- order to create the TVars.
    forall (e :: Effect) (r :: EffectRow) a. Sem r a -> Sem (e : r) a
raise forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) (r :: EffectRow) a.
Member (Final m) r =>
ThroughWeavingToFinal m (Sem r) a -> Sem r a
withWeavingToFinal forall a b. (a -> b) -> a -> b
$ \f ()
s forall x. f (Sem r x) -> IO (f x)
wv forall x. f x -> Maybe x
_ -> forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
      -- See below to understand how this works
      TVar o
tvar   <- forall a. a -> IO (TVar a)
newTVarIO forall a. Monoid a => a
mempty
      TVar Bool
switch <- forall a. a -> IO (TVar a)
newTVarIO Bool
False
      f (f a)
fa     <-
        forall a. IO a -> IO a
restore (forall x. f (Sem r x) -> IO (f x)
wv (forall o (r :: EffectRow) a.
(Member (Final IO) r, Monoid o) =>
(o -> STM ()) -> Sem (Writer o : r) a -> Sem r a
runWriterSTMAction (TVar o -> TVar Bool -> o -> STM ()
writeListen TVar o
tvar TVar Bool
switch) Sem (Writer o : r) (f a)
m' forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ f ()
s))
          forall a b. IO a -> IO b -> IO a
`onException` TVar o -> TVar Bool -> IO o
commitListen TVar o
tvar TVar Bool
switch
      o
o      <- TVar o -> TVar Bool -> IO o
commitListen TVar o
tvar TVar Bool
switch
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) (o
o, ) f (f a)
fa
  Pass Sem rInitial (o -> o, x)
m -> do
    Sem (Writer o : r) (f (o -> o, x))
m'  <- forall (m :: * -> *) a (e :: Effect) (f :: * -> *)
       (r :: EffectRow).
m a -> Sem (WithTactics e f m r) (Sem (e : r) (f a))
runT Sem rInitial (o -> o, x)
m
    Inspector f
ins <- forall (e :: Effect) (f :: * -> *) (m :: * -> *) (r :: EffectRow).
Sem (WithTactics e f m r) (Inspector f)
getInspectorT
    forall (e :: Effect) (r :: EffectRow) a. Sem r a -> Sem (e : r) a
raise forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) (r :: EffectRow) a.
Member (Final m) r =>
ThroughWeavingToFinal m (Sem r) a -> Sem r a
withWeavingToFinal forall a b. (a -> b) -> a -> b
$ \f ()
s forall x. f (Sem r x) -> IO (f x)
wv forall x. f x -> Maybe x
ins' -> forall b. ((forall a. IO a -> IO a) -> IO b) -> IO b
mask forall a b. (a -> b) -> a -> b
$ \forall a. IO a -> IO a
restore -> do
      -- See below to understand how this works
      TVar o
tvar   <- forall a. a -> IO (TVar a)
newTVarIO forall a. Monoid a => a
mempty
      TVar Bool
switch <- forall a. a -> IO (TVar a)
newTVarIO Bool
False
      f (f (o -> o, x))
t      <-
        forall a. IO a -> IO a
restore (forall x. f (Sem r x) -> IO (f x)
wv (forall o (r :: EffectRow) a.
(Member (Final IO) r, Monoid o) =>
(o -> STM ()) -> Sem (Writer o : r) a -> Sem r a
runWriterSTMAction (TVar o -> TVar Bool -> o -> STM ()
writePass TVar o
tvar TVar Bool
switch) Sem (Writer o : r) (f (o -> o, x))
m' forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ f ()
s))
          forall a b. IO a -> IO b -> IO a
`onException` TVar o -> TVar Bool -> (o -> o) -> IO ()
commitPass TVar o
tvar TVar Bool
switch forall a. a -> a
id
      TVar o -> TVar Bool -> (o -> o) -> IO ()
commitPass TVar o
tvar TVar Bool
switch
        (forall b a. b -> (a -> b) -> Maybe a -> b
maybe forall a. a -> a
id forall a b. (a, b) -> a
fst forall a b. (a -> b) -> a -> b
$ forall x. f x -> Maybe x
ins' f (f (o -> o, x))
t forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (f :: * -> *). Inspector f -> forall x. f x -> Maybe x
inspect Inspector f
ins)
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap) forall a b. (a, b) -> b
snd f (f (o -> o, x))
t

  where
    {- KingoftheHomeless:
      'writeListen'/'writePass' is used by the argument computation to a
      'listen' or 'pass' in order to 'tell', rather than directly using
      the provided 'write'.
      This is because we need to temporarily store its
      'tell's locally in order for the 'listen'/'pass' to work
      properly. In the case of 'listen', this is done in parallel with
      the global 'write's. In the case of 'pass', the argument computation
      doesn't use 'write' at all, and instead, when the computation completes,
      commit the changes it made to the local tvar by 'commitPass',
      globally 'write'ing it all at once.
      ('commitListen' serves only as a (likely unneeded)
      safety measure.)

      'commitListen'/'commitPass' is protected by 'mask'+'onException'.
      Combine this with the fact that the 'withWeavingToFinal' can't be
      interrupted by pure errors emitted by effects (since these will be
      represented as part of the functorial state), and we
      guarantee that no writes will be lost if the argument computation
      fails for whatever reason.

      The argument computation to a 'pass' may also spawn
      asynchronous computations which do 'tell's of their own.
      In order to make sure these 'tell's won't be lost once a
      'pass' completes, a switch is used to
      control which tvar 'writePass' writes to. The switch is flipped
      atomically together with commiting the writes of the local tvar
      as part of 'commit'. Once the switch is flipped,
      any asynchronous computations spawned by the argument
      computation will write to the global tvar instead of the local
      tvar (which is no longer relevant), and thus no writes will be
      lost.
    -}

    writeListen :: TVar o
                -> TVar Bool
                -> o
                -> STM ()
    writeListen :: TVar o -> TVar Bool -> o -> STM ()
writeListen TVar o
tvar TVar Bool
switch = \o
o -> do
      Bool
alreadyCommitted <- forall a. TVar a -> STM a
readTVar TVar Bool
switch
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
alreadyCommitted forall a b. (a -> b) -> a -> b
$ do
        o
s <- forall a. TVar a -> STM a
readTVar TVar o
tvar
        forall a. TVar a -> a -> STM ()
writeTVar TVar o
tvar forall a b. (a -> b) -> a -> b
$! o
s forall a. Semigroup a => a -> a -> a
<> o
o
      o -> STM ()
write o
o
    {-# INLINE writeListen #-}

    writePass :: TVar o
              -> TVar Bool
              -> o
              -> STM ()
    writePass :: TVar o -> TVar Bool -> o -> STM ()
writePass TVar o
tvar TVar Bool
switch = \o
o -> do
      Bool
useGlobal <- forall a. TVar a -> STM a
readTVar TVar Bool
switch
      if Bool
useGlobal then
        o -> STM ()
write o
o
      else do
        o
s <- forall a. TVar a -> STM a
readTVar TVar o
tvar
        forall a. TVar a -> a -> STM ()
writeTVar TVar o
tvar forall a b. (a -> b) -> a -> b
$! o
s forall a. Semigroup a => a -> a -> a
<> o
o
    {-# INLINE writePass #-}

    commitListen :: TVar o
                 -> TVar Bool
                 -> IO o
    commitListen :: TVar o -> TVar Bool -> IO o
commitListen TVar o
tvar TVar Bool
switch = forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
      forall a. TVar a -> a -> STM ()
writeTVar TVar Bool
switch Bool
True
      forall a. TVar a -> STM a
readTVar TVar o
tvar
    {-# INLINE commitListen #-}

    commitPass :: TVar o
               -> TVar Bool
               -> (o -> o)
               -> IO ()
    commitPass :: TVar o -> TVar Bool -> (o -> o) -> IO ()
commitPass TVar o
tvar TVar Bool
switch o -> o
f = forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ do
      o
o <- forall a. TVar a -> STM a
readTVar TVar o
tvar
      let !o' :: o
o' = o -> o
f o
o
      -- Likely redundant, but doesn't hurt.
      Bool
alreadyCommitted <- forall a. TVar a -> STM a
readTVar TVar Bool
switch
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
alreadyCommitted forall a b. (a -> b) -> a -> b
$
        o -> STM ()
write o
o'
      forall a. TVar a -> a -> STM ()
writeTVar TVar Bool
switch Bool
True
    {-# INLINE commitPass #-}
{-# 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 :: forall o (e :: Effect) (r :: EffectRow) a.
Monoid o =>
(forall (m :: * -> *) x.
 Monad m =>
 Weaving e (WriterT o m) x -> WriterT o m x)
-> Sem (e : r) a -> Sem r (o, a)
interpretViaLazyWriter forall (m :: * -> *) x.
Monad m =>
Weaving e (WriterT o m) x -> WriterT o m x
f Sem (e : r) a
sem = forall (r :: EffectRow) a.
(forall (m :: * -> *).
 Monad m =>
 (forall x. Union r (Sem r) x -> m x) -> m a)
-> Sem r a
Sem forall a b. (a -> b) -> a -> b
$ \(forall x. Union r (Sem r) x -> m x
k :: forall x. Union r (Sem r) x -> m x) ->
  let
    go :: forall x. Sem (e ': r) x -> Lazy.WriterT o m x
    go :: forall x. Sem (e : r) x -> WriterT o m x
go = forall (m :: * -> *) (r :: EffectRow) a.
Monad m =>
(forall x. Union r (Sem r) x -> m x) -> Sem r a -> m a
usingSem forall a b. (a -> b) -> a -> b
$ \Union (e : r) (Sem (e : r)) x
u -> case forall (e :: Effect) (r :: EffectRow) (m :: * -> *) a.
Union (e : r) m a -> Either (Union r m a) (Weaving e m a)
decomp Union (e : r) (Sem (e : r)) x
u of
      Right (Weaving e (Sem rInitial) a
e f ()
s forall x. f (Sem rInitial x) -> Sem (e : r) (f x)
wv f a -> x
ex forall x. f x -> Maybe x
ins) -> forall (m :: * -> *) x.
Monad m =>
Weaving e (WriterT o m) x -> WriterT o m x
f forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) (e :: Effect) (rInitial :: EffectRow) a
       resultType (mAfter :: * -> *).
Functor f =>
e (Sem rInitial) a
-> f ()
-> (forall x. f (Sem rInitial x) -> mAfter (f x))
-> (f a -> resultType)
-> (forall x. f x -> Maybe x)
-> Weaving e mAfter resultType
Weaving e (Sem rInitial) a
e f ()
s (forall x. Sem (e : r) x -> WriterT o m x
go forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x. f (Sem rInitial x) -> Sem (e : r) (f x)
wv) f a -> x
ex forall x. f x -> Maybe x
ins
      Left Union r (Sem (e : r)) x
g -> forall w (m :: * -> *) a. m (a, w) -> WriterT w m a
Lazy.WriterT forall a b. (a -> b) -> a -> b
$ do
        ~(o
o, x
a) <- forall x. Union r (Sem r) x -> m x
k forall a b. (a -> b) -> a -> b
$
          forall (s :: * -> *) (n :: * -> *) (m :: * -> *) (r :: EffectRow)
       a.
(Functor s, Functor n) =>
s ()
-> (forall x. s (m x) -> n (s x))
-> (forall x. s x -> Maybe x)
-> Union r m a
-> Union r n (s a)
weave
            (forall a. Monoid a => a
mempty, ())
            (\ ~(o
o, Sem (e : r) x
m) -> (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first) (o
o forall a. Semigroup a => a -> a -> a
<>) (forall o (e :: Effect) (r :: EffectRow) a.
Monoid o =>
(forall (m :: * -> *) x.
 Monad m =>
 Weaving e (WriterT o m) x -> WriterT o m x)
-> Sem (e : r) a -> Sem r (o, a)
interpretViaLazyWriter forall (m :: * -> *) x.
Monad m =>
Weaving e (WriterT o m) x -> WriterT o m x
f Sem (e : r) x
m))
            (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd)
            Union r (Sem (e : r)) x
g
        forall (m :: * -> *) a. Monad m => a -> m a
return (x
a, o
o)
    {-# INLINE go #-}
  in do
    ~(a
a,o
s) <- forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
Lazy.runWriterT (forall x. Sem (e : r) x -> WriterT o m x
go Sem (e : r) a
sem)
    forall (m :: * -> *) a. Monad m => a -> m a
return (o
s, a
a)
{-# INLINE interpretViaLazyWriter #-}