Garbage collecting a list while running an IO action over it

Question

I want to write a conjugate gradient solver in Haskell and want to use lazy lists to decouple stopping rule and output of information from the iterations. My code essentially looks like this:

data CGState = CGState { cgx :: Image
                       , cgp :: Image
                       , cgr :: Image
                       , cgr2 :: Double
                       }

cg :: Operator -> Image -> [CGState]
cg = [...]

runCG :: (CGState -> Bool) -> Operator -> Image -> IO Image
runCG stoprule op rhs = do
  let steps = takeWhile (not . stoprule) $ cg op rhs
  fmap last $ forM (zip [(1::Int)..] steps) $ \(n, cgs) -> do
      putStrLn $ show n ++ "  " ++ show (sqrt $ cgr2 cgs)
      return $ cgx cgs

The idea is to iterate over the list, output some information, but only retain the last iterate. However, when running this code, it does not seem to garbage collect the preceding iterates. My guess is that this is connected to IO: if I rewrite the code like

runCG :: (CGState -> Bool) -> Operator -> Image -> IO Image
runCG stoprule op rhs = do
  let steps = takeWhile (not . stoprule) $ cg op rhs
  return $ cgx $ last steps

the problem does not occur, i.e. everything but the final iterate gets garbage collected directly.

How can I achieve the same effect while being able to output some information about the iterates?

Answer 1

Right, I think the problem is with fmap in IO: because IO actions are always executed in strict sequence, the fmap only applies the last after the entire result list from forM has been constructed.

Instead you can probably use Control.Monad.foldM, which folds monadically over a list (untested):

runCG stoprule op rhs = do
  let steps = takeWhile (not . stoprule) $ cg op rhs
  foldM (\ _ (n, cgs) -> do
      putStrLn $ show n ++ "  " ++ show (sqrt $ cgr2 cgs)
      return $ cgx cgs)
    undefined -- initial value for fold is ignored as long as steps is nonempty
    (zip [(1::Int)..] steps)