Why does this simple use of the State monad cause a stack overflow?

Question

I was playing around with the State monad, and I don't know what's causing the stack overflow in this simple piece of code.

import Control.Monad.State.Lazy

tick :: State Int Int
tick = do n <- get
         put $! (n+1)
         return n

million :: Int
million = snd $ runState (mapM_ (const tick) [1..1000000]) 0

main = print million

Note I would just like to know what's causing the problem in this piece of code, the task itself is not important per se.

Answer 1

The problem is that Control.Monad.State.Lazy's (>>=) is so lazy that even the ($!) doesn't help you.
Try Control.Monad.State.Strict, that should reach the ($!).

The (>>=) of the lazy state monad doesn't look at all at the (value,state) pair, so the only way to get some evaluation done before the end is reached is having the f in m >>= f deconstruct the pair. That doesn't happen here, so you get a huge thunk, that is too large for the stack when runState finally wants a result.

Okay, I've eaten, now I can elaborate. Let me use the old (mtl-1.x) definition of the lazy State s monad, it's a bit easier to see there without the inner monad. The new (mtl-2.x) definition type State s = StateT s Identity behaves the same, it's just more writing and reading. The definition of (>>=) was

m >>= k  = State $ \s -> let
    (a, s') = runState m s
    in runState (k a) s'

Now, let bindings are lazy, hence this is

m >>= k = State $ \s -> let
    blob = runState m s
    in runState (k $ fst blob) (snd blob)

only more readable. So the (>>=) lets the blob completely unevaluated. Evaluation is only required if k needs to inspect fst blob to determine how to continue, or k a needs to inspect snd blob.

In replicateM r tick, the computations are chained with (>>), so the k in (>>=)'s definition is const tick. As a constant function, it most definitely doesn't need to inspect its argument. So tick >> tick becomes

State $ \s -> 
    let blob1 = (\n -> let n' = n+1 in seq n' ((),n')) s
        blob2 = (\m -> let m' = m+1 in seq m' ((),m')) (snd blob1)
    in blob2

the seq isn't touched until blobN has to be evaluated. But needing to evaluate it to the outermost constructor - the pair constructor (,) - would be enough to trigger the seq and that would in turn lead to complete evaluation here. Now, in million, nothing requires any evaluation until the final snd after the runState is reached. By that time, a thunk with a million layers has been built. Evaluating that thunk requires pushing many let m' = m+1 in seq m' ((),m') on the stack until the initial state is reached, and if the stack is large enough to hold them, they'd then be popped and applied. So it'd be three traversals, 1. building the thunk, 2. peeling layers from the thunk and pushing them on the stack, 3. consuming the stack.

The (>>=) of Control.Monad.State.Strict is just strict enough to force the seq s on each bind, thus there's only one traversal, no (nontrivial) thunk is built and the computation runs in constant space. The definition is

m >>= k = State $ \s ->
    case runState m s of
      (a, s') -> runState (k a) s'

The important difference is that pattern matching in case expressions is strict, here the blob has to be evaluated to the outermost constructor to match it against the pattern in the case.
With m = tick = State (\m -> let m' = m+1 in seq m' ((),m')) the essential part becomes

case let s' = s+1 in seq s' ((),s') of
    (a, s'') -> runState (k a) s''

The pattern-match demands the evaluation of ((), s') [to the (,) constructor], by the seq that is tied to the evaluation of s' = s+1, everything is fully evaluated on each bind, no thunks, no stack.

However, you still need to be careful. In this case, due to the seq (resp. ($!)) and the shallow structure of the involved types,evaluation kept up with application of (>>). In general, with deeper structured types and/or without seqs, C.M.S.Strict also builds large thunks which can lead to stack overflow. The thunks are just a bit simpler and less entangled than those generated by C.M.S.Lazy in those circumstances.

On the other hand, the laziness of C.M.S.Lazy allows other computations that are impossible with C.M.S.Strict. For example, C.M.S.Lazy provides one of the very few monads where

take 100 <$> mapM_ something [1 .. ]

terminates. [But be aware that the state is then unusable; before it could be used, it would have to travel through the entire infinite list. So, if you do something like that, before you can resume state-dependent computations, you have to put a fresh state.]