Reader-Writer-State

As I said in the introduction to monad transformers, it isn't uncommon to have a computation that needs some state, some read-only context, and the ability to produce a log. To make this easier than stacking three monad transformers on top of each other, we have the reader-writer-state monad RWS. In fact, we may want to be able to stack RWS on top of some other monad, so we also have a transformer version of RWS, called RWST:

newtype RWST r w s m a = RWST { runRWST :: r -> s -> m (a, s, w) }

I know, the number of type variables is getting ridiculous, but the names give you a clue: r is the type of our read-only context. w is our writer monoid, s is the state, m is the underlying monad to be wrapped, and a is the return type of the computation.

The function type that RWST wraps is also fairly natural if you understand ReaderT, WriterT, and StateT. ReaderT wraps a function that takes the context r as argument, but this context, being read-only, does not show up in the return type. WriterT produces a log value of type w, but it does not take a w as argument; the log is write-only. And StateT threads the state through the whole computation because this state is readable and writable, so StateT wraps a function that takes a state of type s as argument and also returns an updated state of type s. RWST does all this in one go. The function it wraps takes a context r and a state s as arguments, and the return value consists of a value of type a, the updated state s, and some addition w to the log.

The Monad instance for RWST implements the behaviours of the Monad instances of all three of ReaderT, WriterT, and StateT:

instance (Monoid w, Monad m) => Monad (RWST r w s m) where
    return x = RWST $ \_ -> s -> return (x, s, mempty)
    x >>= f  = RWST $ \r -> s -> do
        (y, s',  wx) <- runRWST x     r s
        (z, s'', wf) <- runRWST (f y) r s'
        return (z, s'', wx <> wf)

This may all be a bit overwhelming, but it is easy to understand if you focus on one component of the returned tuples. For example, the s component behaves exactly as for StateT: return x leaves the state unchanged. x >>= f runs x on the current state s. This returns a result y and an updated state s'. We run f y on the updated state s'. This produces an updated state s'', which is what we return as part of the result of the computation. The context r is not part of any tuple we return. We just pass it to both x and f y when we run them, just as in the implementation of ReaderT. As for the log, we once again combine the log values produced by x and f using our monoid multiplication (<>). That's the third component of the triple returned by an RWST computation.

ask, asks, get, gets, put, modify, and tell work also for RWST. I won't show their implementations here. The RWS monad is once again produced by stacking RWST on top of the Identity monad:

type RWS r w s = RWST r w s Identity

runRWS f r s = runIdentity $ runRWST f r s

We won't discuss an example of using RWS or RWST here. You use it exactly as you would a monad obtained by stacking ReaderT, WriterT, and StateT, only there is much less lifting involved. If you were to build RWS by hand using

type MyRWS r w s = ReaderT r (WriterT w (State s))

then running a computation of this type requires us to call all three: runReaderT, runWriterT, and State:

runMyRWS :: MyRWS r w s a -> r -> s -> (a, s, w)
runMyRWS f r s = runState (runWriterT (runReaderT f r)) s

and writing a new state would require¹

do
    ...
    lift $ lift $ put s
    ...

because our State monad is the bottom-most monad, hidden under two layers of ReaderT and WriterT. RWS makes this much easier. It comes with a function runRWS that does the same as runMyRWS; we don't have to unwrap each layer of decoration individually. And all the abilities of Reader, Writer, and State are part of the RWS monad, so no lifting is required:

do
    ...
    ctxt <- ask
    ...
    put s
    ...
    tell w
    ...

---

1. Well, this would be true if we didn't have the MonadState, MonadReader, and MonadWriter type classes. As we will discuss shortly, any monad that is an instance of MonadState supports get, gets, put, and modify. Any monad that is an instance of MonadReader supports ask and asks. And any monad that is an instance of MonadWriter supports tell.

   When we stack ReaderT and WriterT on top of a monad that is an instance of MonadState, which State is, then the resulting monad is also an instance of MonadState. So, we can in fact use get and put without the need for any lifting in our manually built reader-writer-state monad. Still, working with RWS is more convenient, and more efficient, than using a manually constructed transformer stack. ↩