Collect CCs in CorePrep, including CCs in unfoldings
[ghc.git] / compiler / coreSyn / CorePrep.hs
index 74de5af..46474bb 100644 (file)
@@ -14,6 +14,8 @@ module CorePrep (
 
 #include "HsVersions.h"
 
+import GhcPrelude
+
 import OccurAnal
 
 import HscTypes
@@ -58,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.
@@ -122,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.
@@ -167,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))
@@ -175,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
 
@@ -185,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
 
@@ -273,7 +287,7 @@ b) The top-level binding is marked NoCafRefs.  This really happens
       $fApplicativeSTM [NoCafRefs] = D:Alternative sat ...blah...
 
    So, gruesomely, we must set the NoCafRefs flag on the sat bindings,
-   *and* substutite the modified 'sat' into the old RHS.
+   *and* substitute the modified 'sat' into the old RHS.
 
    It should be the case that 'sat' is itself [NoCafRefs] (a value, no
    cafs) else the original top-level binding would not itself have been
@@ -451,7 +465,7 @@ cpeBind top_lvl env (Rec pairs)
   where
     (bndrs, rhss) = unzip pairs
 
-        -- Flatten all the floats, and the currrent
+        -- Flatten all the floats, and the current
         -- group into a single giant Rec
     add_float (FloatLet (NonRec b r)) prs2 = (b,r) : prs2
     add_float (FloatLet (Rec prs1))   prs2 = prs1 ++ prs2
@@ -645,8 +659,21 @@ cpeRhsE env (Case scrut bndr ty alts)
        ; let bndr1 = bndr `setIdUnfolding` evaldUnfolding
             -- Record that the case binder is evaluated in the alternatives
        ; (env', bndr2) <- cpCloneBndr env bndr1
-       ; alts' <- mapM (sat_alt env') alts
-       ; return (floats, Case scrut' bndr2 ty alts') }
+       ; let alts'
+                 -- This flag is intended to aid in debugging strictness
+                 -- analysis bugs. These are particularly nasty to chase down as
+                 -- they may manifest as segmentation faults. When this flag is
+                 -- enabled we instead produce an 'error' expression to catch
+                 -- the case where a function we think should bottom
+                 -- unexpectedly returns.
+               | gopt Opt_CatchBottoms (cpe_dynFlags env)
+               , not (altsAreExhaustive alts)
+               = addDefault alts (Just err)
+               | otherwise = alts
+               where err = mkRuntimeErrorApp rUNTIME_ERROR_ID ty
+                                             "Bottoming expression returned"
+       ; alts'' <- mapM (sat_alt env') alts'
+       ; return (floats, Case scrut' bndr2 ty alts'') }
   where
     sat_alt env (con, bs, rhs)
        = do { (env2, bs') <- cpCloneBndrs env bs
@@ -774,7 +801,7 @@ cpeApp top_env expr
     collect_args :: CoreExpr -> (CoreExpr, [ArgInfo], Int)
     collect_args e = go e [] 0
       where
-        go (App fun arg)      as depth
+        go (App fun arg)      as !depth
             = go fun (CpeApp arg : as)
                 (if isTyCoArg arg then depth else depth + 1)
         go (Cast fun co)      as depth
@@ -1565,14 +1592,59 @@ newVar ty
 
 -- | Like wrapFloats, but only wraps tick floats
 wrapTicks :: Floats -> CoreExpr -> (Floats, CoreExpr)
-wrapTicks (Floats flag floats0) expr = (Floats flag floats1, expr')
-  where (floats1, expr') = foldrOL go (nilOL, expr) floats0
-        go (FloatTick t) (fs, e) = ASSERT(tickishPlace t == PlaceNonLam)
-                                   (mapOL (wrap t) fs, mkTick t e)
-        go other         (fs, e) = (other `consOL` fs, e)
+wrapTicks (Floats flag floats0) expr =
+    (Floats flag (toOL $ reverse floats1), foldr mkTick expr (reverse ticks1))
+  where (floats1, ticks1) = foldlOL go ([], []) $ floats0
+        -- Deeply nested constructors will produce long lists of
+        -- redundant source note floats here. We need to eliminate
+        -- those early, as relying on mkTick to spot it after the fact
+        -- can yield O(n^3) complexity [#11095]
+        go (floats, ticks) (FloatTick t)
+          = ASSERT(tickishPlace t == PlaceNonLam)
+            (floats, if any (flip tickishContains t) ticks
+                     then ticks else t:ticks)
+        go (floats, ticks) f
+          = (foldr wrap f (reverse ticks):floats, ticks)
+
         wrap t (FloatLet bind)    = FloatLet (wrapBind t bind)
         wrap t (FloatCase b r ok) = FloatCase b (mkTick t r) ok
         wrap _ other              = pprPanic "wrapTicks: unexpected float!"
                                              (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