A Stricter Control.Monad.Trans.Writer.Strict

Question

So we have:

import Control.Monad.Writer.Strict

type M a = Writer (Map Key Val) a

for some Key and Val.

Everything works okay as long as we don't look at the collected outputs:

report comp = do
  let (a,w) = runWriter comp
  putStrLn a

However if we want to examine w, we get stack overflows.

 report comp = do
   let (a,w) = runWriter comp
   guard (not $ null w) $ do -- forcing w causes a stack overflow
     reportOutputs w
   putStrLn a

The reason, I think, is because (>>=) for Writer is defined as:

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

If I have a large Writer a computation, it builds up a long sequence of mappends: w <> (w' <> (w'' <> ...)) and in this case that's a Map.union which is strict in the spine of the map. So if I build up a big sequence of unions, the whole thing has to be evaluated as soon as I force the Map which overflows the stack.

What we want is to perform the unions early. We want a stricter Strict.Writer:

m >>= k = WriterT $ do
    (a, w) <- runWriterT m
    (b, w') <- runWriterT (k a)
    let w'' = w `mappend` w'
    w'' `seq` return (b, w'')

So my question is: does this exist in some "standard" library? If not, why not?

Answer 1

The immediate answer to your question is: no, there is no standard library that provides this. Also, the version you proposed will still leak. The only version I know that does not leak is to simulate WriterT using a strict StateT. I wrote up a very detailed e-mail about this to the Haskell libraries mailing list comparing the strictness and performance of several implementations. Long story short: implementing WriterT as a strict StateT not only eliminates space leaks, but also generates very efficient code.

Here is the implementation that worked:

newtype WriterT w m a = WriterT { unWriterT :: w -> m (a, w) }

instance (Monad m, Monoid w) => Monad (WriterT w m) where
    return a = WriterT $ \w -> return (a, w)
    m >>= f  = WriterT $ \w -> do
        (a, w') <- unWriterT m w
        unWriterT (f a) w'

runWriterT :: (Monoid w) => WriterT w m a -> m (a, w)
runWriterT m = unWriterT m mempty

tell :: (Monad m, Monoid w) => w -> WriterT w m ()
tell w = WriterT $ \w' ->
    let wt = w `mappend` w'
     in wt `seq` return ((), w `mappend` w')

I would like to see this added to transformers at some point, but there are some minor issues that need to be resolved (i.e. what the module name should be, for one).