Source notes (CorePrep and Stg support)

This is basically just about continuing maintaining source notes after
the Core stage. Unfortunately, this is more involved as it might seem,
as there are more restrictions on where ticks are allowed to show up.

Notes:

* We replace the StgTick / StgSCC constructors with a unified StgTick
  that can carry any tickish.

* For handling constructor or lambda applications, we generally float
  ticks out.

* Note that thanks to the NonLam placement, we know that source notes
  can never appear on lambdas. This means that as long as we are
  careful to always use mkTick, we will never violate CorePrep
  invariants.

* This is however not automatically true for eta expansion, which
  needs to somewhat awkwardly strip, then re-tick the expression in
  question.

* Where CorePrep floats out lets, we make sure to wrap them in the
  same spirit as FloatOut.

* Detecting selector thunks becomes a bit more involved, as we can run
  into ticks at multiple points.

(From Phabricator D169)

2 files changed, 36 insertions, 24 deletions

diff --git a/compiler/codeGen/StgCmmBind.hs b/compiler/codeGen/StgCmmBind.hs
index 444112f967..3c17160750 100644
--- a/compiler/codeGen/StgCmmBind.hs
+++ b/compiler/codeGen/StgCmmBind.hs
@@ -31,7 +31,7 @@ import StgCmmClosure
 import StgCmmForeign    (emitPrimCall)
 
 import MkGraph
-import CoreSyn          ( AltCon(..) )
+import CoreSyn          ( AltCon(..), tickishIsCode )
 import SMRep
 import Cmm
 import CmmInfo
@@ -50,7 +50,6 @@ import Outputable
 import FastString
 import DynFlags
 
-import Data.Maybe
 import Control.Monad
 
 #if __GLASGOW_HASKELL__ >= 709
@@ -268,14 +267,22 @@ mkRhsClosure    dflags bndr _cc _bi
                 [NonVoid the_fv]                -- Just one free var
                 upd_flag                -- Updatable thunk
                 []                      -- A thunk
-                (StgCase (StgApp scrutinee [{-no args-}])
-                      _ _ _ _   -- ignore uniq, etc.
-                      (AlgAlt _)
-                      [(DataAlt _, params, _use_mask,
-                            (StgApp selectee [{-no args-}]))])
-  |  the_fv == scrutinee                -- Scrutinee is the only free variable
-  && isJust maybe_offset                -- Selectee is a component of the tuple
-  && offset_into_int <= mAX_SPEC_SELECTEE_SIZE dflags -- Offset is small enough
+                expr
+  | let strip = snd . stripStgTicksTop (not . tickishIsCode)
+  , StgCase (StgApp scrutinee [{-no args-}])
+         _ _ _ _   -- ignore uniq, etc.
+         (AlgAlt _)
+         [(DataAlt _, params, _use_mask, sel_expr)] <- strip expr
+  , StgApp selectee [{-no args-}] <- strip sel_expr
+  , the_fv == scrutinee                -- Scrutinee is the only free variable
+
+  , let (_, _, params_w_offsets) = mkVirtConstrOffsets dflags (addIdReps params)
+                                   -- Just want the layout
+  , Just the_offset <- assocMaybe params_w_offsets (NonVoid selectee)
+
+  , let offset_into_int = bytesToWordsRoundUp dflags the_offset
+                          - fixedHdrSizeW dflags
+  , offset_into_int <= mAX_SPEC_SELECTEE_SIZE dflags -- Offset is small enough
   = -- NOT TRUE: ASSERT(is_single_constructor)
     -- The simplifier may have statically determined that the single alternative
     -- is the only possible case and eliminated the others, even if there are
@@ -284,16 +291,8 @@ mkRhsClosure    dflags bndr _cc _bi
     -- will evaluate to.
     --
     -- srt is discarded; it must be empty
-    cgRhsStdThunk bndr lf_info [StgVarArg the_fv]
-  where
-    lf_info               = mkSelectorLFInfo bndr offset_into_int
-                                 (isUpdatable upd_flag)
-    (_, _, params_w_offsets) = mkVirtConstrOffsets dflags (addIdReps params)
-                               -- Just want the layout
-    maybe_offset          = assocMaybe params_w_offsets (NonVoid selectee)
-    Just the_offset       = maybe_offset
-    offset_into_int       = bytesToWordsRoundUp dflags the_offset
-                             - fixedHdrSizeW dflags
+    let lf_info = mkSelectorLFInfo bndr offset_into_int (isUpdatable upd_flag)
+    in cgRhsStdThunk bndr lf_info [StgVarArg the_fv]
 
 ---------- Note [Ap thunks] ------------------
 mkRhsClosure    dflags bndr _cc _bi
diff --git a/compiler/codeGen/StgCmmExpr.hs b/compiler/codeGen/StgCmmExpr.hs
index b2b64f8650..9097e7fa12 100644
--- a/compiler/codeGen/StgCmmExpr.hs
+++ b/compiler/codeGen/StgCmmExpr.hs
@@ -66,10 +66,7 @@ cgExpr (StgOpApp (StgPrimOp SeqOp) [StgVarArg a, _] _res_ty) =
 
 cgExpr (StgOpApp op args ty) = cgOpApp op args ty
 cgExpr (StgConApp con args)  = cgConApp con args
-cgExpr (StgSCC cc tick push expr) = do { emitSetCCC cc tick push; cgExpr expr }
-cgExpr (StgTick m n expr) = do dflags <- getDynFlags
-                               emit (mkTickBox dflags m n)
-                               cgExpr expr
+cgExpr (StgTick t e)         = cgTick t >> cgExpr e
 cgExpr (StgLit lit)       = do cmm_lit <- cgLit lit
                                emitReturn [CmmLit cmm_lit]
 
@@ -852,3 +849,19 @@ emitEnter fun = do
        ; return (ReturnedTo lret off)
        }
   }
+
+------------------------------------------------------------------------
+--              Ticks
+------------------------------------------------------------------------
+
+-- | Generate Cmm code for a tick. Depending on the type of Tickish,
+-- this will either generate actual Cmm instrumentation code, or
+-- simply pass on the annotation as a @CmmTickish@.
+cgTick :: Tickish Id -> FCode ()
+cgTick tick
+  = do { dflags <- getDynFlags
+       ; case tick of
+           ProfNote   cc t p -> emitSetCCC cc t p
+           HpcTick    m n    -> emit (mkTickBox dflags m n)
+           _other            -> return () -- ignore
+       }