Search code examples
haskellspace-leakwriter-monad

Space leaks, and Writers, and Sums (oh my!)


I've been playing with the Writer Monad recently, and I've run into what appears to be a space leak. I can't say I fully understand these things yet, so I'd like to know what's happening here, and how to fix it.

First, here's how I can trigger this error:

import Control.Monad.Writer
import Data.Monoid

foo :: Integer -> Writer (Sum Integer) Integer
foo 0 = return 0
foo x = tell (Sum x) >> foo (pred x)

main = print $ runWriter $ foo 1000000

I get:

Stack space overflow: current size 8388608 bytes.
Use `+RTS -Ksize -RTS' to increase it.

To understand this better, I've reimplemented similar functionality without Writer or Sum, and if I keep things nice and lazy, I get the same error:

bar :: Integer -> (Integer, Integer)
bar x = bar' 0 x
    where bar' c 0 = (0, c)
          bar' c x = bar' (c + x) (pred x)

But I can remedy this by adding seq to the equation:

bar' c x = c `seq` bar' (c + x) (pred x)

I've tried seqing various bits of my foo function, but that doesn't seem to help. Also, I've tried using Control.Monad.Writer.Strict but that doesn't make a difference either.

Does Sum need to be strict somehow? Or am I missing something completely different?

Notes

  • I may have my terminology wrong here. According to Space leak zoo, my problem would be classified as a 'stack overflow', and if that's the case, how would I convert foo to a more iterative style? Is my manual recursion the problem?
  • After reading Haskell Space Overflow, I had the idea to compile with -O2, just to see what happens. This may be a topic for another question, but with optimizations, even my seq'd bar function fails to run. Update: This issue goes away if I add -fno-full-laziness.

Solution

  • The problem with Writer monad is that it's >>= is not tail-recursive:

    instance (Monoid w, Monad m) => Monad (WriterT w m) where
    m >>= k  = WriterT $ do
        (a, w)  <- runWriterT m
        (b, w') <- runWriterT (k a)
        return (b, w `mappend` w')
    

    As you can see it has to evaluate both m and k a to evaluate mappend which means the entire stack of recursive calls has to be forced before the first mappend can be evaluated. I believe that regardless of the strictness Writer monad will cause stack overflow in your definition (or can it be avoided with lazy version somehow?).

    If you still want to use monads you can try State which is tail-recursive. Either strict version of it with strict put:

    import Control.Monad.State.Strict
    
    foo :: Integer -> State (Sum Integer) Integer
    foo 0 = return 0
    foo x = do
       w <- get
       put $! w `mappend` (Sum x)
       foo (pred x)
    
    main = print $ (`execState` Sum 0) $ foo 1000000
    

    Or lazy version with continuation passing style (CPS):

    import Control.Monad.Cont
    import Control.Monad.State.Lazy
    
    foo :: Integer -> ContT Integer (State (Sum Integer)) Integer
    foo 0 = return 0
    foo x = do
      w <- get
      put $! w `mappend` (Sum x)
      foo (pred x)
    
    main = print $ ((`execState`Sum 0) . (`runContT`return)) $ foo 1000000
    

    Handy analog for tell:

    stell :: (MonadState w m, Monoid w) => w -> m ()
    stell a = get >>= \w -> put $! w `mappend` a
    

    I suspect that if it was possible to use ContT with Writer CPS would help us with Writer as well, but looks like it isn't possible to define ContT for MonadWriter: