Infinite recursion in IO wrapped with StateT causing stack space overflow

Question

I am developing a graphical (game-like) program (using SDL for graphics) in Haskell. As part of this, there is necessarily a 'main' infinite loop, which handles the state updates and drawing, keeping track of time in the process. I previously had it working where it would take the state as a parameter like the following simplified example:

loop :: State -> Int -> IO ()
loop state prevTime = do
  time <- getTicks
  let state' = updateState state (time - prevTime)
  loop state' time

In this, the updateState function itself used a monad transformer stack including a StateT for state, ReaderT for config, a WriterT for logging, and a RandT for random number generation, so I had to pass all of the parameters required through to updateState, and take the random number generator back. And there were multiple different updateState-related functions, all of which needed some IO done in between them, and so a lot of the code in my main loop was just threading state and parameters through.

Of course, I quickly realised that that kind of code was horrible and decided I could probably do away with it by making the whole loop run in a different transformer stack, this one with IO at the bottom. I can then run code in my pure monad and have it automatically thread its state through with the following function:

promoteContext :: PureContext a -> IOContext a
promoteContext ctx = do
  state <- get
  config <- ask
  gen <- liftIO getStdGen
  let (result, state', log, gen') = runContext ctx gen rules state
  liftIO $ setStdGen gen'
  put state'
  tell log
  return result

This cleaned up my code massively, and it seemed like a much better solution at first, but I soon encountered a massive issue: running my program for around a minute or so suddenly started causing a stack space overflow.

I managed to somewhat work around this issue while keeping my code cleaner than it was before, by making the main loop method run in the IO monad directly and thread the state through, meaning I only have to deal with that boilerplate once per loop. This clearly forces the state to be fully evaluated every time it loops, meaning the stack doesn't fill up, but it still feels 'dirty'.

My question is this: is there any way I can force a full evaluation of the state every time I loop, without resorting to falling back to IO and explicitly leaving/re-entering the transformer stack?

EDIT: Thanks to suggestions by Petr Pudlák, I spent some time trying to isolate the issue, and eventually arrived at this code which causes the problem:

import Control.Monad.Writer.Strict
import Control.Monad.State.Strict

import Control.DeepSeq
import Exception

type ContextT s m = StateT s (WriterT [String] m)

evalContext ctx state = do
  (a, log) <- runWriterT (evalStateT ctx state)
  liftIO $ evaluate (rnf log)
  return a

problem :: ContextT Double IO ()
problem = do modify (+ 0.001)
             s <- get
             liftIO $ print s
             problem

main :: IO ()
main = evalContext problem 0

{A small edit to the code: I added the forced evaluation of the log and the problem still happens.}

Running this causes the stack space to overflow by the time the state reaches about 500. Getting rid of the WriterT, however, stops the overflow from happening, suggesting that it was the Writer's fault all along. I don't understand how this can be, though, as in this isolated code the writer isn't even used. I guess my question is now why on Earth does the presence of a WriterT cause this to happen?

Answer 1

Update: I wasn't able to reproduce the issue, different GHC versions might optimize code differently. Nevertheless some ideas below:

The forced evaluation you added doesn't do anything - it's never executed, as problem runs indefinitely.

The reason why adding WriterT to the stack causes the memory leak could be precisely because it's never used for anything. This is a common problem with Writers. With State you can force the current state at any point, but not with a Writer.

What I'd do is to create a separate primitive for looping forever in ContextT that ensures that after each loop both the state and the log are fully evaluated:

 import Control.Monad
 import Control.Monad.IO.Class
 import Control.Monad.RWS.Strict

 import Control.Exception (evaluate)
 import Control.DeepSeq

 import qualified Data.Sequence as Seq

 type ContextT s m = RWST () (Seq.Seq String) s m

 evaluate' :: (MonadIO m, NFData a) => a -> m a
 evaluate' = liftIO . evaluate . force

 forever' :: (MonadIO m, NFData s) => ContextT s m a -> ContextT s m b
 forever' k = RWST $ \_ -> loop mempty
   where
     loop w s = do
         (_, s', w') <- runRWST k () s
         let w'' = w <> w'
         evaluate' s'
         evaluate' w''
         loop w'' s'

 evalContext ctx state = do
   (a, _, _) <- runRWST ctx () state
   return a

 problem :: ContextT Double IO ()
 problem = do modify (+ 0.001)
              s <- get
              liftIO $ print s

 main :: IO ()
 main = evalContext (forever' problem) 0

Some notes and caveats:

• Separating the one-step computation from looping gives you more control how to ensure strictness or other properties of the loop.
• Using [String] as the monoid inside a Writer will most likely have O(n^2) performance as you're appending small lists from the right to the growing list of current entries. Therefore I'd use Seq to get it down to O(n log n).
• However, forcing the log at every loop will still be O(n^2) - we'll traverse the full log each time. Solving this is probably out of the scope of this answer, I'd probably implement my own monoid with a relaxed tree structure for this.
• I'd suggest to use newtype for ContextT and not to export its internals outside the module. Even though this requires little boiler-plate code to add all the required instances, I see two advantages:
  • Encapsulation - this forces every other part of your code to use only supplied primitives, which allows you to change the internals later (like changing StateT + WriterT to RWST) without touching anything outside the module.
  • You have more control of the monadic operations - for example by defining a customized >>= you can enforce that the state is evaluated to its NF after each monadic operation.

---

Old answer: If possible, please include a self-contained code sample that demonstrates the issue, otherwise it's hard to just guess the cause.

Some general suggestions I'd try out:

• The leak might be also caused by the writer part, so try out tell $! log or tell $!! log as well.
• If you're in the IO monad, a canonical way how to force evaluation is evaluate. If you want a normal form then you'd use evaluate (rnf value).
• If your stack uses StateT, ReaderT as well as WriterT, you might use `RWST instead.

• Your promoteContext looks very close to hoist. To be more specific: You could create a monad transformer parametrized by the base monad:

  newtype ContextT m a = ...
  

  And also implement MFunctor. Then the pure computations would be run inside ContextT Identity, the IO-based would be run in ContextT IO and you can hoist from ContextT Identity to ContextT IO using hoist generalize.

• Another approach could be to run the pure parts in ContextT m for an arbitrary m (or constrained by a type class such as MonadRandom m => ...). If m is arbitrary then it's possible to substitute Identity for m, so the code is indeed pure, but you can combine it together with the IO based parts seamlessly (also IO implements MonadRandom).

• I'd try to separate the part that forces evaluation of all the internal state and writer output into a separate function, something as

  forceContext :: ContextT IO () -- or MonadIO m => ContextT m ()
  forceContext = do
    -- state
    s <- -- read the internal state
    liftIO $ evaluate (rnf s)
    ...
  

  Then you could express the main loop as something similar to

  main = runContextT (forever (step >> forceContext))
  

Answer 2

The problem is that Writer doesn't have proper tail calls; its >>= is something like (ignoring minor details like newtype constructors / extractors):

a >>= f = do
    (w0, x) <- a
    (w1, y) <- f x
    return (w0 `mappend` w1, y)

So f isn't really in tail context here; as such, infinite tail-recursive loops like you're using will accumulate stack space on every loop.

Writer is a problem even though you 'aren't using it' because do is always de-sugared using the >>= of whatever type your computation is in, and it's the >>= you're picking up that's the problem here.