Space leak in Pipes with RWST

Question

A memory analysis of the following program shows that the noleak functions runs in constant memory while the leak function leaks memory in a linear fashion. dflemstr indicated that this might be due to RWST causing an infinite chain of allocations. Is this the case and what other solutions exists? I actually dont need the Writer monad.

Environment:

GHC 7.8.3 on ARCH 64 bit

ghc Pipe.hs -o Pipe -prof

import Control.Concurrent (threadDelay)
import Control.Monad (forever)

import Pipes
import Control.Monad.Trans.RWS.Strict

main = leak

effectLeak :: Effect (RWST () () () IO) ()
effectLeak =
  (forever $ do
      liftIO . threadDelay $ 10000 * 1
      yield "Space") >->
  (forever $ do
      text <- await
      yield $ text ++ (" leak" :: String)) >->
  (forever $ do
      text <- await
      liftIO . print $ text
  )

effectNoleak :: Effect IO ()
effectNoleak =
  (forever $ do
      lift . threadDelay $ 10000 * 1
      yield "Space") >->
  (forever $ do
      text <- await
      yield $ text ++ (" leak" :: String)) >->
  (forever $ do
      text <- await
      lift . print $ text
  )

leak = (\e -> runRWST e () ()) . runEffect $ effectLeak

noleak = runEffect $ effectNoleak

Answer 1

Zeta is right, and the space leak is because of WriterT. WriterT and RWST (both the "strict" and lazy versions) always leak space no matter what monoid you use.

I wrote up a longer explanation of this here, but here is the summary: the only way to not leak space is to simulate WriterT using a StateT monad where tell is simulated using a strict put, like this:

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 ((), wt)

That's basically equivalent to:

type WriterT = StateT

runWriterT m = runStateT m mempty

tell w = do
    w' <- get
    put $! mappend w w'

Answer 2

It seems like the Writer part of RWST is actually the culprit:

instance (Monoid w, Monad m) => Monad (RWST r w s m) where
    return a = RWST $ \ _ s -> return (a, s, mempty)
    m >>= k  = RWST $ \ r s -> do
        (a, s', w)  <- runRWST m r s
        (b, s'',w') <- runRWST (k a) r s'
        return (b, s'', w `mappend` w') -- mappend
    fail msg = RWST $ \ _ _ -> fail msg

As you can see, the writer uses a plain mappend. Since (,,) isn't strict in its arguments, w `mappend` w' builds a series of thunks, even tough the Monoid instance of () is rather trivial:

instance Monoid () where
        -- Should it be strict?
        mempty        = ()
        _ `mappend` _ = ()
        mconcat _     = ()

In order to fix this, you need to add strictness to w `mappend` w' in the tuple:

        let wt = w `mappend` w'
        wt `seq` return (b, s'', wt) 

However, if you don't need the Writer anyway, you can simply use ReaderT r (StateT st m) instead:

import Control.Monad.Trans.Reader
import Control.Monad.Trans.State.Strict

type RST r st m = ReaderT r (StateT st m)

runRST :: Monad m => RST r st m a -> r -> st -> m (a,st)
runRST rst r st = flip runStateT st . flip runReaderT r $ rst

However, given that this will force you to lift the computations to the correct monad, you might want to use the mtl package instead. The code will stay the same, but the imports will be the following in this case

import Control.Monad.Reader
import Control.Monad.State.Strict