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

Question

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.

Answer 1

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:

Answer 2

Look at the source for the strict writer monad: http://hackage.haskell.org/packages/archive/transformers/0.2.2.0/doc/html/src/Control-Monad-Trans-Writer-Strict.html#line-122

The difference from the lazy writer is that in the latter, the pattern matches are lazy -- but in neither case is the mappend operation or the "state" of the writer thus far forced. To solve your issue, you'd need a "super-strict" writer.

Answer 3

I think your understanding is correct.

I'm interested in how these functions:

bar :: Integer -> (Integer, Integer)
bar x = bar' 0 x
    where bar' c 0 = (0, c)
          bar' c x = c `seq` bar' (c + x) (pred x)
      --  bar' c x = let c' = c+x in c' `seq` bar' c' (pred x)
      --  bar' !c !x = bar' (c+x) (pred x)

produces a stack overflow when compiled with optimizations, although the related functions:

bar2 :: Integer -> (Integer, Integer)
bar2 x = bar' 0 x
     where bar' c 0 = (0, c)
           bar' !c !x = let c' = c+x in c' `seq` bar' c' (pred x)

bar3 :: Integer -> Integer
bar3 x = bar' 0 x
     where bar' c 0 = c
           bar' c x = c `seq` bar' (c + x) (pred x)

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

do not.

I think this looks like a bug in GHC's optimizer, and you should report it. From looking at the core (produced with -ddump-simpl), the c argument isn't strictly evaluated in all cases for the overflowing functions. bar2 is the closest working version I found to the original, and it seems over-specified to me.

Since you have a very simple test case, there's a good chance it'll get fixed.