Lecture 23: Monad Transformer Library: Implementation

Today I'd like to review the monad transformer library (mtl) with an eye to its overall construction. There is some very clever engineering here, and you can learn a lot from the careful study of the work of masters. But there is another reason. Some of the cleverness leaks out, in that the documentation of the unadorned State monad and similar monads can be confusing to the uninitiated. For starters, if you look where you expect in the documentation to be in Control.Monad.State, you get pointed to Control.Monad.State.Lazy, but it starts out by talking about the type class MonadState, which is something different than the State monad, although they seem somehow related. And then, the documentation for State and StateT doesn't contain links to source code, and that makes us unhappy. But then, if you dig around a bit more, you'll find Control.Monad.Trans.State, which points us to Control.Monad.Trans.State.Lazy, and that module contains what seem to be duplicate definitions of State and StateT, this time with source links, but MonadState is now nowhere to be found. But what about our imports? We've been importing Control.Monad.State, and it's been working for us. Should we keep doing that? Should we import Control.Monad.Trans.State, or maybe even Control.Monad.Trans.State.Lazy? Or should we just go home, have a good cry, and reconcile ourselves to a life of writing JavaScript? No, not that. The short answer is that we should import Control.Monad.State, but get our documentation from Control.Monad.State.Lazy. The rest of the lecture constitutes the long answer.

The fundamental idea behind the mtl is that derived types can preserve structure (i.e., capabilities of the types from which they're derived), and that preserved structure can be realized programmatically via automatically derived instances of typeclasses. This has the effect at library use time of considerable functionality “coming along for free,” as we saw when we implemented Parser.hs using monad transformers: not only did we get put and friends (as you'd expect from something built using StateT), we got Functor, Applicative, Monad, and other typeclass instances.

IdentityT

There is, unsurprisingly, an IdentityT monad transformer. The definition of IdentityT can't be surprising:

newtype IdentityT m a = IdentityT { runIdentityT :: m a }

This is just an untransformed m a in a (virtual) IdentityT box. We'll look at IdentityT as a simple setting in which we can get our bearings. First of all, IdentityT is a monad transformer, and so is itself an instance of the MonadTrans class:

class MonadTrans t where lift :: Monad m => m a -> t m a

i.e., IdentityT possesses a lift function that maps the untransformed monad m a into the transformed monad IdentityT m a. Instances t of MonadTrans are expected to satisfy a couple of laws:

These look mysterious, in part because the types are hidden from us, but they're well motivated and so are easy to understand. E.g., if we start with a value object of type a, there are two natural ways to embed it into the transformed monad t m a. The first is directly via t m's return function, the second is indirectly, from a to m a via m's return function, and from there to t m a via t m's lift. The first law says that two ways must produce the same value in t m a.

Lift Law Digram

The second law says something similar. If we have a binding m >>= f in the untransformed m a, there are two natural ways that we can go about lifting the result into the transformed monad t m a. The first, as before, is directly via t m a's lift function. The second is by lifting (>>=)'s arguments, resulting in lift m and lift . f, and applying (>>=) in the transformed monad t m a. Again, these two different natural ways of lifting a binding must produce the same value in t m a.

The instance definition for MonadTrans IdentityT is straightforward—we just take a value in m a, and drop it in a (virtual) IdentityT box.

instance MonadTrans IdentityT where lift = IdentityT

More complicated monad transformers will have more complicated definitions of lift.

Next up, we have instances (when appropriate) for IdentityT m of Functor, Applicative, Foldable, and of course Monad, MonadPlus, and a few other type classes we haven't met yet. These are all pretty simple, with all of the activity taking the form of traffic in (virtual) boxing and unboxing. To this end, the following helper functions are defined:

mapIdentity :: (m a -> n b) -> IdentityT m a -> IdentityT n a mapIdentity f = IdentityT . f . runIdentityT lift2IdentityT :: (m a -> n b -> p c) -> IdentityT m a -> IdentityT n b -> IdentityT p c lift2IdentityT f a b = IdentityT (f (runIdentityT a) (runIdentityT b))

Note that lift2IdentityT is just a binary version of mapIdentity, it's impressive type notwithstanding, and both simply lift functions on untransformed monads to functions on transformed monads. We're now ready to define some instances:

instance (Functor m) => Functor (IdentityT m) where fmap f = mapIdentityT (fmap f) instance (Foldable f) => Foldable (IdentityT f) where foldMap f (IdentityT a) = foldMap f a instance (Applicative m) => Applicative (IdentityT m) where pure x = IdentityT (pure x) (<*>) = lift2IdentityT (<*>) instance (Monad m) => Monad (IdentityT m) where return = IdentityT . return m >>= k = IdentityT $ runIdentityT . k =<< runIdentityT m

The most interesting thing here is the definition of (>>=) in the monad instance. Note the use of a right-to-left version (=<<) of bind, and how this mixes better with (.). We solved the same problem earlier, but by resolving the follow of compositions and binds in the other direction, using (>>>). The critical thing to notice is how the we're using the (>>=) from the untransformed monad to define (>>=) in the transformed monad. Our code is all about making sure that things get boxed or unboxed as needed. Now, let's look at a couple of more interesting monad transformers.

StateT

The StateT monad transformer has everything going for it so far as this lecture is concerned. It's powerful and even somewhat familiar. But there are a few interesting twists and turns in its implementation, and understanding it will go a long way to helping you understand how the monad transformer library as a whole is implemented, and why. We'll start by considering the definition:

newtype StateT s m a = StateT { runStateT :: s -> m (a,s) }

This type isn't really surprising, the only question being “where do we stick the inner monad m into the state transformation function?” There are at least three plausible answers,

Perhaps unsurprisingly, the choice taken is one that can't be obtained by simply stacking untransformed monads. Of course, we use StateT as a way to define functions that we're latter going to want to run via runStateT or one of its friends. But there's a trickiness with any of these, because the result returned is going to be in terms of the untransformed monad m. This is explicit in runStateT :: StateT s m a -> s -> m (a,s), but requires a bit more work with the others, e.g.,

evalStateT :: (Monad m) => StateT s m a -> s -> m a evalStateT m s = do (a, _) <- runStateT m s return a

Note here that I'm pulling a few minor cheats on you, as I will throughout this lecture (and indeed, already had). This is the implementation of evalStateT from Control.Monad.Trans.State.Strict, not Lazy. The difference is a minor technicality in pattern matching that's not important for this lecture. I'm making a few other pedagogical simplifications too, but nothing you can't puzzle out from the actual code.

Of course, we want types that result from applying StateT s to a monad m to be monads themselves, hence

instance Monad m => Monad (StateT s m) where return = lift . return m >>= f = StateT $ \s -> do (a,t) <- runStateT m s runStateT (f a) t

Here, the focus should be on the definition of (>>=), in which the body of the lambda form \s -> ... is an object of type m (a,s). As before, we rely on (>>=) in the untransformed monad m (a,s) (here hidden by do sugar) to extract the (result,newState) pair, which is then handed off to runStateT again. This is essentially the same as the naïve code for (>>=) in State, albeit with a monadic binding (>>=) replacing a let binding. We'll see this again.

Instances of Functor, Applicative, Alternative, and MonadPlus and similar type classes are also provided, when the untransformed monad has the corresponding instance. It's worth noting the that in the particular case of Applicative, it suffices that the inner type (not necessarily a monad) merely be a functor.

We make StateT into an instance of a MonadTrans by providing an appropriate lift:

instance MonadTrans (StateT s) where lift m = StateT $ \s -> do a <- m return (a,s)

Now that Haskell requires that every monad is a functor, we can have write lift more concisely as

instance MonadTrans StateT where lift m = StateT$ \s -> (`(,)` s) <*> m

The implementations of get, put and modify, are built on that of state:

state :: (s -> (a,s)) -> StateT s m a state f = StateT $ return . f

i.e., state puts in a StateT box the result of applying the function f to a state, obtaining a (result,newState) pair, which is then lifted into the monad m (a,s) via return. With this:

get :: StateT s m s get = state $ \s -> (s,s) put :: s -> StateT s m () put t = state $ \s -> ((),t) modify :: (s -> s) -> StateT s m () modify f = state $ \s -> ((),f s)

Now we encounter that extra layer of trickiness, alluded to at the beginning of the lecture. The class definition for StateT isn't in Control.Monad.State, or even Control.Monad.State.Lazy, it's in Control.Monad.Trans.State.Lazy, for two reasons. First of all, Control.Monad.State.Lazy relies on functional dependencies, an extension to Haskell implemented in GHC, but not in all Haskell systems. Non-GHC systems have to include Control.Monad.Trans.State, and lose the advantages of the new facilities in Control.Monad.State.Lazy. Second, for GHC, we're going to hide the definitions of put, get, and state from Control.Monad.Trans.State.Lazy, exposing instead the definitions that come from the MonadState type class:

class Monad m => MonadState s m | m -> s where get :: m s put :: s -> m () state :: (s -> (a, s)) -> m a

Note the functional dependency (aka, fundep) | m -> s in the class declaration. This is indicates a commitment from the programmer to GHC that there can't be distinct s and t for which there are MonadState s m and MonadState t m instances.To define an instance of MonadState, it is necessary to define either get and put, or state. The instance for StateT s m is particularly easy:

import qualified Control.Monad.Trans.State.Lazy as Lazy (StateT, get, put, state) instance Monad m => MonadState s (Lazy.StateT s m) where get = Lazy.get put = Lazy.put state = Lazy.state

But the virtue of MonadState becomes apparent if we take a look back at IdentityT. The punchline here is the following instance declaration:

instance MonadState s m => MonadState s (IdentityT m) where get = lift get put = lift . put state = lift . state

The point here is that if m is a state monad with state s, then IdentityT m is also a state monad with state s. Or, to say the same thing in slightly different words, IdentityT preserves the MonadState s property. And this saves us from having to remember how many layers deep the StateT actually is, and therefore how many times we have to apply lift to get the state in or out of the monad. We can just use get, put, state, and their friends on the transformed monad, is if it were a plain old-fashioned state monad. That is very convenient.

WriterT

The Writer monad is usually used to implement logging. I think of a Writer as being half of a State, in that we're producing something like a state on output, we're just not consuming state on input. And like the State monad's state, the log we're building is hidden.

We'll do a quick example once we made some introductions:

class Writer w a = { runWriter :: (a,w) }

Our first problem is “how do we make Writer into a monad.” The hard part of the job is usually in defining (>>=). In the state monad, (>>=) handled the states by hooking up the output state of the first argument's state-transition function to the input state of the second monad, which we got by applying the second argument, a Kleisli arrow, to the value that came through the binding. That way, we ended up with one input state, and one output state, and all was well. This can't work the same way with writers, because there's no input on the second monad to attach the output of the first. The plumbing is just different. Instead, we'll combine the logs, and this requires some sort of binary operator, which is going to have to be associative if (>>=) is to be associative. The implementation of return will require a identity. So we're going to need to require that w is a Monoid, because that's precisely what monoids give us: an associative binary operator with identity.

instance (Monoid w) => Monad (Writer w) where return a = Writer $ (a,mempty) m >>= f = Writer $ let (a,w1) = runWriter m (b,w2) = runWriter $ f a in (b,w1 <> w2)

The practical use of writer monads relies, much like state monads, on a scant handful of functions that enable us to interface with the monad, without dwelling on its implementation. The first of these is tell, which is similar to put, in that it writes its argument to the log:

tell :: (Monoid w) => w -> Writer w () tell w = Writer ((), w)

The listen function is analogous to State's get, but different in that it takes a writer as an argument, and it returns the result returned by the argument monad along with its log:

listen :: (Monoid w) => Writer w a -> Writer w (a,w) listen m = Writer $ let (a, w) = runWriter m in return ((a, w), w)

The censor function is analogous to State's modify, in much the same way that listen is analogous to put, in that it takes a writer monad as an argument, and applies the argument function to its argument's log, packaging the result.

censor :: (Monoid w, Monad m) => (w -> w) -> Writer w a -> Writer w a censor f m = Writer $ let (a, w) = runWriter m in return (a, f w)

The following, somewhat contrived, example illustrates the use of Writer:

import Control.Monad.Writer type LogApply = Writer [String] logFunction :: (Show a, Show b) => String -> (a -> b) -> a -> LogApply b logFunction name f = \a -> do let b = f a tell $ ["Applying " ++ name ++ " to " ++ show a ++ " resulting in " ++ show b ++ "."] return b logSquare, logDouble, logSuccessor, logPredecessor :: (Show n,Num n) => n -> LogApply n logSquare = logFunction "square" (^2) logDouble = logFunction "double" (*2) logSuccessor = logFunction "successor" (+1) logPredecessor = logFunction "predecessor" (subtract 1) computation :: LogApply Integer computation = return 10 >>= logPredecessor >>= logSquare >>= logSuccessor >>= logDouble main :: IO () main = do let log = execWriter computation putStr $ unlines log

When run, this produces the output:

Applying predecessor to 10 resulting in 9. Applying square to 9 resulting in 81. Applying successor to 81 resulting in 82. Applying double to 82 resulting in 164.

This is, perhaps, not truly inspiring, but it does illustrate the sort of thing that Writer makes easy.

Of course, this whole discussion about Writer glosses over the same sort of complexities we saw with State. In reality, there's a WriterT monad transformer, where

newtype WriterT w m a = WriterT { runWriterT :: m (a,w) } type Writer w = WriterT w Identity

In this case the contribution of WriterT isn't in the type (which we could arrange by ordinary stacking), but in the plumbing,

instance (Monad m, Monoid w) => Monad (WriterT w m) where return a = WriterT $ return (a,mempty) m >>= g = WriterT $ do (a,w1) <- runWriterT a (b,w2) <- runWriterT (g a) return (b,w1 <> w2)

Note how this adapts the definition we gave for (>>=) for an ordinary writer, by replacing the let bindings by monadic bindings, packaging our results with return.

Similar changes get made to tell, listen and censor, e.g.,

listen :: (Monad m, Monoid w) => WriterT w m a -> WriterT w m (a,w) listen m = Writer $ do (a, w) <- runWriter m return ((a, w), w)

As with State, these types and functions get defined in Control.Monad.Trans.Writer.Lazy. And in Control.Monad.Writer.Lazy, contains the MonadWriter type class, and associated functions and instances:

class (Monoid w, Monad m) => MonadWriter w m | m -> w where writer :: (a,w) -> m a tell :: w -> m () listen :: m a -> m (a, w) censor :: MonadWriter w m => (w -> w) -> m a -> m a censor f m = do (a,w) <- runWriterT m return (a, f w)

Naturally, monads built by the WriterT w monad transformer are instances of MonadWriter w:

instance (Monoid w, Monad m) => MonadWriter w (Lazy.WriterT w m) where writer = Lazy.writer tell = Lazy.tell listen = Lazy.listen

Likewise, any monad m with a MonadWriter w m instance remains a MonadWriter w when wrapped with IdentityT or a StateT s:

instance MonadWriter w m => MonadWriter w (IdentityT m) where writer = lift . writer tell = lift . tell listen = Identity.mapIdentityT listen instance MonadWriter w m => MonadWriter w (Lazy.StateT s m) where writer = lift . writer tell = lift . tell listen = Lazy.liftListen listen

Likewise, any monad m with a MonadState s m instance remains a MonadState s m when wrapped with a WriterT w, and the associated code is found in Control.Monad.State.Lazy.

And that's the way the whole mlt is built. Associated with each monad transformer is a type class that defines the properties that transformer adds. For each pair of a transformer and another transformer's associated type class, there's an instance of that type class, when possible, that preserves that property in a transformed monad. The exception to all of this is IO, as it so often is, but that's a puzzle for another day.