profiling: Look in RHS of rules for cost centre ticks

There are some obscure situations where the RHS of a rule can contain a
tick which is not mentioned anywhere else in the program. If this
happens you end up with an obscure linker error. The solution is quite
simple, traverse the RHS of rules to also look for ticks. It turned out
to be easier to implement if the traversal was moved into CoreTidy
rather than at the start of code generation because there we still had
easy access to the rules.

./StreamD.o(.text+0x1b9f2): error: undefined reference to 'StreamK_mkStreamFromStream_HPC_cc'
./MArray.o(.text+0xbe83): error: undefined reference to 'StreamK_mkStreamFromStream_HPC_cc'
Main.o(.text+0x6fdb): error: undefined reference to 'StreamK_mkStreamFromStream_HPC_cc'

5 files changed, 71 insertions, 52 deletions

diff --git a/compiler/GHC/CoreToStg/Prep.hs b/compiler/GHC/CoreToStg/Prep.hs
index 6c86ef990a..4529bc7d1b 100644
--- a/compiler/GHC/CoreToStg/Prep.hs
+++ b/compiler/GHC/CoreToStg/Prep.hs
@@ -20,7 +20,6 @@ where
 import GHC.Prelude
 
 import GHC.Platform
-import GHC.Platform.Ways
 
 import GHC.Driver.Session
 import GHC.Driver.Env
@@ -74,14 +73,12 @@ import GHC.Types.SrcLoc ( SrcSpan(..), realSrcLocSpan, mkRealSrcLoc )
 import GHC.Types.Literal
 import GHC.Types.Tickish
 import GHC.Types.TyThing
-import GHC.Types.CostCentre ( CostCentre, ccFromThisModule )
 import GHC.Types.Unique.Supply
 
 import GHC.Data.Pair
 import Data.List        ( unfoldr )
 import Data.Functor.Identity
 import Control.Monad
-import qualified Data.Set as S
 
 {-
 -- ---------------------------------------------------------------------------
@@ -241,20 +238,15 @@ type CpeRhs  = CoreExpr    -- Non-terminal 'rhs'
 -}
 
 corePrepPgm :: HscEnv -> Module -> ModLocation -> CoreProgram -> [TyCon]
-            -> IO (CoreProgram, S.Set CostCentre)
+            -> IO CoreProgram
 corePrepPgm hsc_env this_mod mod_loc binds data_tycons =
     withTiming logger
                (text "CorePrep"<+>brackets (ppr this_mod))
-               (\(a,b) -> a `seqList` b `seq` ()) $ do
+               (\a -> a `seqList` ()) $ do
     us <- mkSplitUniqSupply 's'
     initialCorePrepEnv <- mkInitialCorePrepEnv hsc_env
 
-    let cost_centres
-          | WayProf `S.member` ways dflags
-          = collectCostCentres this_mod binds
-          | otherwise
-          = S.empty
-
+    let
         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
@@ -265,7 +257,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, cost_centres)
+    return binds_out
   where
     dflags = hsc_dflags hsc_env
     logger = hsc_logger hsc_env
@@ -2120,41 +2112,6 @@ wrapTicks (Floats flag floats0) expr =
         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 (Alt _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/GHC/Driver/Main.hs b/compiler/GHC/Driver/Main.hs
index 2f40d7a00b..9a55807c0d 100644
--- a/compiler/GHC/Driver/Main.hs
+++ b/compiler/GHC/Driver/Main.hs
@@ -1565,6 +1565,7 @@ hscGenHardCode hsc_env cgguts location output_filename = do
                     -- From now on, we just use the bits we need.
                     cg_module   = this_mod,
                     cg_binds    = core_binds,
+                    cg_ccs      = local_ccs,
                     cg_tycons   = tycons,
                     cg_foreign  = foreign_stubs0,
                     cg_foreign_files = foreign_files,
@@ -1582,7 +1583,7 @@ hscGenHardCode hsc_env cgguts location output_filename = do
         -------------------
         -- PREPARE FOR CODE GENERATION
         -- Do saturation and convert to A-normal form
-        (prepd_binds, local_ccs) <- {-# SCC "CorePrep" #-}
+        (prepd_binds) <- {-# SCC "CorePrep" #-}
                        corePrepPgm hsc_env this_mod location
                                    core_binds data_tycons
 
@@ -1595,7 +1596,7 @@ hscGenHardCode hsc_env cgguts location output_filename = do
                    (myCoreToStg logger dflags (hsc_IC hsc_env) this_mod location prepd_binds)
 
         let cost_centre_info =
-              (S.toList local_ccs ++ caf_ccs, caf_cc_stacks)
+              (local_ccs ++ caf_ccs, caf_cc_stacks)
             platform = targetPlatform dflags
             prof_init
               | sccProfilingEnabled dflags = profilingInitCode platform this_mod cost_centre_info
@@ -1661,7 +1662,7 @@ hscInteractive hsc_env cgguts location = 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
 
     (stg_binds, _infotable_prov, _caf_ccs__caf_cc_stacks)
@@ -1978,7 +1979,7 @@ hscParsedDecls hsc_env decls = runInteractiveHsc hsc_env $ do
 
     {- 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
 
     (stg_binds, _infotable_prov, _caf_ccs__caf_cc_stacks)
diff --git a/compiler/GHC/HsToCore.hs b/compiler/GHC/HsToCore.hs
index 1c11c17ac1..73c6accff4 100644
--- a/compiler/GHC/HsToCore.hs
+++ b/compiler/GHC/HsToCore.hs
@@ -45,13 +45,13 @@ import GHC.Core.FVs       ( exprsSomeFreeVarsList )
 import GHC.Core.SimpleOpt ( simpleOptPgm, simpleOptExpr )
 import GHC.Core.Utils
 import GHC.Core.Unfold.Make
-import GHC.Core.Ppr
 import GHC.Core.Coercion
 import GHC.Core.DataCon ( dataConWrapId )
 import GHC.Core.Make
 import GHC.Core.Rules
 import GHC.Core.Opt.Monad ( CoreToDo(..) )
 import GHC.Core.Lint     ( endPassIO )
+import GHC.Core.Ppr
 
 import GHC.Builtin.Names
 import GHC.Builtin.Types.Prim
diff --git a/compiler/GHC/Iface/Tidy.hs b/compiler/GHC/Iface/Tidy.hs
index 101d470bdc..3285fb88e5 100644
--- a/compiler/GHC/Iface/Tidy.hs
+++ b/compiler/GHC/Iface/Tidy.hs
@@ -2,6 +2,7 @@
 {-# LANGUAGE DeriveFunctor #-}
 
 {-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
+{-# LANGUAGE NamedFieldPuns #-}
 
 {-
 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
@@ -77,6 +78,9 @@ import GHC.Data.Maybe
 import Control.Monad
 import Data.Function
 import Data.List        ( sortBy, mapAccumL )
+import qualified Data.Set as S
+import GHC.Platform.Ways
+import GHC.Types.CostCentre
 
 {-
 Constructing the TypeEnv, Instances, Rules from which the
@@ -429,6 +433,13 @@ tidyProgram hsc_env  (ModGuts { mg_module           = mod
               --  (c) Constructors even if they are not exported (the
               --      tidied TypeEnv has trimmed these away)
               ; alg_tycons = filter isAlgTyCon tcs
+
+
+              ; local_ccs
+                  | WayProf `S.member` ways dflags
+                        = collectCostCentres mod all_tidy_binds tidy_rules
+                  | otherwise
+                        = S.empty
               }
 
         ; endPassIO hsc_env print_unqual CoreTidy all_tidy_binds tidy_rules
@@ -454,6 +465,7 @@ tidyProgram hsc_env  (ModGuts { mg_module           = mod
         ; return (CgGuts { cg_module   = mod,
                            cg_tycons   = alg_tycons,
                            cg_binds    = all_tidy_binds,
+                           cg_ccs      = S.toList local_ccs,
                            cg_foreign  = add_spt_init_code foreign_stubs,
                            cg_foreign_files = foreign_files,
                            cg_dep_pkgs = dep_direct_pkgs deps,
@@ -474,6 +486,53 @@ tidyProgram hsc_env  (ModGuts { mg_module           = mod
     dflags = hsc_dflags hsc_env
     logger = hsc_logger hsc_env
 
+
+------------------------------------------------------------------------------
+-- Collecting cost centres
+-- ---------------------------------------------------------------------------
+
+-- | Collect cost centres defined in the current module, including those in
+-- unfoldings.
+collectCostCentres :: Module -> CoreProgram -> [CoreRule] -> S.Set CostCentre
+collectCostCentres mod_name binds rules
+  = foldl' go_bind (go_rules S.empty) binds
+  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 (Alt _con _bndrs e) -> go cs e)
+
+    go_bind :: S.Set CostCentre -> CoreBind -> S.Set CostCentre
+    go_bind cs (NonRec b e) =
+      go (do_binder cs b) e
+    go_bind cs (Rec bs) =
+      foldl' (\cs' (b, e) -> go (do_binder cs' b) e) cs bs
+
+    do_binder cs b = maybe cs (go cs) (get_unf b)
+
+
+    -- Unfoldings may have cost centres that in the original definion are
+    -- optimized away, see #5889.
+    get_unf = maybeUnfoldingTemplate . realIdUnfolding
+
+    -- Have to look at the RHS of rules as well, as these may contain ticks which
+    -- don't appear anywhere else. See #19894
+    go_rules cs = foldl' go cs (mapMaybe get_rhs rules)
+
+    get_rhs Rule { ru_rhs } = Just ru_rhs
+    get_rhs BuiltinRule {} = Nothing
+
 --------------------------
 trimId :: Bool -> Id -> Id
 -- With -O0 we now trim off the arity, one-shot-ness, strictness
diff --git a/compiler/GHC/Unit/Module/ModGuts.hs b/compiler/GHC/Unit/Module/ModGuts.hs
index b40c980744..e799ebf2a1 100644
--- a/compiler/GHC/Unit/Module/ModGuts.hs
+++ b/compiler/GHC/Unit/Module/ModGuts.hs
@@ -34,6 +34,7 @@ import GHC.Types.Name.Reader
 import GHC.Types.SafeHaskell
 import GHC.Types.SourceFile ( HscSource(..), hscSourceToIsBoot )
 import GHC.Types.SrcLoc
+import GHC.Types.CostCentre
 
 
 -- | A ModGuts is carried through the compiler, accumulating stuff as it goes
@@ -131,6 +132,7 @@ data CgGuts
                 -- data constructor workers; reason: we regard them
                 -- as part of the code-gen of tycons
 
+        cg_ccs       :: [CostCentre], -- List of cost centres used in bindings and rules
         cg_foreign   :: !ForeignStubs,   -- ^ Foreign export stubs
         cg_foreign_files :: ![(ForeignSrcLang, FilePath)],
         cg_dep_pkgs  :: ![UnitId], -- ^ Dependent packages, used to