Make the old codegen run in constant space too

2 files changed, 41 insertions, 29 deletions

diff --git a/compiler/codeGen/CgMonad.lhs b/compiler/codeGen/CgMonad.lhs
index 302d8ac652..59f6accf9d 100644
--- a/compiler/codeGen/CgMonad.lhs
+++ b/compiler/codeGen/CgMonad.lhs
@@ -14,7 +14,7 @@ module CgMonad (
         Code,
         FCode,
 
-        initC, thenC, thenFC, listCs, listFCs, mapCs, mapFCs,
+        initC, runC, thenC, thenFC, listCs, listFCs, mapCs, mapFCs,
         returnFC, fixC, fixC_, checkedAbsC, 
         stmtC, stmtsC, labelC, emitStmts, nopC, whenC, newLabelC,
         newUnique, newUniqSupply, 
@@ -379,13 +379,12 @@ instance Monad FCode where
 The Abstract~C is not in the environment so as to improve strictness.
 
 \begin{code}
-initC :: DynFlags -> Module -> FCode a -> IO a
+initC :: IO CgState
+initC  = do { uniqs <- mkSplitUniqSupply 'c'
+            ; return (initCgState uniqs) }
 
-initC dflags mod (FCode code)
-  = do  { uniqs <- mkSplitUniqSupply 'c'
-        ; case code (initCgInfoDown dflags mod) (initCgState uniqs) of
-              (res, _) -> return res
-        }
+runC :: DynFlags -> Module -> CgState -> FCode a -> (a,CgState)
+runC dflags mod st (FCode code) = code (initCgInfoDown dflags mod) st
 
 returnFC :: a -> FCode a
 returnFC val = FCode (\_ state -> (val, state))
diff --git a/compiler/codeGen/CodeGen.lhs b/compiler/codeGen/CodeGen.lhs
index aa561c4f40..f8898450ef 100644
--- a/compiler/codeGen/CodeGen.lhs
+++ b/compiler/codeGen/CodeGen.lhs
@@ -45,6 +45,13 @@ import TyCon
 import Module
 import ErrUtils
 import Panic
+import Outputable
+
+import OrdList
+import Stream (Stream, liftIO)
+import qualified Stream
+
+import Data.IORef
 
 codeGen :: DynFlags
         -> Module                     -- Module we are compiling
@@ -52,32 +59,38 @@ codeGen :: DynFlags
         -> CollectedCCs               -- (Local/global) cost-centres needing declaring/registering.
         -> [(StgBinding,[(Id,[Id])])] -- Bindings to convert, with SRTs
         -> HpcInfo                    -- Profiling info
-        -> IO [CmmGroup]
+        -> Stream IO CmmGroup ()
               -- N.B. returning '[Cmm]' and not 'Cmm' here makes it
               -- possible for object splitting to split up the
               -- pieces later.
 
-codeGen dflags this_mod data_tycons cost_centre_info stg_binds hpc_info = do
-    showPass dflags "CodeGen"
-    code_stuff <-
-        initC dflags this_mod $ do
-            cmm_binds  <- mapM (getCmm . cgTopBinding dflags) stg_binds
-            cmm_tycons <- mapM cgTyCon data_tycons
-            cmm_init   <- getCmm (mkModuleInit dflags cost_centre_info this_mod hpc_info)
-            return (cmm_init : cmm_binds ++ cmm_tycons)
-                -- Put datatype_stuff after code_stuff, because the
-                -- datatype closure table (for enumeration types) to
-                -- (say) PrelBase_True_closure, which is defined in
-                -- code_stuff
-
-                -- Note [codegen-split-init] the cmm_init block must
-                -- come FIRST.  This is because when -split-objs is on
-                -- we need to combine this block with its
-                -- initialisation routines; see Note
-                -- [pipeline-split-init].
-
-    dumpIfSet_dyn dflags Opt_D_dump_cmm "Cmm" (pprCmms (targetPlatform dflags) code_stuff)
-    return code_stuff
+codeGen dflags this_mod data_tycons cost_centre_info stg_binds hpc_info
+
+   = do { liftIO $ showPass dflags "CodeGen"
+
+        ; cgref <- liftIO $ newIORef =<< initC
+        ; let cg :: FCode CmmGroup -> Stream IO CmmGroup ()
+              cg fcode = do
+                cmm <- liftIO $ do
+                         st <- readIORef cgref
+                         let (a,st') = runC dflags this_mod st fcode
+
+                         dumpIfSet_dyn dflags Opt_D_dump_cmm "Cmm" $
+                             pprPlatform (targetPlatform dflags) a
+
+                         -- NB. stub-out cgs_tops and cgs_stmts.  This fixes
+                         -- a big space leak.  DO NOT REMOVE!
+                         writeIORef cgref $! st'{ cgs_tops = nilOL,
+                                                  cgs_stmts = nilOL }
+                         return a
+                Stream.yield cmm
+
+        ; cg (getCmm $ mkModuleInit dflags cost_centre_info this_mod hpc_info)
+
+        ; mapM_ (cg . getCmm . cgTopBinding dflags) stg_binds
+
+        ; mapM_ (cg . cgTyCon) data_tycons
+        }
 
 mkModuleInit
         :: DynFlags