Stricter Strict State Monad

The strict state monad is defined using:

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

But this can still leak memory, because a and s' are left unevaluated. For example, we might have a function f that takes a large object as input and quickly returns (a, s'), but as long as a is left unevaluated the input to f cannot be GC'ed.

One potential solution is to have f return seq a (a, s'), but this isn't always possible if we are using something like MonadRandom, and the state is encapsulated away from f. Is there a version that is defined like this:

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

Does this exist in a library anywhere already?

Solution

According to the monad identity laws,

return a >>= const b = const b a = b

Thus in particular,

return undefined >>= const b = b

If the >>= operation were strict in the result value, that would break this law, so you shouldn't do that.

Suppose you instead do this:

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

Now we face another identity law:

m >>= return = m

For example,

return a >>= return = return a

So if return a >>= return is strict in the state, then we must also have return a strict in the state! So we need to redefine return as well:

return a = State $ \ !s -> (a, s)

Note that you don't really need to do any of this; if you want, you can use the usual strict state monad, and write things like

!_ <- get

in the spots where you want to force the state. You could even write an action to do this:

forceState :: Monad m => StateT s m ()
forceState = get >>= \ !_ -> return ()

Edit

Even this definition feels a little bit strange to me; I would expect the lambda to force the state, rather than the case. I'm not sure if not doing that leads to some kind of breakage, but it wouldn't surprise me if it did.