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.
Each line where we call a function "in the state monad" behaves as follows: the arguments we give it on that line are applied, yielding a partially applied function. This is because each of the state monad functions yields a State function (which still needs an argument of type s before it can yields it's result tuple.
The IO monad in Haskell is often explained as a state monad where the state is the world. So a value of type IO a monad is viewed as something like worldState -> (a, worldState) .
The Haskell type State describes functions that consume a state and produce both a result and an updated state, which are given back in a tuple. Here, s is the type of the state, and a the type of the produced result.
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 seq
s, 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.]
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With