{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections   #-}

module Polysemy.Writer
  ( -- * Effect
    Writer (..)

    -- * Actions
  , tell
  , listen
  , censor

    -- * Interpretations
  , runWriter

    -- * Interpretations for Other Effects
  , runOutputAsWriter
  ) where

import Polysemy
import Polysemy.Output
import Polysemy.State


------------------------------------------------------------------------------
-- | An effect capable of emitting and intercepting messages.
data Writer o m a where
  Tell   :: o -> Writer o m ()
  Listen :: ∀ o m a. m a -> Writer o m (o, a)
  Censor :: (o -> o) -> m a -> Writer o m a

makeSem ''Writer


------------------------------------------------------------------------------
-- | Transform an 'Output' effect into a 'Writer' effect.
runOutputAsWriter :: Member (Writer o) r => Sem (Output o ': r) a -> Sem r a
runOutputAsWriter = interpret $ \case
  Output o -> tell o
{-# INLINE runOutputAsWriter #-}


------------------------------------------------------------------------------
-- | Run a 'Writer' effect in the style of 'Control.Monad.Trans.Writer.WriterT'
-- (but without the nasty space leak!)
runWriter
    :: Monoid o
    => Sem (Writer o ': r) a
    -> Sem r (o, a)
runWriter = runState mempty . reinterpretH
  (\case
      Tell o -> do
        modify (`mappend` o) >>= pureT
      Listen m -> do
        mm <- runT m
        -- TODO(sandy): this is fucking stupid
        (o, fa) <- raise $ runWriter mm
        pure $ fmap (o, ) fa
      Censor f m -> do
        mm <- runT m
        ~(o, a) <- raise $ runWriter mm
        modify (`mappend` f o)
        pure a
  )