How do I write this more general version of `Control.Monad.Writer.censor`?

Question

The standard Control.Monad.Writer.censor is dual to the standard Control.Monad.Reader.local, with censor modifying the writer state after the computation and local modifying the reader state before the computation:

censor :: (w -> w) -> Writer w a -> Writer w a
local  :: (r -> r) -> Reader r a -> Reader r a

However, the Reader and Writer monads are not entirely symmetric. Namely, a writer computation produces a result, in addition to a writer state, and I'm trying to write an alternative version of censor that takes advantage of this asymmetry. I want to write a function

censorWithResult :: (a -> w -> w) -> Writer w a -> Writer w a

which takes a transformer of type a -> w -> w that receives the result of the computation in addition to the writer state. I don't see how to write this function using tell, listen, and pass.

The precise behavior I expect of censorWithResult is that if

ma :: Writer w a
f  :: a -> w -> w

and

runWriter ma = (r , y)

then

runWriter (censorWithResult f ma) = (r , f r y)

whereas

runWriter (censor g ma) = (r , g y)

when g :: w -> w.

Example

This shouldn't be necessary to understand the question, but here's a simplified version of the motivating example:

import Control.Applicative
import Control.Monad.Writer

-- Call-trace data type for functions from 'Int' to 'Int'.
--
-- A 'Call x subs r' is for a call with argument 'x', sub calls
-- 'subs', and result 'r'.
data Trace = Call Int Forest Int
type Forest = [Trace]

-- A writer monad for capturing call traces.
type M a = Writer Forest a

-- Recursive traced negation function.
--
-- E.g. we expect that
--
--   runWriter (t 2) = (-2 , Call 2 [Call 1 [Call 0 [] 0] -1] -2)
t , n :: Int -> M Int
t x = trace n x
n x = if x <= 0 then pure 0 else subtract 1 <$> t (x - 1)

trace :: (Int -> M Int) -> (Int -> M Int)
trace h x = do
  censorWithResult (\r subs -> [Call x subs r]) (h x)

-- The idea is that if 'ma :: Writer w a' and 'runWriter ma = (r , y)'
-- then 'runWriter (censorWithResult f ma) = (r , f r y)'. I.e.,
-- 'censorWithResult' is like 'Control.Monad.Writer.censor', except it
-- has access to the result of the 'Writer' computation, in addition
-- to the written data.
censorWithResult :: (a -> w -> w) -> Writer w a -> Writer w a
censorWithResult = undefined

Answer 1

The precise behavior I expect of censorWithResult is that if

ma :: Writer w a
f  :: a -> w -> w

and

runWriter ma = (r , y)

then

runWriter (censorWithResult f ma) = (r , f r y)

Okay, let's do that, then. The only thing you need to know is that writer is a left inverse for runWriter. Then we get the following chain of equalities, first by appling writer to both sides, then by eliminating the left inverse.

        runWriter (censorWithResult f ma)  =        (r, f r y)
writer (runWriter (censorWithResult f ma)) = writer (r, f r y)
                   censorWithResult f ma   = writer (r, f r y)

The only thing we need to do now is plug in your equation runWriter ma = (r, y):

censorWithResult f ma = let (r, y) = runWriter ma in writer (r, f r y)

Ain't equational reasoning grand?

Answer 2

If we are allowed to use only tell, pass and listen, the only function that is able to access output is

-- | `pass m` is an action that executes the action `m`, which returns a value
-- and a function, and returns the value, applying the function to the output. 
pass :: (MonadWriter w m) => m (a, w -> w) -> m a

So for censorWithResult, we need to partially apply a given function of type a -> w -> w to get w -> w and handle it to pass. This can be accomplished as

censorWithResult :: (MonadWriter w m) => (a -> w -> w) -> m a -> m a
censorWithResult f m = pass $ do
    a <- m
    return (a, f a)

The action inside pass executes the given action, partially applies f to it and pass then modifies the output accordingly.