The saga of my quest to understand the State monad

Look carefully at the right hand side of the equation they give for p >>= k. Right at the start: \s0 -> .... So already it can't be a pair, but must be some type of function. We're hoping that it's a function of type s -> (s,b), so let's say that

s0 :: s

We already know from the type signature of (>>=) that in this case,

p :: StateTrans s a
   = s -> (s,a) -- expanding the definition of StateTrans
k :: (a -> StateTrans s b)
   = a -> s -> (s,b) -- expanding

So, we know that the return value has type s -> u
where u is the type of:

let (s1,a) = p s0
    q = k a
in q s1

Let's unravel that piece by piece, I'll take the smallest steps possible, and document the reasons for things. (The things I didn't give a reason for are things we already know from above.)

p :: s -> (s,a)          -- 1
s0 :: s                  -- 2
p s0 :: (s,a)            -- 3, by 1 and 2
(s1,a) :: (s,a)          -- 4, by 3 and the binding (s1,a) = p s0
s1 :: s                  -- 5, by 4
a :: a                   -- 6, by 4
k :: a -> s -> (s,b)     -- 7
k a :: s -> (s,b)        -- 8, by 6 and 7
q :: s -> (s,b)          -- 9, by 8 and the binding q = k a
q s1 :: (s,b)            -- 10, by 5 and 9.
let ... in q s1 :: (s,b) -- 11, by 10

So, we've determined that our type u = (s,b), so the full type of the right hand side is s -> (s,b) = StateTrans s b. So the thing checks out.

Formal proofs aside, intuitively, this should also make sense. What is m >>= k really doing? Well, it's taking a state computation m, and a function from possible results of that state computation, to other state computations, k (you might even call it a continuation). It's forming a new state computation which, when applied to an initial state of type s, is going to perform the first state computation, getting a result of type a and a *new* state value of type s. It then applies that result of type a to the function k, which gives us back a state computation to continue with. It then applies the new state value to that continuation, getting a final state of type s and a final result of type b.

But none of that happens immediately -- it's all inside the function, and none of it can happen until the actual initial state is passed in.

Now, that's pretty wordy, but it's exactly what the code says it does if you look carefully. However, the following code, which is equivalent, might be easier to read:

newtype State s a = State (s -> (s,a))

runState :: State s a -> s -> (s,a)
runState (State f) s = f s

instance Monad (State s) where
    return x = State $ \s -> (s,x)
    m >>= k = State $ \s -> -- given an initial state
                let (t, a) = runState m s -- we run the computation m,
                                          -- getting a new state and
                                          -- a value
                in runState (k a) t -- then run the resulting future
                                    -- given to us by k, passing in
                                    -- the new state

3. q s1 provides the desired state, causing q to return a 2-tuple

This returned 2-tuple is _not_ a state transition and therefore violates the given type signature.

...but the return value of ">>=" is not the return value of "q". Note the lambda function starting with \s0 ->... Maybe adding some parenthesis will help...

(>>=) :: StateTrans s a -> (a -> StateTrans s b) -> StateTrans s b
p >>= k = \s0 -> (let (s1, a) = p s0
                            q = k a
                  in q s1)

Type-compatibility between the composed functions should have told me how m and onVal were related: m returns (a, s) therefore I only had to look at

1. can m "digest" (a,s)
2. can (onVal m) "digest" (a,s)

Since f can only take a and not
(a,s) we see that it cannot be composed with m.