Composing state with functions and lenses

Here’s the scenario: you’re writing some stateful code. Maybe it’s threaded state (a State monad), or maybe it’s just fixed shared state (a Reader monad). So you’ve got a lot of types flying around like:

Reader X a Reader X a -> Reader X b State X a State X a -> State X b

But then you take your stateful code, and try to compose it with someone else’s stateful code, and their state is different. That is, they have:

Reader Y a Reader Y a -> Reader Y b State Y a State Y a -> State Y b

Question: how do you get these pieces of code to work together?

Clearly you’ll need some kind of relationship between the types X and Y, or you have no hope. But what kind of relationship do you need here? We’ll consider each of the types in turn.

Case 1: Reader X a / Reader Y a

In this case, you’ve got a Reader X a, and a Reader Y a, and you want the combine them. It turns out all you need here is a function from one to the other, and you can turn these into compatible types to compose them nicely. The following is in the mtl package already.

withReader :: (p -> q) -> Reader q a -> Reader p a

That’s not surprising, actually. After all, Reader x y is conceptually just a newtype wrapper around x -> y, so withReader is a fancy name for function composition!

withReader f r = reader (runReader r . f)

Note the contravariance there… you pass in a function p -> q, but what you get back is Reader q a -> Reader p a, in the opposite order. That makes a lot of sense, though, if you think it through. (Exercise for the reader: think it through.)

Case 2: (Reader X a -> Reader X b) / (Reader Y a -> Reader Y b)

Another situation that comes up is that we’ve got a way of wrapping reader monads. This happens particularly often if you’re building up values by making changes to other values. For example, one of the two primitives from the MonadReader class, local, gives you precisely this kind of map between Reader monads.

The first thing we notice here is that a function from one state type to the other cannot possibly be good enough, because a conversion doesn’t even have any clear meaning on those types. What turns out to work for us, though, is a lens. A lens can be thought of as a getter/setter pair, and I’ll use the definition from the excellent fclabels package. Here’s what you need to know:

data a :-> b lens :: (a -> b) -> (b -> a -> a) -> (a :-> b) getL :: (a :-> b) -> (a -> b) setL :: (a :-> b) -> (b -> a -> a)

In other words, a :-> b (note the colon) is the type of lenses from a to b. You construct them by providing a getter and a setter to the lens function, and you can extract the getter and setter from getL and setL. They can also be composed like functions, and have identities (in other words, they form the arrows of a category).

With both getters and setters in mind, we can set out to compose the types earlier.

wrapWithReader :: (x :-> y) -> (Reader y a -> Reader y b) -> (Reader x a -> Reader x b) wrapWithReader l f r = reader (\x -> runReader (f (reader (\y -> runReader r (setL l y x)))) (getL l x))

This may look complex, but mostly the complexity is in constructing and deconstructing the Reader monad newtype. The definition is straight-forward aside from that: to turn the Reader x a into a corresponding Reader x b, you simply consider the Reader y a that results from fixing the non-y bits of the input, transform it, and then map it back.

Case 3: State X a / State Y a

The third case is where we have a state monad rather than a reader monad. Since changes of state in the state monad are almost interchangeable with a lens, it turns out a lens is what we need here, too. We can implement this without too much trouble.

withState :: (x :-> y) -> State y a -> State x a withState l s = state (\x -> let (a,y) = runState s (getL l x) in (a, setL l y x))

In other words, we pull the y out of the x, run the state computation with it, and then push the resulting y back into the x to get a modified x. Works like a charm.

Case 4: (State X a -> State X b) / (State Y a -> State Y b)

The final case, and the most complicated one yet, arises if you have a function to modify state types, and need to change the type of the state. Sadly, even a lens is not sufficient to assign a reasonable meaning to this conversion. To make sense of such a transformation, you need to know something even stronger: we’ll do it where there is an isomorphism between the types X and Y. Then the composition can be seen as transforming the functions by simply treating one as the other.

Fortunately, fclabels has the types we need still!

data x :<->: y (:<->:) :: (x -> y) -> (y -> x) -> (x :<->: y) fw :: (x :<->: y) -> (x -> y) bw :: (x :<->: y) -> (y -> x)

An isomorphism is just a pair of inverse functions between the types, meaning they are essentially interchangeable. Then it’s easy to build the state wrapper converter, which interchanges them:

wrapWithState :: (x :<->: y) -> (State y a -> State y a) -> (State x a -> State x a) wrapWithState iso f = t (bw iso) (fw iso) . f . t (fw iso) (bw iso) where t fw bw s = state (second fw . runState s . bw)

And voila, composable state types!

(Side note: It is possible to build something with the type given above for wrapWithState but using just a lens instead of an isomorphism. Sadly, it doesn’t act as you’d expect for doing the composition. Also, Gregory Collins pointed out to me that you can implement wrapWithState with the expected behavior and just a lens, if you give it a rank 2 type and require that the wrapper function be universally quantified on the result type. Neither of these quite what we’re looking for, though, and the isomorphism is needed to get something with the obvious meaning.)

(Second side note: I’ve done this with pure state and reader monads for simplicity; but it’s easy to generalize to StateT and ReaderT, if that’s what you want.)