Make the old codegen run in constant space too

compiler/cmm/CmmParse.y

@@ -1074,7 +1074,8 @@ parseCmmFile dflags filename = do
         let msg = mkPlainErrMsg span err
         return ((emptyBag, unitBag msg), Nothing)
     POk pst code -> do
-        cmm <- initC dflags no_module (getCmm (unEC code initEnv [] >> return ()))
+        st <- initC
+        let (cmm,_) = runC dflags no_module st (getCmm (unEC code initEnv [] >> return ()))
         let ms = getMessages pst
         if (errorsFound dflags ms)
          then return (ms, Nothing)

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))

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

compiler/main/HscMain.hs

@@ -1214,9 +1214,9 @@ hscGenHardCode cgguts mod_summary = do
                              cost_centre_info
                              stg_binds hpc_info
                     else {-# SCC "CodeGen" #-}
-                         codeGen dflags this_mod data_tycons
+                         return (codeGen dflags this_mod data_tycons
                                cost_centre_info
-                               stg_binds hpc_info >>= return . Stream.fromList
+                               stg_binds hpc_info)
 
 
         ------------------  Code output -----------------------