Collect CCs in CorePrep, including CCs in unfoldings

This patch includes two changes:

1. Move cost centre collection from `SCCfinal` to `CorePrep`, to be able
   to collect cost centres in unfoldings. `CorePrep` drops unfoldings, so
   that's the latest stage in the compilation pipeline for this.

   After this change `SCCfinal` no longer collects all cost centres, but
   it still generates & collects CAF cost centres + updates cost centre
   stacks of `StgRhsClosure` and `StgRhsCon`s.

   This fixes #5889.

2. Initialize cost centre stack fields of `StgRhs` in `coreToStg`. With
   this we no longer need to update cost centre stack fields in
   `SCCfinal`, so that module is removed.

   Cost centre initialization explained in Note [Cost-centre
   initialization plan].

   Because with -fcaf-all we need to attach a new cost-centre to each
   CAF, `coreTopBindToStg` now returns `CollectedCCs`.

Test Plan: validate

Reviewers: simonpj, bgamari, simonmar

Reviewed By: simonpj, bgamari

Subscribers: rwbarton, thomie, carter

GHC Trac Issues: #5889

Differential Revision: https://phabricator.haskell.org/D4325

10 files changed, 258 insertions, 429 deletions

diff --git a/compiler/coreSyn/CorePrep.hs b/compiler/coreSyn/CorePrep.hs
index 2bfb558526..75301791b4 100644
--- a/compiler/coreSyn/CorePrep.hs
+++ b/compiler/coreSyn/CorePrep.hs
@@ -60,12 +60,14 @@ import Name             ( NamedThing(..), nameSrcSpan )
 import SrcLoc           ( SrcSpan(..), realSrcLocSpan, mkRealSrcLoc )
 import Data.Bits
 import MonadUtils       ( mapAccumLM )
-import Data.List        ( mapAccumL )
+import Data.List        ( mapAccumL, foldl' )
 import Control.Monad
+import CostCentre       ( CostCentre, ccFromThisModule )
+import qualified Data.Set as S
 
 {-
 -- ---------------------------------------------------------------------------
--- Overview
+-- Note [CorePrep Overview]
 -- ---------------------------------------------------------------------------
 
 The goal of this pass is to prepare for code generation.
@@ -124,6 +126,10 @@ The goal of this pass is to prepare for code generation.
     (non-type) applications where we can, and make sure that we
     annotate according to scoping rules when floating.
 
+12. Collect cost centres (including cost centres in unfoldings) if we're in
+    profiling mode. We have to do this here beucase we won't have unfoldings
+    after this pass (see `zapUnfolding` and Note [Drop unfoldings and rules].
+
 This is all done modulo type applications and abstractions, so that
 when type erasure is done for conversion to STG, we don't end up with
 any trivial or useless bindings.
@@ -169,7 +175,7 @@ type CpeRhs  = CoreExpr    -- Non-terminal 'rhs'
 -}
 
 corePrepPgm :: HscEnv -> Module -> ModLocation -> CoreProgram -> [TyCon]
-            -> IO CoreProgram
+            -> IO (CoreProgram, S.Set CostCentre)
 corePrepPgm hsc_env this_mod mod_loc binds data_tycons =
     withTiming (pure dflags)
                (text "CorePrep"<+>brackets (ppr this_mod))
@@ -177,7 +183,13 @@ corePrepPgm hsc_env this_mod mod_loc binds data_tycons =
     us <- mkSplitUniqSupply 's'
     initialCorePrepEnv <- mkInitialCorePrepEnv dflags hsc_env
 
-    let implicit_binds = mkDataConWorkers dflags mod_loc data_tycons
+    let cost_centres
+          | WayProf `elem` ways dflags
+          = collectCostCentres this_mod binds
+          | otherwise
+          = S.empty
+
+        implicit_binds = mkDataConWorkers dflags mod_loc data_tycons
             -- NB: we must feed mkImplicitBinds through corePrep too
             -- so that they are suitably cloned and eta-expanded
 
@@ -187,7 +199,7 @@ corePrepPgm hsc_env this_mod mod_loc binds data_tycons =
                       return (deFloatTop (floats1 `appendFloats` floats2))
 
     endPassIO hsc_env alwaysQualify CorePrep binds_out []
-    return binds_out
+    return (binds_out, cost_centres)
   where
     dflags = hsc_dflags hsc_env
 
@@ -1683,3 +1695,39 @@ wrapTicks (Floats flag floats0) expr =
                                              (ppr other)
         wrapBind t (NonRec binder rhs) = NonRec binder (mkTick t rhs)
         wrapBind t (Rec pairs)         = Rec (mapSnd (mkTick t) pairs)
+
+------------------------------------------------------------------------------
+-- Collecting cost centres
+-- ---------------------------------------------------------------------------
+
+-- | Collect cost centres defined in the current module, including those in
+-- unfoldings.
+collectCostCentres :: Module -> CoreProgram -> S.Set CostCentre
+collectCostCentres mod_name
+  = foldl' go_bind S.empty
+  where
+    go cs e = case e of
+      Var{} -> cs
+      Lit{} -> cs
+      App e1 e2 -> go (go cs e1) e2
+      Lam _ e -> go cs e
+      Let b e -> go (go_bind cs b) e
+      Case scrt _ _ alts -> go_alts (go cs scrt) alts
+      Cast e _ -> go cs e
+      Tick (ProfNote cc _ _) e ->
+        go (if ccFromThisModule cc mod_name then S.insert cc cs else cs) e
+      Tick _ e -> go cs e
+      Type{} -> cs
+      Coercion{} -> cs
+
+    go_alts = foldl' (\cs (_con, _bndrs, e) -> go cs e)
+
+    go_bind :: S.Set CostCentre -> CoreBind -> S.Set CostCentre
+    go_bind cs (NonRec b e) =
+      go (maybe cs (go cs) (get_unf b)) e
+    go_bind cs (Rec bs) =
+      foldl' (\cs' (b, e) -> go (maybe cs' (go cs') (get_unf b)) e) cs bs
+
+    -- Unfoldings may have cost centres that in the original definion are
+    -- optimized away, see #5889.
+    get_unf = maybeUnfoldingTemplate . realIdUnfolding
diff --git a/compiler/deSugar/Coverage.hs b/compiler/deSugar/Coverage.hs
index 862e564aed..b2e9ea2cf6 100644
--- a/compiler/deSugar/Coverage.hs
+++ b/compiler/deSugar/Coverage.hs
@@ -370,14 +370,7 @@ bindTick density name pos fvs = do
 
 -- Note [inline sccs]
 --
--- It should be reasonable to add ticks to INLINE functions; however
--- currently this tickles a bug later on because the SCCfinal pass
--- does not look inside unfoldings to find CostCentres.  It would be
--- difficult to fix that, because SCCfinal currently works on STG and
--- not Core (and since it also generates CostCentres for CAFs,
--- changing this would be difficult too).
---
--- Another reason not to add ticks to INLINE functions is that this
+-- The reason not to add ticks to INLINE functions is that this is
 -- sometimes handy for avoiding adding a tick to a particular function
 -- (see #6131)
 --
diff --git a/compiler/ghc.cabal.in b/compiler/ghc.cabal.in
index d6d55bf01e..384a50ff7e 100644
--- a/compiler/ghc.cabal.in
+++ b/compiler/ghc.cabal.in
@@ -391,7 +391,6 @@ Library
         TysWiredIn
         CostCentre
         ProfInit
-        SCCfinal
         RnBinds
         RnEnv
         RnExpr
diff --git a/compiler/main/HscMain.hs b/compiler/main/HscMain.hs
index 975c96fbf4..39c2748ad5 100644
--- a/compiler/main/HscMain.hs
+++ b/compiler/main/HscMain.hs
@@ -1309,15 +1309,17 @@ hscGenHardCode hsc_env cgguts mod_summary output_filename = do
         -------------------
         -- PREPARE FOR CODE GENERATION
         -- Do saturation and convert to A-normal form
-        prepd_binds <- {-# SCC "CorePrep" #-}
+        (prepd_binds, local_ccs) <- {-# SCC "CorePrep" #-}
                        corePrepPgm hsc_env this_mod location
                                    core_binds data_tycons
         -----------------  Convert to STG ------------------
-        (stg_binds, cost_centre_info)
+        (stg_binds, (caf_ccs, caf_cc_stacks))
             <- {-# SCC "CoreToStg" #-}
                myCoreToStg dflags this_mod prepd_binds
 
-        let prof_init = profilingInitCode this_mod cost_centre_info
+        let cost_centre_info =
+              (S.toList local_ccs ++ caf_ccs, caf_cc_stacks)
+            prof_init = profilingInitCode this_mod cost_centre_info
             foreign_stubs = foreign_stubs0 `appendStubC` prof_init
 
         ------------------  Code generation ------------------
@@ -1374,7 +1376,7 @@ hscInteractive hsc_env cgguts mod_summary = do
     -------------------
     -- PREPARE FOR CODE GENERATION
     -- Do saturation and convert to A-normal form
-    prepd_binds <- {-# SCC "CorePrep" #-}
+    (prepd_binds, _) <- {-# SCC "CorePrep" #-}
                    corePrepPgm hsc_env this_mod location core_binds data_tycons
     -----------------  Generate byte code ------------------
     comp_bc <- byteCodeGen hsc_env this_mod prepd_binds data_tycons mod_breaks
@@ -1478,15 +1480,15 @@ doCodeGen hsc_env this_mod data_tycons
 
 myCoreToStg :: DynFlags -> Module -> CoreProgram
             -> IO ( [StgTopBinding] -- output program
-                  , CollectedCCs) -- cost centre info (declared and used)
+                  , CollectedCCs )  -- CAF cost centre info (declared and used)
 myCoreToStg dflags this_mod prepd_binds = do
-    let stg_binds
+    let (stg_binds, cost_centre_info)
          = {-# SCC "Core2Stg" #-}
            coreToStg dflags this_mod prepd_binds
 
-    (stg_binds2, cost_centre_info)
+    stg_binds2
         <- {-# SCC "Stg2Stg" #-}
-           stg2stg dflags this_mod stg_binds
+           stg2stg dflags stg_binds
 
     return (stg_binds2, cost_centre_info)
 
@@ -1612,7 +1614,7 @@ hscDeclsWithLocation hsc_env0 str source linenumber =
 
     {- Prepare For Code Generation -}
     -- Do saturation and convert to A-normal form
-    prepd_binds <- {-# SCC "CorePrep" #-}
+    (prepd_binds, _) <- {-# SCC "CorePrep" #-}
       liftIO $ corePrepPgm hsc_env this_mod iNTERACTIVELoc core_binds data_tycons
 
     {- Generate byte code -}
diff --git a/compiler/profiling/CostCentre.hs b/compiler/profiling/CostCentre.hs
index f89654dc00..0043fd4bbc 100644
--- a/compiler/profiling/CostCentre.hs
+++ b/compiler/profiling/CostCentre.hs
@@ -4,9 +4,9 @@ module CostCentre (
                 -- All abstract except to friend: ParseIface.y
 
         CostCentreStack,
-        CollectedCCs,
-        noCCS, currentCCS, dontCareCCS,
-        noCCSAttached, isCurrentCCS,
+        CollectedCCs, emptyCollectedCCs, collectCC,
+        currentCCS, dontCareCCS,
+        isCurrentCCS,
         maybeSingletonCCS,
 
         mkUserCC, mkAutoCC, mkAllCafsCC,
@@ -160,9 +160,7 @@ mkAllCafsCC m loc = AllCafsCC { cc_mod = m, cc_loc = loc }
 --        pre-defined CCSs, see below).
 
 data CostCentreStack
-  = NoCCS
-
-  | CurrentCCS          -- Pinned on a let(rec)-bound
+  = CurrentCCS          -- Pinned on a let(rec)-bound
                         -- thunk/function/constructor, this says that the
                         -- cost centre to be attached to the object, when it
                         -- is allocated, is whatever is in the
@@ -185,20 +183,20 @@ type CollectedCCs
     , [CostCentreStack]  -- pre-defined "singleton" cost centre stacks
     )
 
+emptyCollectedCCs :: CollectedCCs
+emptyCollectedCCs = ([], [])
+
+collectCC :: CostCentre -> CostCentreStack -> CollectedCCs -> CollectedCCs
+collectCC cc ccs (c, cs) = (cc : c, ccs : cs)
 
-noCCS, currentCCS, dontCareCCS :: CostCentreStack
+currentCCS, dontCareCCS :: CostCentreStack
 
-noCCS                   = NoCCS
 currentCCS              = CurrentCCS
 dontCareCCS             = DontCareCCS
 
 -----------------------------------------------------------------------------
 -- Predicates on Cost-Centre Stacks
 
-noCCSAttached :: CostCentreStack -> Bool
-noCCSAttached NoCCS                     = True
-noCCSAttached _                         = False
-
 isCurrentCCS :: CostCentreStack -> Bool
 isCurrentCCS CurrentCCS                 = True
 isCurrentCCS _                          = False
@@ -222,7 +220,6 @@ mkSingletonCCS cc = SingletonCCS cc
 -- expression.
 
 instance Outputable CostCentreStack where
-  ppr NoCCS             = text "NO_CCS"
   ppr CurrentCCS        = text "CCCS"
   ppr DontCareCCS       = text "CCS_DONT_CARE"
   ppr (SingletonCCS cc) = ppr cc <> text "_ccs"
diff --git a/compiler/profiling/SCCfinal.hs b/compiler/profiling/SCCfinal.hs
deleted file mode 100644
index 8a2513fd16..0000000000
--- a/compiler/profiling/SCCfinal.hs
+++ /dev/null
@@ -1,284 +0,0 @@
--- (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
-{-# LANGUAGE CPP #-}
-
------------------------------------------------------------------------------
--- Modify and collect code generation for final STG program
-
-{-
- This is now a sort-of-normal STG-to-STG pass (WDP 94/06), run by stg2stg.
-
-  - Traverses the STG program collecting the cost centres. These are required
-    to declare the cost centres at the start of code generation.
-
-    Note: because of cross-module unfolding, some of these cost centres may be
-    from other modules.
-
-  - Puts on CAF cost-centres if the user has asked for individual CAF
-    cost-centres.
--}
-
-module SCCfinal ( stgMassageForProfiling ) where
-
-#include "HsVersions.h"
-
-import GhcPrelude
-
-import StgSyn
-
-import CostCentre       -- lots of things
-import Id
-import Name
-import Module
-import UniqSupply       ( UniqSupply )
-import Outputable
-import DynFlags
-import CoreSyn          ( Tickish(..) )
-import FastString
-import SrcLoc
-import Util
-
-import Control.Monad (liftM, ap)
-
-stgMassageForProfiling
-        :: DynFlags
-        -> Module                       -- module name
-        -> UniqSupply                   -- unique supply
-        -> [StgTopBinding]              -- input
-        -> (CollectedCCs, [StgTopBinding])
-
-stgMassageForProfiling dflags mod_name _us stg_binds
-  = let
-        ((local_ccs, cc_stacks),
-         stg_binds2)
-          = initMM mod_name (do_top_bindings stg_binds)
-
-        (fixed_ccs, fixed_cc_stacks)
-          = if gopt Opt_AutoSccsOnIndividualCafs dflags
-            then ([],[])  -- don't need "all CAFs" CC
-            else ([all_cafs_cc], [all_cafs_ccs])
-
-        local_ccs_no_dups  = nubSort local_ccs
-    in
-    ((fixed_ccs ++ local_ccs_no_dups,
-      fixed_cc_stacks ++ cc_stacks), stg_binds2)
-  where
-
-    span = mkGeneralSrcSpan (mkFastString "<entire-module>") -- XXX do better
-    all_cafs_cc  = mkAllCafsCC mod_name span
-    all_cafs_ccs = mkSingletonCCS all_cafs_cc
-
-    ----------
-    do_top_bindings :: [StgTopBinding] -> MassageM [StgTopBinding]
-
-    do_top_bindings [] = return []
-
-    do_top_bindings (StgTopLifted (StgNonRec b rhs) : bs) = do
-        rhs' <- do_top_rhs b rhs
-        bs' <- do_top_bindings bs
-        return (StgTopLifted (StgNonRec b rhs') : bs')
-
-    do_top_bindings (StgTopLifted (StgRec pairs) : bs) = do
-        pairs2 <- mapM do_pair pairs
-        bs' <- do_top_bindings bs
-        return (StgTopLifted (StgRec pairs2) : bs')
-      where
-        do_pair (b, rhs) = do
-             rhs2 <- do_top_rhs b rhs
-             return (b, rhs2)
-
-    do_top_bindings (b@StgTopStringLit{} : bs) = do
-        bs' <- do_top_bindings bs
-        return (b : bs')
-
-    ----------
-    do_top_rhs :: Id -> StgRhs -> MassageM StgRhs
-
-    do_top_rhs _ (StgRhsClosure _ _ _ _ []
-                     (StgTick (ProfNote _cc False{-not tick-} _push)
-                              (StgConApp con args _)))
-      | not (isDllConApp dflags mod_name con args)
-        -- Trivial _scc_ around nothing but static data
-        -- Eliminate _scc_ ... and turn into StgRhsCon
-
-        -- isDllConApp checks for LitLit args too
-      = return (StgRhsCon dontCareCCS con args)
-
-    do_top_rhs binder (StgRhsClosure _ bi fv u [] body)
-      = do
-        -- Top level CAF without a cost centre attached
-        -- Attach CAF cc (collect if individual CAF ccs)
-        caf_ccs <- if gopt Opt_AutoSccsOnIndividualCafs dflags
-                   then let cc = mkAutoCC binder modl CafCC
-                            ccs = mkSingletonCCS cc
-                                   -- careful: the binder might be :Main.main,
-                                   -- which doesn't belong to module mod_name.
-                                   -- bug #249, tests prof001, prof002
-                            modl | Just m <- nameModule_maybe (idName binder) = m
-                                 | otherwise = mod_name
-                        in do
-                        collectNewCC  cc
-                        collectCCS ccs
-                        return ccs
-                   else
-                        return all_cafs_ccs
-        body' <- do_expr body
-        return (StgRhsClosure caf_ccs bi fv u [] body')
-
-    do_top_rhs _ (StgRhsClosure _no_ccs bi fv u args body)
-      = do body' <- do_expr body
-           return (StgRhsClosure dontCareCCS bi fv u args body')
-
-    do_top_rhs _ (StgRhsCon _ con args)
-        -- Top-level (static) data is not counted in heap
-        -- profiles; nor do we set CCCS from it; so we
-        -- just slam in dontCareCostCentre
-      = return (StgRhsCon dontCareCCS con args)
-
-    ------
-    do_expr :: StgExpr -> MassageM StgExpr
-
-    do_expr (StgLit l) = return (StgLit l)
-
-    do_expr (StgApp fn args)
-      = return (StgApp fn args)
-
-    do_expr (StgConApp con args ty_args)
-      = return (StgConApp con args ty_args)
-
-    do_expr (StgOpApp con args res_ty)
-      = return (StgOpApp con args res_ty)
-
-    do_expr (StgTick note@(ProfNote cc _ _) expr) = do
-        -- Ha, we found a cost centre!
-        collectCC cc
-        expr' <- do_expr expr
-        return (StgTick note expr')
-
-    do_expr (StgTick ti expr) = do
-        expr' <- do_expr expr
-        return (StgTick ti expr')
-
-    do_expr (StgCase expr bndr alt_type alts) = do
-        expr' <- do_expr expr
-        alts' <- mapM do_alt alts
-        return (StgCase expr' bndr alt_type alts')
-      where
-        do_alt (id, bs, e) = do
-            e' <- do_expr e
-            return (id, bs, e')
-
-    do_expr (StgLet b e) = do
-          (b,e) <- do_let b e
-          return (StgLet b e)
-
-    do_expr (StgLetNoEscape b e) = do
-          (b,e) <- do_let b e
-          return (StgLetNoEscape b e)
-
-    do_expr other = pprPanic "SCCfinal.do_expr" (ppr other)
-
-    ----------------------------------
-
-    do_let (StgNonRec b rhs) e = do
-        rhs' <- do_rhs rhs
-        e' <- do_expr e
-        return (StgNonRec b rhs',e')
-
-    do_let (StgRec pairs) e = do
-        pairs' <- mapM do_pair pairs
-        e' <- do_expr e
-        return (StgRec pairs', e')
-      where
-        do_pair (b, rhs) = do
-             rhs2 <- do_rhs rhs
-             return (b, rhs2)
-
-    ----------------------------------
-    do_rhs :: StgRhs -> MassageM StgRhs
-        -- We play much the same game as we did in do_top_rhs above;
-        -- but we don't have to worry about cafs etc.
-
-        -- throw away the SCC if we don't have to count entries.  This
-        -- is a little bit wrong, because we're attributing the
-        -- allocation of the constructor to the wrong place (XXX)
-        -- We should really attach (PushCC cc CurrentCCS) to the rhs,
-        -- but need to reinstate PushCC for that.
-    do_rhs (StgRhsClosure _closure_cc _bi _fv _u []
-               (StgTick (ProfNote cc False{-not tick-} _push)
-                        (StgConApp con args _)))
-      = do collectCC cc
-           return (StgRhsCon currentCCS con args)
-
-    do_rhs (StgRhsClosure _ bi fv u args expr) = do
-        expr' <- do_expr expr
-        return (StgRhsClosure currentCCS bi fv u args expr')
-
-    do_rhs (StgRhsCon _ con args)
-      = return (StgRhsCon currentCCS con args)
-
-
--- -----------------------------------------------------------------------------
--- Boring monad stuff for this
-
-newtype MassageM result
-  = MassageM {
-      unMassageM :: Module              -- module name
-                 -> CollectedCCs
-                 -> (CollectedCCs, result)
-    }
-
-instance Functor MassageM where
-      fmap = liftM
-
-instance Applicative MassageM where
-      pure x = MassageM (\_ ccs -> (ccs, x))
-      (<*>) = ap
-      (*>) = thenMM_
-
-instance Monad MassageM where
-    (>>=) = thenMM
-    (>>)  = (*>)
-
--- the initMM function also returns the final CollectedCCs
-
-initMM :: Module        -- module name, which we may consult
-       -> MassageM a
-       -> (CollectedCCs, a)
-
-initMM mod_name (MassageM m) = m mod_name ([],[])
-
-thenMM  :: MassageM a -> (a -> MassageM b) -> MassageM b
-thenMM_ :: MassageM a -> (MassageM b) -> MassageM b
-
-thenMM expr cont = MassageM $ \mod ccs ->
-    case unMassageM expr mod ccs of { (ccs2, result) ->
-    unMassageM (cont result) mod ccs2 }
-
-thenMM_ expr cont = MassageM $ \mod ccs ->
-    case unMassageM expr mod ccs of { (ccs2, _) ->
-    unMassageM cont mod ccs2 }
-
-
-collectCC :: CostCentre -> MassageM ()
-collectCC cc
- = MassageM $ \mod_name (local_ccs, ccss)
-  -> if (cc `ccFromThisModule` mod_name) then
-        ((cc : local_ccs, ccss), ())
-     else
-        ((local_ccs, ccss), ())
-
--- Version of collectCC used when we definitely want to declare this
--- CC as local, even if its module name is not the same as the current
--- module name (eg. the special :Main module) see bug #249, #1472,
--- test prof001,prof002.
-collectNewCC :: CostCentre -> MassageM ()
-collectNewCC cc
- = MassageM $ \_mod_name (local_ccs, ccss)
-              -> ((cc : local_ccs, ccss), ())
-
-collectCCS :: CostCentreStack -> MassageM ()
-
-collectCCS ccs
- = MassageM $ \_mod_name (local_ccs, ccss)
-              -> ASSERT(not (noCCSAttached ccs))
-                       ((local_ccs, ccs : ccss), ())
diff --git a/compiler/simplStg/SimplStg.hs b/compiler/simplStg/SimplStg.hs
index 2af53e4877..6bdc1c9573 100644
--- a/compiler/simplStg/SimplStg.hs
+++ b/compiler/simplStg/SimplStg.hs
@@ -14,28 +14,23 @@ import GhcPrelude
 
 import StgSyn
 
-import CostCentre       ( CollectedCCs )
-import SCCfinal         ( stgMassageForProfiling )
 import StgLint          ( lintStgTopBindings )
 import StgStats         ( showStgStats )
 import UnariseStg       ( unarise )
 import StgCse           ( stgCse )
 
 import DynFlags
-import Module           ( Module )
 import ErrUtils
 import SrcLoc
-import UniqSupply       ( mkSplitUniqSupply, splitUniqSupply )
+import UniqSupply       ( mkSplitUniqSupply )
 import Outputable
 import Control.Monad
 
 stg2stg :: DynFlags                  -- includes spec of what stg-to-stg passes to do
-        -> Module                    -- module name (profiling only)
         -> [StgTopBinding]           -- input...
-        -> IO ( [StgTopBinding]      -- output program...
-              , CollectedCCs)        -- cost centre information (declared and used)
+        -> IO [StgTopBinding]        -- output program
 
-stg2stg dflags module_name binds
+stg2stg dflags binds
   = do  { showPass dflags "Stg2Stg"
         ; us <- mkSplitUniqSupply 'g'
 
@@ -43,23 +38,21 @@ stg2stg dflags module_name binds
                (putLogMsg dflags NoReason SevDump noSrcSpan
                   (defaultDumpStyle dflags) (text "VERBOSE STG-TO-STG:"))
 
-        ; (binds', us', ccs) <- end_pass us "Stg2Stg" ([],[]) binds
+        ; binds' <- end_pass "Stg2Stg" binds
 
                 -- Do the main business!
-        ; let (us0, us1) = splitUniqSupply us'
-        ; (processed_binds, _, cost_centres)
-                <- foldM do_stg_pass (binds', us0, ccs) (getStgToDo dflags)
+        ; processed_binds <- foldM do_stg_pass binds' (getStgToDo dflags)
 
         ; dumpIfSet_dyn dflags Opt_D_dump_stg "Pre unarise:"
                         (pprStgTopBindings processed_binds)
 
         ; let un_binds = stg_linter True "Unarise"
-                         $ unarise us1 processed_binds
+                         $ unarise us processed_binds
 
         ; dumpIfSet_dyn dflags Opt_D_dump_stg "STG syntax:"
                         (pprStgTopBindings un_binds)
 
-        ; return (un_binds, cost_centres)
+        ; return un_binds
    }
 
   where
@@ -68,38 +61,24 @@ stg2stg dflags module_name binds
       | otherwise                    = \ _whodunnit binds -> binds
 
     -------------------------------------------
-    do_stg_pass (binds, us, ccs) to_do
+    do_stg_pass binds to_do
       = case to_do of
           D_stg_stats ->
              trace (showStgStats binds)
-             end_pass us "StgStats" ccs binds
-
-          StgDoMassageForProfiling ->
-             {-# SCC "ProfMassage" #-}
-             let
-                 (us1, us2) = splitUniqSupply us
-                 (collected_CCs, binds3)
-                   = stgMassageForProfiling dflags module_name us1 binds
-             in
-             end_pass us2 "ProfMassage" collected_CCs binds3
+             end_pass "StgStats" binds
 
           StgCSE ->
              {-# SCC "StgCse" #-}
              let
                  binds' = stgCse binds
              in
-             end_pass us "StgCse" ccs binds'
+             end_pass "StgCse" binds'
 
-    end_pass us2 what ccs binds2
+    end_pass what binds2
       = do -- report verbosely, if required
            dumpIfSet_dyn dflags Opt_D_verbose_stg2stg what
               (vcat (map ppr binds2))
-           let linted_binds = stg_linter False what binds2
-           return (linted_binds, us2, ccs)
-            -- return: processed binds
-            --         UniqueSupply for the next guy to use
-            --         cost-centres to be declared/registered (specialised)
-            --         add to description of what's happened (reverse order)
+           return (stg_linter False what binds2)
 
 -- -----------------------------------------------------------------------------
 -- StgToDo:  abstraction of stg-to-stg passes to run.
@@ -107,14 +86,12 @@ stg2stg dflags module_name binds
 -- | Optional Stg-to-Stg passes.
 data StgToDo
   = StgCSE
-  | StgDoMassageForProfiling  -- should be (next to) last
   | D_stg_stats
 
 -- | Which optional Stg-to-Stg passes to run. Depends on flags, ways etc.
 getStgToDo :: DynFlags -> [StgToDo]
 getStgToDo dflags
   = [ StgCSE                   | gopt Opt_StgCSE dflags] ++
-    [ StgDoMassageForProfiling | WayProf `elem` ways dflags] ++
     [ D_stg_stats              | stg_stats ]
   where
         stg_stats = gopt Opt_StgStats dflags
diff --git a/compiler/stgSyn/CoreToStg.hs b/compiler/stgSyn/CoreToStg.hs
index 900c52e2a6..671f3eb5b5 100644
--- a/compiler/stgSyn/CoreToStg.hs
+++ b/compiler/stgSyn/CoreToStg.hs
@@ -11,7 +11,7 @@
 -- And, as we have the info in hand, we may convert some lets to
 -- let-no-escapes.
 
-module CoreToStg ( coreToStg, coreExprToStg ) where
+module CoreToStg ( coreToStg ) where
 
 #include "HsVersions.h"
 
@@ -29,10 +29,10 @@ import MkId             ( coercionTokenId )
 import Id
 import IdInfo
 import DataCon
-import CostCentre       ( noCCS )
+import CostCentre
 import VarEnv
 import Module
-import Name             ( isExternalName, nameOccName )
+import Name             ( isExternalName, nameOccName, nameModule_maybe )
 import OccName          ( occNameFS )
 import BasicTypes       ( Arity )
 import TysWiredIn       ( unboxedUnitDataCon )
@@ -46,6 +46,7 @@ import ForeignCall
 import Demand           ( isUsedOnce )
 import PrimOp           ( PrimCall(..) )
 import UniqFM
+import SrcLoc           ( mkGeneralSrcSpan )
 
 import Data.Maybe    (isJust, fromMaybe)
 import Control.Monad (liftM, ap)
@@ -196,61 +197,97 @@ import Control.Monad (liftM, ap)
 --     in
 --         ...(x b)...
 
+-- Note [Cost-centre initialization plan]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+--
+-- Previously `coreToStg` was initializing cost-centre stack fields as `noCCS`,
+-- and the fields were then fixed by a seperate pass `stgMassageForProfiling`.
+-- We now initialize these correctly. The initialization works like this:
+--
+--   - For non-top level bindings always use `currentCCS`.
+--
+--   - For top-level bindings, check if the binding is a CAF
+--
+--     - CAF:      If -fcaf-all is enabled, create a new CAF just for this CAF
+--                 and use it. Note that these new cost centres need to be
+--                 collected to be able to generate cost centre initialization
+--                 code, so `coreToTopStgRhs` now returns `CollectedCCs`.
+--
+--                 If -fcaf-all is not enabled, use "all CAFs" cost centre.
+--
+--     - Non-CAF:  Top-level (static) data is not counted in heap profiles; nor
+--                 do we set CCCS from it; so we just slam in
+--                 dontCareCostCentre.
+
 -- --------------------------------------------------------------
 -- Setting variable info: top-level, binds, RHSs
 -- --------------------------------------------------------------
 
-coreToStg :: DynFlags -> Module -> CoreProgram -> [StgTopBinding]
+coreToStg :: DynFlags -> Module -> CoreProgram
+          -> ([StgTopBinding], CollectedCCs)
 coreToStg dflags this_mod pgm
-  = pgm'
-  where (_, _, pgm') = coreTopBindsToStg dflags this_mod emptyVarEnv pgm
+  = (pgm', final_ccs)
+  where
+    (_, _, (local_ccs, local_cc_stacks), pgm')
+      = coreTopBindsToStg dflags this_mod emptyVarEnv emptyCollectedCCs pgm
 
-coreExprToStg :: CoreExpr -> StgExpr
-coreExprToStg expr
-  = new_expr where (new_expr,_) = initCts emptyVarEnv (coreToStgExpr expr)
+    prof = WayProf `elem` ways dflags
+
+    final_ccs
+      | prof && gopt Opt_AutoSccsOnIndividualCafs dflags
+      = (local_ccs,local_cc_stacks)  -- don't need "all CAFs" CC
+      | prof
+      = (all_cafs_cc:local_ccs, all_cafs_ccs:local_cc_stacks)
+      | otherwise
+      = emptyCollectedCCs
 
+    (all_cafs_cc, all_cafs_ccs) = getAllCAFsCC this_mod
 
 coreTopBindsToStg
     :: DynFlags
     -> Module
     -> IdEnv HowBound           -- environment for the bindings
+    -> CollectedCCs
     -> CoreProgram
-    -> (IdEnv HowBound, FreeVarsInfo, [StgTopBinding])
+    -> (IdEnv HowBound, FreeVarsInfo, CollectedCCs, [StgTopBinding])
 
-coreTopBindsToStg _      _        env [] = (env, emptyFVInfo, [])
-coreTopBindsToStg dflags this_mod env (b:bs)
-  = (env2, fvs2, b':bs')
+coreTopBindsToStg _      _        env ccs []
+  = (env, emptyFVInfo, ccs, [])
+coreTopBindsToStg dflags this_mod env ccs (b:bs)
+  = (env2, fvs2, ccs2, b':bs')
   where
         -- Notice the mutually-recursive "knot" here:
         --   env accumulates down the list of binds,
         --   fvs accumulates upwards
-        (env1, fvs2, b' ) = coreTopBindToStg dflags this_mod env fvs1 b
-        (env2, fvs1, bs') = coreTopBindsToStg dflags this_mod env1 bs
+        (env1, fvs2, ccs1, b' ) =
+          coreTopBindToStg dflags this_mod env fvs1 ccs b
+        (env2, fvs1, ccs2, bs') =
+          coreTopBindsToStg dflags this_mod env1 ccs1 bs
 
 coreTopBindToStg
         :: DynFlags
         -> Module
         -> IdEnv HowBound
         -> FreeVarsInfo         -- Info about the body
+        -> CollectedCCs
         -> CoreBind
-        -> (IdEnv HowBound, FreeVarsInfo, StgTopBinding)
+        -> (IdEnv HowBound, FreeVarsInfo, CollectedCCs, StgTopBinding)
 
-coreTopBindToStg _ _ env body_fvs (NonRec id (Lit (MachStr str)))
+coreTopBindToStg _ _ env body_fvs ccs (NonRec id (Lit (MachStr str)))
   -- top-level string literal
   = let
         env' = extendVarEnv env id how_bound
         how_bound = LetBound TopLet 0
-    in (env', body_fvs, StgTopStringLit id str)
+    in (env', body_fvs, ccs, StgTopStringLit id str)
 
-coreTopBindToStg dflags this_mod env body_fvs (NonRec id rhs)
+coreTopBindToStg dflags this_mod env body_fvs ccs (NonRec id rhs)
   = let
         env'      = extendVarEnv env id how_bound
         how_bound = LetBound TopLet $! manifestArity rhs
 
-        (stg_rhs, fvs') =
-            initCts env $ do
-              (stg_rhs, fvs') <- coreToTopStgRhs dflags this_mod body_fvs (id,rhs)
-              return (stg_rhs, fvs')
+        (stg_rhs, fvs', ccs') =
+            initCts env $
+              coreToTopStgRhs dflags ccs this_mod body_fvs (id,rhs)
 
         bind = StgTopLifted $ StgNonRec id stg_rhs
     in
@@ -259,9 +296,9 @@ coreTopBindToStg dflags this_mod env body_fvs (NonRec id rhs)
       --     as well as 'id', but that led to a black hole
       --     where printing the assertion error tripped the
       --     assertion again!
-    (env', fvs' `unionFVInfo` body_fvs, bind)
+    (env', fvs' `unionFVInfo` body_fvs, ccs', bind)
 
-coreTopBindToStg dflags this_mod env body_fvs (Rec pairs)
+coreTopBindToStg dflags this_mod env body_fvs ccs (Rec pairs)
   = ASSERT( not (null pairs) )
     let
         binders = map fst pairs
@@ -270,16 +307,21 @@ coreTopBindToStg dflags this_mod env body_fvs (Rec pairs)
                      | (b, rhs) <- pairs ]
         env' = extendVarEnvList env extra_env'
 
-        (stg_rhss, fvs')
+        -- generate StgTopBindings, accumulate body_fvs and CAF cost centres
+        -- created for CAFs
+        ((fvs', ccs'), stg_rhss)
           = initCts env' $ do
-               (stg_rhss, fvss') <- mapAndUnzipM (coreToTopStgRhs dflags this_mod body_fvs) pairs
-               let fvs' = unionFVInfos fvss'
-               return (stg_rhss, fvs')
+               mapAccumLM (\(fvs, ccs) rhs -> do
+                            (rhs', fvs', ccs') <-
+                              coreToTopStgRhs dflags ccs this_mod body_fvs rhs
+                            return ((fvs' `unionFVInfo` fvs, ccs'), rhs'))
+                          (body_fvs, ccs)
+                          pairs
 
         bind = StgTopLifted $ StgRec (zip binders stg_rhss)
     in
     ASSERT2(consistentCafInfo (head binders) bind, ppr binders)
-    (env', fvs' `unionFVInfo` body_fvs, bind)
+    (env', fvs' `unionFVInfo` body_fvs, ccs', bind)
 
 
 -- Assertion helper: this checks that the CafInfo on the Id matches
@@ -299,18 +341,23 @@ consistentCafInfo id bind
 
 coreToTopStgRhs
         :: DynFlags
+        -> CollectedCCs
         -> Module
         -> FreeVarsInfo         -- Free var info for the scope of the binding
         -> (Id,CoreExpr)
-        -> CtsM (StgRhs, FreeVarsInfo)
+        -> CtsM (StgRhs, FreeVarsInfo, CollectedCCs)
 
-coreToTopStgRhs dflags this_mod scope_fv_info (bndr, rhs)
+coreToTopStgRhs dflags ccs this_mod scope_fv_info (bndr, rhs)
   = do { (new_rhs, rhs_fvs) <- coreToStgExpr rhs
 
-       ; let stg_rhs   = mkTopStgRhs dflags this_mod rhs_fvs bndr bndr_info new_rhs
-             stg_arity = stgRhsArity stg_rhs
+       ; let (stg_rhs, ccs') =
+               mkTopStgRhs dflags this_mod ccs rhs_fvs bndr bndr_info new_rhs
+             stg_arity =
+               stgRhsArity stg_rhs
+
        ; return (ASSERT2( arity_ok stg_arity, mk_arity_msg stg_arity) stg_rhs,
-                 rhs_fvs) }
+                 rhs_fvs,
+                 ccs') }
   where
     bndr_info = lookupFVInfo scope_fv_info bndr
 
@@ -333,14 +380,6 @@ coreToTopStgRhs dflags this_mod scope_fv_info (bndr, rhs)
                 text "Id arity:" <+> ppr id_arity,
                 text "STG arity:" <+> ppr stg_arity]
 
-mkTopStgRhs :: DynFlags -> Module -> FreeVarsInfo
-            -> Id -> StgBinderInfo -> StgExpr
-            -> StgRhs
-
-mkTopStgRhs dflags this_mod = mkStgRhs' con_updateable
-        -- Dynamic StgConApps are updatable
-  where con_updateable con args = isDllConApp dflags this_mod con args
-
 -- ---------------------------------------------------------------------------
 -- Expressions
 -- ---------------------------------------------------------------------------
@@ -720,36 +759,86 @@ coreToStgRhs scope_fv_info (bndr, rhs) = do
   where
     bndr_info = lookupFVInfo scope_fv_info bndr
 
-mkStgRhs :: FreeVarsInfo -> Id -> StgBinderInfo -> StgExpr -> StgRhs
-mkStgRhs = mkStgRhs' con_updateable
-  where con_updateable _ _ = False
+-- Generate a top-level RHS. Any new cost centres generated for CAFs will be
+-- appended to `CollectedCCs` argument.
+mkTopStgRhs :: DynFlags -> Module -> CollectedCCs
+            -> FreeVarsInfo -> Id -> StgBinderInfo -> StgExpr
+            -> (StgRhs, CollectedCCs)
 
-mkStgRhs' :: (DataCon -> [StgArg] -> Bool)
-            -> FreeVarsInfo -> Id -> StgBinderInfo -> StgExpr -> StgRhs
-mkStgRhs' con_updateable rhs_fvs bndr binder_info rhs
+mkTopStgRhs dflags this_mod ccs rhs_fvs bndr binder_info rhs
   | StgLam bndrs body <- rhs
-  = StgRhsClosure noCCS binder_info
-                   (getFVs rhs_fvs)
-                   ReEntrant
-                   bndrs body
-  | isJoinId bndr -- must be nullary join point
-  = ASSERT(idJoinArity bndr == 0)
-    StgRhsClosure noCCS binder_info
-                   (getFVs rhs_fvs)
-                   ReEntrant -- ignored for LNE
-                   [] rhs
+  = -- StgLam can't have empty arguments, so not CAF
+    ASSERT(not (null bndrs))
+    ( StgRhsClosure dontCareCCS binder_info
+                    (getFVs rhs_fvs)
+                    ReEntrant
+                    bndrs body
+    , ccs )
+
   | StgConApp con args _ <- unticked_rhs
-  , not (con_updateable con args)
+  , -- Dynamic StgConApps are updatable
+    not (isDllConApp dflags this_mod con args)
   = -- CorePrep does this right, but just to make sure
     ASSERT2( not (isUnboxedTupleCon con || isUnboxedSumCon con)
            , ppr bndr $$ ppr con $$ ppr args)
-    StgRhsCon noCCS con args
+    ( StgRhsCon dontCareCCS con args, ccs )
+
+  -- Otherwise it's a CAF, see Note [Cost-centre initialization plan].
+  | gopt Opt_AutoSccsOnIndividualCafs dflags
+  = ( StgRhsClosure caf_ccs binder_info
+                    (getFVs rhs_fvs)
+                    upd_flag [] rhs
+    , collectCC caf_cc caf_ccs ccs )
+
   | otherwise
-  = StgRhsClosure noCCS binder_info
-                   (getFVs rhs_fvs)
-                   upd_flag [] rhs
- where
+  = ( StgRhsClosure all_cafs_ccs binder_info
+                    (getFVs rhs_fvs)
+                    upd_flag [] rhs
+    , ccs )
 
+  where
+    (_, unticked_rhs) = stripStgTicksTop (not . tickishIsCode) rhs
+
+    upd_flag | isUsedOnce (idDemandInfo bndr) = SingleEntry
+             | otherwise                      = Updatable
+
+    -- CAF cost centres generated for -fcaf-all
+    caf_cc = mkAutoCC bndr modl CafCC
+    caf_ccs = mkSingletonCCS caf_cc
+           -- careful: the binder might be :Main.main,
+           -- which doesn't belong to module mod_name.
+           -- bug #249, tests prof001, prof002
+    modl | Just m <- nameModule_maybe (idName bndr) = m
+         | otherwise = this_mod
+
+    -- default CAF cost centre
+    (_, all_cafs_ccs) = getAllCAFsCC this_mod
+
+-- Generate a non-top-level RHS. Cost-centre is always currentCCS,
+-- see Note [Cost-centre initialzation plan].
+mkStgRhs :: FreeVarsInfo -> Id -> StgBinderInfo -> StgExpr -> StgRhs
+mkStgRhs rhs_fvs bndr binder_info rhs
+  | StgLam bndrs body <- rhs
+  = StgRhsClosure currentCCS binder_info
+                  (getFVs rhs_fvs)
+                  ReEntrant
+                  bndrs body
+
+  | isJoinId bndr -- must be a nullary join point
+  = ASSERT(idJoinArity bndr == 0)
+    StgRhsClosure currentCCS binder_info
+                  (getFVs rhs_fvs)
+                  ReEntrant -- ignored for LNE
+                  [] rhs
+
+  | StgConApp con args _ <- unticked_rhs
+  = StgRhsCon currentCCS con args
+
+  | otherwise
+  = StgRhsClosure currentCCS binder_info
+                  (getFVs rhs_fvs)
+                  upd_flag [] rhs
+  where
     (_, unticked_rhs) = stripStgTicksTop (not . tickishIsCode) rhs
 
     upd_flag | isUsedOnce (idDemandInfo bndr) = SingleEntry
@@ -905,6 +994,14 @@ lookupBinding env v = case lookupVarEnv env v of
                         Just xx -> xx
                         Nothing -> ASSERT2( isGlobalId v, ppr v ) ImportBound
 
+getAllCAFsCC :: Module -> (CostCentre, CostCentreStack)
+getAllCAFsCC this_mod =
+    let
+      span = mkGeneralSrcSpan (mkFastString "<entire-module>") -- XXX do better
+      all_cafs_cc  = mkAllCafsCC this_mod span
+      all_cafs_ccs = mkSingletonCCS all_cafs_cc
+    in
+      (all_cafs_cc, all_cafs_ccs)
 
 -- ---------------------------------------------------------------------------
 -- Free variable information
diff --git a/testsuite/tests/profiling/should_compile/all.T b/testsuite/tests/profiling/should_compile/all.T
index 068b43b77e..155206ab7b 100644
--- a/testsuite/tests/profiling/should_compile/all.T
+++ b/testsuite/tests/profiling/should_compile/all.T
@@ -4,4 +4,4 @@ test('prof001', [only_ways(['normal']), req_profiling], compile_and_run, ['-prof
 test('prof002', [only_ways(['normal']), req_profiling], compile_and_run, ['-prof -fprof-cafs'])
 
 test('T2410', [only_ways(['normal']), req_profiling], compile, ['-O2 -prof -fprof-cafs'])
-test('T5889', [expect_broken(5889), only_ways(['normal']), req_profiling, extra_files(['T5889/A.hs', 'T5889/B.hs'])], multimod_compile, ['A B', '-O -prof -fno-prof-count-entries -v0'])
+test('T5889', [only_ways(['normal']), req_profiling, extra_files(['T5889/A.hs', 'T5889/B.hs'])], multimod_compile, ['A B', '-O -prof -fno-prof-count-entries -v0'])
diff --git a/testsuite/tests/simplCore/should_compile/noinline01.stderr b/testsuite/tests/simplCore/should_compile/noinline01.stderr
index 53db7dac3f..ce01fcc863 100644
--- a/testsuite/tests/simplCore/should_compile/noinline01.stderr
+++ b/testsuite/tests/simplCore/should_compile/noinline01.stderr
@@ -15,7 +15,7 @@ Noinline01.$trModule4 :: GHC.Prim.Addr#
 
 Noinline01.$trModule3 :: GHC.Types.TrName
 [GblId, Caf=NoCafRefs, Str=m1, Unf=OtherCon []] =
-    NO_CCS GHC.Types.TrNameS! [Noinline01.$trModule4];
+    CCS_DONT_CARE GHC.Types.TrNameS! [Noinline01.$trModule4];
 
 Noinline01.$trModule2 :: GHC.Prim.Addr#
 [GblId, Caf=NoCafRefs, Unf=OtherCon []] =
@@ -23,12 +23,12 @@ Noinline01.$trModule2 :: GHC.Prim.Addr#
 
 Noinline01.$trModule1 :: GHC.Types.TrName
 [GblId, Caf=NoCafRefs, Str=m1, Unf=OtherCon []] =
-    NO_CCS GHC.Types.TrNameS! [Noinline01.$trModule2];
+    CCS_DONT_CARE GHC.Types.TrNameS! [Noinline01.$trModule2];
 
 Noinline01.$trModule :: GHC.Types.Module
 [GblId, Caf=NoCafRefs, Str=m, Unf=OtherCon []] =
-    NO_CCS GHC.Types.Module! [Noinline01.$trModule3
-                              Noinline01.$trModule1];
+    CCS_DONT_CARE GHC.Types.Module! [Noinline01.$trModule3
+                                     Noinline01.$trModule1];
 
 
 
@@ -48,7 +48,7 @@ Noinline01.$trModule4 :: GHC.Prim.Addr#
 
 Noinline01.$trModule3 :: GHC.Types.TrName
 [GblId, Caf=NoCafRefs, Str=m1, Unf=OtherCon []] =
-    NO_CCS GHC.Types.TrNameS! [Noinline01.$trModule4];
+    CCS_DONT_CARE GHC.Types.TrNameS! [Noinline01.$trModule4];
 
 Noinline01.$trModule2 :: GHC.Prim.Addr#
 [GblId, Caf=NoCafRefs, Unf=OtherCon []] =
@@ -56,11 +56,11 @@ Noinline01.$trModule2 :: GHC.Prim.Addr#
 
 Noinline01.$trModule1 :: GHC.Types.TrName
 [GblId, Caf=NoCafRefs, Str=m1, Unf=OtherCon []] =
-    NO_CCS GHC.Types.TrNameS! [Noinline01.$trModule2];
+    CCS_DONT_CARE GHC.Types.TrNameS! [Noinline01.$trModule2];
 
 Noinline01.$trModule :: GHC.Types.Module
 [GblId, Caf=NoCafRefs, Str=m, Unf=OtherCon []] =
-    NO_CCS GHC.Types.Module! [Noinline01.$trModule3
-                              Noinline01.$trModule1];
+    CCS_DONT_CARE GHC.Types.Module! [Noinline01.$trModule3
+                                     Noinline01.$trModule1];