What is the State monad? functional programming with state
This post was written after working through the State chapter in the excellent Functional Programming in Scala with a group of colleagues, and intends to capture some of our discussions towards gaining an intuition about the State monad.
The main idea:
State[S, A]is a computation – it can produce a value of type
A, or do some more computation with intermediate state of type
State is a functional way of describing multi-step computation where you have some internal intermediate state. There’s no mutation of the state, but the internal immutable state that’s carried around with the computation can change at each step.
You can think of the
A type parameter as describing the value that you can get out of the state during the course of a function. Similarly, the
S type represents the internal computation state. (You end up saying “state” a lot when talking about this. Try to disambiguate between the
State type that models the computation and the internal state that is maintained while running the computation.)
Of course, as is clear from the definition below, there’s no state and no values stored in the
State itself. Rather, you can use the things of type
A in the course of composing functions that are applied to it, as you build up a computation. The state and values exist only as function arguments that get threaded through the computation.
Commonly this is used in a for-comprehension, where it looks like a sequence of imperative operations. In reality the internal state is threaded through, sort of behind the scenes.
The implementation in the book looks like this. The Scalaz implementation is a lot less straightforward, but captures the same idea.
You can use
State as a pure way to run a computation where you need to modify state but want to hide the intermediate state of a computation.
State eliminates a certain class of bugs. The motivation in the book is in using an immutable random number generator by manually composing functions that have a similar signature
Rand[A] => (A, Rand[A]). You need to pass the right version of the state through to the next function, else you lose an update to your internal state. Using
State, especially in a for-comprehension, this all happens automatically, and you can’t get it wrong.
On the other hand I’ve heard comments like “I’ve never actually used it in production”, usually accompanied by descriptions of using related structures. So I’m thinking of this as a pedagogical tool for learning about building a program by composing functions, expecting that as we as functional programming novices progress, we will soon learn about more sophisticated structures based on the same ideas (I’m looking at you,
ReaderWriterStateT) that have more direct application. Thinking about a program in this way is at least a new view on how to compose functionality together.
Example: swapping head of a stack
Imagine swapping the top two elements of a stack. At each line in the for-comprehension the right-hand-side is some combinator acting on the state, and the type of the thing on the left is the value type from the
State on the right-hand-side.
(I’m totally ignoring the case where the input stack has few elements. Just pretend, and we’ll cover the error handling aspect another day.)
As with many constructions around monads, nothing actually happens until right at the end. Most of your code is describing the program that will be run, effectively building a single function that takes your starting state and spits out an end state. Only afterwards do you actually run the function. This is very different to imperative programming where you can imagine that each operation is run as it is encountered.
If you want to visualise exactly what’s happening, then do so through the
run functions that are called. The
pop functions look like
h :: t => (h, t), where the input stack is decomposed into head (new value) and tail (new state). After the second invocation of this the “value” is the second item on the input stack.
These two values are used later by the
push commands, which act as if they have side-effects and so “return”
Unit. They don’t have side-effects in any traditional sense (everything is immutable, right?), and they return a value only in the for-comprehension pull-a-value-out-of-the-right sense. This value is available later as an argument.
The last function,
get turns the internal state into the value of the computation, so that it can be accessed after running the for-comprehension.
What is the type of the whole for-comprehension? The non-fixed type parameter is the same as the type of the
yielded thing, so we get
State[S, S]. When we finally run the function that we’ve built up,
f.run, we get a pair containing the value and the current internal state, so we pick off the first value of that as our result
List(2, 1, 3, 4, 5)
“What is passed to
get, if we’re ignoring the previous output?”
Back to the stack example. Aren’t we discarding the output of the
get? No – don’t think about it too imperatively. The for-comprehension all boils down to maps and flatMaps, and while it’s easier to visualise the sequence of events in the for-comprehension, it’s important to understand how the computation state is threaded through the sequence of function compositions.
If you de-sugar the for-comprehension and inline
flatMap you get the
f2 function here. This makes it explicit how the state is threaded through. The only parts of this you write yourself in the for-comprehension are the
get calls and (sort of) the reference to
It looks a lot less like magic when written out this way.
The stateless operations on
There are a few standard operations we can use on
State, which modify the internal state. Each of these is itself a
State wrapping some function that composes with the rest of your computation.
getmakes the internal state accessible – by turning it into a
State[S, S]the internal state is the argument to the next map/flatMap call.
setreplaces all existing state with a new
S. The resulting
Statehas a second type parameter of
Unitbecause you don’t have a “value” any more.
mapbut acts on the state instead of the values.
For completeness, here are the implementations of map and flatMap for