Use Info Table Provenances to decode cloned stack (#18163)

Emit an Info Table Provenance Entry (IPE) for every stack represeted info table
if -finfo-table-map is turned on.

To decode a cloned stack, lookupIPE() is used. It provides a mapping between
info tables and their source location.

Please see these notes for details:

- [Stacktraces from Info Table Provenance Entries (IPE based stack unwinding)]
- [Mapping Info Tables to Source Positions]

Metric Increase:
    T12545

14 files changed, 356 insertions, 81 deletions

diff --git a/compiler/GHC/Builtin/primops.txt.pp b/compiler/GHC/Builtin/primops.txt.pp
index 58d4ec91f3..2aa1eefb4b 100644
--- a/compiler/GHC/Builtin/primops.txt.pp
+++ b/compiler/GHC/Builtin/primops.txt.pp
@@ -3648,19 +3648,9 @@ primop  SetThreadAllocationCounter "setThreadAllocationCounter#" GenPrimOp
    out_of_line      = True
 
 primtype StackSnapshot#
-
-primop  CloneMyStack "cloneMyStack#" GenPrimOp
-   State# RealWorld -> (# State# RealWorld, StackSnapshot# #)
-   { Clones the stack of the current (active) Haskell thread. A cloned stack is
-     represented by {\tt StackSnapshot# } and is not evaluated any further
-     (i.e. it's "cold"). This is useful for stack decoding (backtraces) and
-     analyses because there are no concurrent mutations on a cloned stack.
-     The module {\tt GHC.Stack.CloneStack } contains related funcions.
-     Please see Note [Stack Cloning] for technical details. }
-   with
-   has_side_effects = True
-   out_of_line      = True
-
+   { Haskell representation of a {\tt StgStack*} that was created (cloned)
+     with a function in {\tt GHC.Stack.CloneStack}. Please check the
+     documentation in this module for more detailed explanations. }
 
 ------------------------------------------------------------------------
 section "Safe coercions"
diff --git a/compiler/GHC/Cmm.hs b/compiler/GHC/Cmm.hs
index 3a461fa03c..893ca556db 100644
--- a/compiler/GHC/Cmm.hs
+++ b/compiler/GHC/Cmm.hs
@@ -167,12 +167,12 @@ data CmmInfoTable
         -- place to convey this information from the code generator to
         -- where we build the static closures in
         -- GHC.Cmm.Info.Build.doSRTs.
-    }
+    } deriving Eq
 
 data ProfilingInfo
   = NoProfilingInfo
   | ProfilingInfo ByteString ByteString -- closure_type, closure_desc
-
+  deriving Eq
 -----------------------------------------------------------------------------
 --              Static Data
 -----------------------------------------------------------------------------
@@ -288,4 +288,3 @@ instance OutputableP env instr => OutputableP env (GenBasicBlock instr) where
 pprBBlock :: Outputable stmt => GenBasicBlock stmt -> SDoc
 pprBBlock (BasicBlock ident stmts) =
     hang (ppr ident <> colon) 4 (vcat (map ppr stmts))
-
diff --git a/compiler/GHC/Cmm/Ppr.hs b/compiler/GHC/Cmm/Ppr.hs
index 0f846bad1b..455a7d639a 100644
--- a/compiler/GHC/Cmm/Ppr.hs
+++ b/compiler/GHC/Cmm/Ppr.hs
@@ -64,6 +64,8 @@ import GHC.Cmm.Dataflow.Graph
 
 -------------------------------------------------
 -- Outputable instances
+instance OutputableP Platform InfoProvEnt where
+  pdoc platform (InfoProvEnt clabel _ _ _ _) = pdoc platform clabel
 
 instance Outputable CmmStackInfo where
     ppr = pprStackInfo
diff --git a/compiler/GHC/Driver/GenerateCgIPEStub.hs b/compiler/GHC/Driver/GenerateCgIPEStub.hs
new file mode 100644
index 0000000000..e0b0deaa83
--- /dev/null
+++ b/compiler/GHC/Driver/GenerateCgIPEStub.hs
@@ -0,0 +1,266 @@
+{-# LANGUAGE GADTs #-}
+
+module GHC.Driver.GenerateCgIPEStub (generateCgIPEStub) where
+
+import qualified Data.Map.Strict as Map
+import Data.Maybe (catMaybes, listToMaybe)
+import GHC.Cmm
+import GHC.Cmm.CLabel (CLabel)
+import GHC.Cmm.Dataflow (Block, C, O)
+import GHC.Cmm.Dataflow.Block (blockSplit, blockToList)
+import GHC.Cmm.Dataflow.Collections (mapToList)
+import GHC.Cmm.Dataflow.Label (Label)
+import GHC.Cmm.Info.Build (emptySRT)
+import GHC.Cmm.Pipeline (cmmPipeline)
+import GHC.Cmm.Utils (toBlockList)
+import GHC.Data.Maybe (firstJusts)
+import GHC.Data.Stream (Stream, liftIO)
+import qualified GHC.Data.Stream as Stream
+import GHC.Driver.Env (hsc_dflags)
+import GHC.Driver.Flags (GeneralFlag (Opt_InfoTableMap))
+import GHC.Driver.Session (gopt, targetPlatform)
+import GHC.Plugins (HscEnv, NonCaffySet)
+import GHC.Prelude
+import GHC.Runtime.Heap.Layout (isStackRep)
+import GHC.Settings (Platform, platformUnregisterised)
+import GHC.StgToCmm.Monad (getCmm, initC, runC)
+import GHC.StgToCmm.Prof (initInfoTableProv)
+import GHC.StgToCmm.Types (CgInfos (..), ModuleLFInfos)
+import GHC.Types.IPE (InfoTableProvMap (provInfoTables), IpeSourceLocation)
+import GHC.Types.Tickish (GenTickish (SourceNote))
+import GHC.Unit.Types (Module)
+
+{-
+Note [Stacktraces from Info Table Provenance Entries (IPE based stack unwinding)]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Stacktraces can be created from return frames as they are pushed to stack for every case scrutinee.
+But to make them readable / meaningful, one needs to know the source location of each return frame.
+
+Every return frame has a distinct info table and thus a distinct code pointer (for tables next to
+code) or at least a distict address itself. Info Table Provernance Entries (IPE) are searchable by
+this pointer and contain a source location.
+
+The info table / info table code pointer to source location map is described in:
+Note [Mapping Info Tables to Source Positions]
+
+To be able to lookup IPEs for return frames one needs to emit them during compile time. This is done
+by `generateCgIPEStub`.
+
+This leads to the question: How to figure out the source location of a return frame?
+
+While the lookup algorithms for registerised and unregisterised builds differ in details, they have in
+common that we want to lookup the `CmmNode.CmmTick` (containing a `SourceNote`) that is nearest
+(before) the usage of the return frame's label. (Which label and label type is used differs between
+these two use cases.)
+
+Registerised
+~~~~~~~~~~~~~
+
+Let's consider this example:
+```
+ Main.returnFrame_entry() { //  [R2]
+         { info_tbls: [(c18g,
+                        label: block_c18g_info
+                        rep: StackRep []
+                        srt: Just GHC.CString.unpackCString#_closure),
+                       (c18r,
+                        label: Main.returnFrame_info
+                        rep: HeapRep static { Fun {arity: 1 fun_type: ArgSpec 5} }
+                        srt: Nothing)]
+           stack_info: arg_space: 8
+         }
+     {offset
+
+      [...]
+
+       c18u: // global
+           //tick src<Main.hs:(7,1)-(16,15)>
+           I64[Hp - 16] = sat_s16B_info;
+           P64[Hp] = _s16r::P64;
+           _c17j::P64 = Hp - 16;
+           //tick src<Main.hs:8:25-39>
+           I64[Sp - 8] = c18g;
+           R3 = _c17j::P64;
+           R2 = GHC.IO.Unsafe.unsafePerformIO_closure;
+           R1 = GHC.Base.$_closure;
+           Sp = Sp - 8;
+           call stg_ap_pp_fast(R3,
+                               R2,
+                               R1) returns to c18g, args: 8, res: 8, upd: 8;
+```
+
+The return frame `block_c18g_info` has the label `c18g` which is used in the call to `stg_ap_pp_fast`
+(`returns to c18g`) as continuation (`cml_cont`). The source location we're after, is the nearest
+`//tick` before the call (`//tick src<Main.hs:8:25-39>`).
+
+In code the Cmm program is represented as a Hoopl graph. Hoopl distinguishes nodes by defining if they
+are open or closed on entry (one can fallthrough to them from the previous instruction) and if they are
+open or closed on exit (one can fallthrough from them to the next node).
+
+Please refer to the paper "Hoopl: A Modular, Reusable Library for Dataflow Analysis and Transformation"
+for a detailed explanation.
+
+Here we use the fact, that calls (represented by `CmmNode.CmmCall`) are always closed on exit
+(`CmmNode O C`, `O` means open, `C` closed). In other words, they are always at the end of a block.
+
+So, given a stack represented info table (likely representing a return frame, but this isn't completely
+sure as there are e.g. update frames, too) with it's label (`c18g` in the example above) and a `CmmGraph`:
+  - Look at the end of every block, if it's a `CmmNode.CmmCall` returning to the continuation with the
+    label of the return frame.
+  - If there's such a call, lookup the nearest `CmmNode.CmmTick` by traversing the middle part of the block
+    backwards (from end to beginning).
+  - Take the first `CmmNode.CmmTick` that contains a `Tickish.SourceNote` and return it's payload as
+    `IpeSourceLocation`. (There are other `Tickish` constructors like `ProfNote` or `HpcTick`, these are
+    ignored.)
+
+Unregisterised
+~~~~~~~~~~~~~
+
+In unregisterised builds there is no return frame / continuation label in calls. The continuation (i.e. return
+frame) is set in an explicit Cmm assignment. Thus the tick lookup algorithm has to be slightly different.
+
+```
+ sat_s16G_entry() { //  [R1]
+         { info_tbls: [(c18O,
+                        label: sat_s16G_info
+                        rep: HeapRep { Thunk }
+                        srt: Just _u18Z_srt)]
+           stack_info: arg_space: 0
+         }
+     {offset
+       c18O: // global
+           _s16G::P64 = R1;
+           if ((Sp + 8) - 40 < SpLim) (likely: False) goto c18P; else goto c18Q;
+       c18P: // global
+           R1 = _s16G::P64;
+           call (stg_gc_enter_1)(R1) args: 8, res: 0, upd: 8;
+       c18Q: // global
+           I64[Sp - 16] = stg_upd_frame_info;
+           P64[Sp - 8] = _s16G::P64;
+           //tick src<Main.hs:20:9-13>
+           I64[Sp - 24] = block_c18M_info;
+           R1 = GHC.Show.$fShow[]_closure;
+           P64[Sp - 32] = GHC.Show.$fShowChar_closure;
+           Sp = Sp - 32;
+           call stg_ap_p_fast(R1) args: 16, res: 8, upd: 24;
+     }
+ },
+ _blk_c18M() { //  [R1]
+         { info_tbls: [(c18M,
+                        label: block_c18M_info
+                        rep: StackRep []
+                        srt: Just System.IO.print_closure)]
+           stack_info: arg_space: 0
+         }
+     {offset
+       c18M: // global
+           _s16F::P64 = R1;
+           R1 = System.IO.print_closure;
+           P64[Sp] = _s16F::P64;
+           call stg_ap_p_fast(R1) args: 32, res: 0, upd: 24;
+     }
+ },
+```
+
+In this example we have to lookup `//tick src<Main.hs:20:9-13>` for the return frame `c18M`.
+Notice, that this cannot be done with the `Label` `c18M`, but with the `CLabel` `block_c18M_info`
+(`label: block_c18M_info` is actually a `CLabel`).
+
+The find the tick:
+  - Every `Block` is checked from top (first) to bottom (last) node for an assignment like
+   `I64[Sp - 24] = block_c18M_info;`. The lefthand side is actually ignored.
+  - If such an assignment is found the search is over, because the payload (content of
+    `Tickish.SourceNote`, represented as `IpeSourceLocation`) of last visited tick is always
+    remembered in a `Maybe`.
+-}
+
+generateCgIPEStub :: HscEnv -> Module -> InfoTableProvMap -> Stream IO CmmGroupSRTs (NonCaffySet, ModuleLFInfos) -> Stream IO CmmGroupSRTs CgInfos
+generateCgIPEStub hsc_env this_mod denv s = do
+  let dflags = hsc_dflags hsc_env
+      platform = targetPlatform dflags
+  cgState <- liftIO initC
+
+  -- Collect info tables, but only if -finfo-table-map is enabled, otherwise labeledInfoTablesWithTickishes is empty.
+  let collectFun = if gopt Opt_InfoTableMap dflags then collect platform else collectNothing
+  (labeledInfoTablesWithTickishes, (nonCaffySet, moduleLFInfos)) <- Stream.mapAccumL_ collectFun [] s
+
+  -- Yield Cmm for Info Table Provenance Entries (IPEs)
+  let denv' = denv {provInfoTables = Map.fromList (map (\(_, i, t) -> (cit_lbl i, t)) labeledInfoTablesWithTickishes)}
+      ((ipeStub, ipeCmmGroup), _) = runC dflags this_mod cgState $ getCmm (initInfoTableProv (map sndOfTriple labeledInfoTablesWithTickishes) denv' this_mod)
+
+  (_, ipeCmmGroupSRTs) <- liftIO $ cmmPipeline hsc_env (emptySRT this_mod) ipeCmmGroup
+  Stream.yield ipeCmmGroupSRTs
+
+  return CgInfos {cgNonCafs = nonCaffySet, cgLFInfos = moduleLFInfos, cgIPEStub = ipeStub}
+  where
+    collect :: Platform -> [(Label, CmmInfoTable, Maybe IpeSourceLocation)] -> CmmGroupSRTs -> IO ([(Label, CmmInfoTable, Maybe IpeSourceLocation)], CmmGroupSRTs)
+    collect platform acc cmmGroupSRTs = do
+      let labelsToInfoTables = collectInfoTables cmmGroupSRTs
+          labelsToInfoTablesToTickishes = map (\(l, i) -> (l, i, lookupEstimatedTick platform cmmGroupSRTs l i)) labelsToInfoTables
+      return (acc ++ labelsToInfoTablesToTickishes, cmmGroupSRTs)
+
+    collectNothing :: [a] -> CmmGroupSRTs -> IO ([a], CmmGroupSRTs)
+    collectNothing _ cmmGroupSRTs = pure ([], cmmGroupSRTs)
+
+    sndOfTriple :: (a, b, c) -> b
+    sndOfTriple (_, b, _) = b
+
+    collectInfoTables :: CmmGroupSRTs -> [(Label, CmmInfoTable)]
+    collectInfoTables cmmGroup = concat $ catMaybes $ map extractInfoTables cmmGroup
+
+    extractInfoTables :: GenCmmDecl RawCmmStatics CmmTopInfo CmmGraph -> Maybe [(Label, CmmInfoTable)]
+    extractInfoTables (CmmProc h _ _ _) = Just $ mapToList (info_tbls h)
+    extractInfoTables _ = Nothing
+
+    lookupEstimatedTick :: Platform -> CmmGroupSRTs -> Label -> CmmInfoTable -> Maybe IpeSourceLocation
+    lookupEstimatedTick platform cmmGroup infoTableLabel infoTable = do
+      -- All return frame info tables are stack represented, though not all stack represented info
+      -- tables have to be return frames.
+      if (isStackRep . cit_rep) infoTable
+        then do
+          let findFun =
+                if platformUnregisterised platform
+                  then findCmmTickishForForUnregistered (cit_lbl infoTable)
+                  else findCmmTickishForRegistered infoTableLabel
+              blocks = concatMap toBlockList (graphs cmmGroup)
+          firstJusts $ map findFun blocks
+        else Nothing
+    graphs :: CmmGroupSRTs -> [CmmGraph]
+    graphs = foldl' go []
+      where
+        go :: [CmmGraph] -> GenCmmDecl d h CmmGraph -> [CmmGraph]
+        go acc (CmmProc _ _ _ g) = g : acc
+        go acc _ = acc
+
+    findCmmTickishForRegistered :: Label -> Block CmmNode C C -> Maybe IpeSourceLocation
+    findCmmTickishForRegistered label block = do
+      let (_, middleBlock, endBlock) = blockSplit block
+
+      isCallWithReturnFrameLabel endBlock label
+      lastTickInBlock middleBlock
+      where
+        isCallWithReturnFrameLabel :: CmmNode O C -> Label -> Maybe ()
+        isCallWithReturnFrameLabel (CmmCall _ (Just l) _ _ _ _) clabel | l == clabel = Just ()
+        isCallWithReturnFrameLabel _ _ = Nothing
+
+        lastTickInBlock block =
+          listToMaybe $
+            catMaybes $
+              map maybeTick $ (reverse . blockToList) block
+
+        maybeTick :: CmmNode O O -> Maybe IpeSourceLocation
+        maybeTick (CmmTick (SourceNote span name)) = Just (span, name)
+        maybeTick _ = Nothing
+
+    findCmmTickishForForUnregistered :: CLabel -> Block CmmNode C C -> Maybe IpeSourceLocation
+    findCmmTickishForForUnregistered cLabel block = do
+      let (_, middleBlock, _) = blockSplit block
+      find cLabel (blockToList middleBlock) Nothing
+      where
+        find :: CLabel -> [CmmNode O O] -> Maybe IpeSourceLocation -> Maybe IpeSourceLocation
+        find label (b : blocks) lastTick = case b of
+          (CmmStore _ (CmmLit (CmmLabel l))) -> if label == l then lastTick else find label blocks lastTick
+          (CmmTick (SourceNote span name)) -> find label blocks $ Just (span, name)
+          _ -> find label blocks lastTick
+        find _ [] _ = Nothing
diff --git a/compiler/GHC/Driver/Hooks.hs b/compiler/GHC/Driver/Hooks.hs
index e4f4262d5e..ded0683ec0 100644
--- a/compiler/GHC/Driver/Hooks.hs
+++ b/compiler/GHC/Driver/Hooks.hs
@@ -49,7 +49,6 @@ import GHC.Types.CostCentre
 import GHC.Types.IPE
 import GHC.Types.Meta
 import GHC.Types.HpcInfo
-import GHC.Types.ForeignStubs
 
 import GHC.Unit.Module
 import GHC.Unit.Module.ModSummary
@@ -146,7 +145,7 @@ data Hooks = Hooks
                                                           -> IO (Maybe HValue)))
   , createIservProcessHook :: !(Maybe (CreateProcess -> IO ProcessHandle))
   , stgToCmmHook           :: !(Maybe (DynFlags -> Module -> InfoTableProvMap -> [TyCon] -> CollectedCCs
-                                 -> [CgStgTopBinding] -> HpcInfo -> Stream IO CmmGroup (CStub, ModuleLFInfos)))
+                                 -> [CgStgTopBinding] -> HpcInfo -> Stream IO CmmGroup ModuleLFInfos))
   , cmmToRawCmmHook        :: !(forall a . Maybe (DynFlags -> Maybe Module -> Stream IO CmmGroupSRTs a
                                  -> IO (Stream IO RawCmmGroup a)))
   }
diff --git a/compiler/GHC/Driver/Main.hs b/compiler/GHC/Driver/Main.hs
index a01c559c80..3605b4ac5a 100644
--- a/compiler/GHC/Driver/Main.hs
+++ b/compiler/GHC/Driver/Main.hs
@@ -2,6 +2,7 @@
 
 {-# LANGUAGE NondecreasingIndentation #-}
 {-# LANGUAGE TupleSections #-}
+{-# LANGUAGE GADTs #-}
 
 {-# OPTIONS_GHC -fprof-auto-top #-}
 
@@ -236,6 +237,9 @@ import Control.DeepSeq (force)
 import Data.Bifunctor (first)
 import GHC.Data.Maybe
 import GHC.Driver.Env.KnotVars
+import GHC.Types.Name.Set (NonCaffySet)
+import GHC.Driver.GenerateCgIPEStub (generateCgIPEStub)
+
 
 {- **********************************************************************
 %*                                                                      *
@@ -1756,7 +1760,7 @@ doCodeGen hsc_env this_mod denv data_tycons
                         Nothing -> StgToCmm.codeGen logger tmpfs
                         Just h  -> h
 
-    let cmm_stream :: Stream IO CmmGroup (CStub, ModuleLFInfos)
+    let cmm_stream :: Stream IO CmmGroup ModuleLFInfos
         -- See Note [Forcing of stg_binds]
         cmm_stream = stg_binds_w_fvs `seqList` {-# SCC "StgToCmm" #-}
             stg_to_cmm dflags this_mod denv data_tycons cost_centre_info stg_binds_w_fvs hpc_info
@@ -1774,21 +1778,21 @@ doCodeGen hsc_env this_mod denv data_tycons
 
         ppr_stream1 = Stream.mapM dump1 cmm_stream
 
-        pipeline_stream :: Stream IO CmmGroupSRTs CgInfos
+        pipeline_stream :: Stream IO CmmGroupSRTs (NonCaffySet, ModuleLFInfos)
         pipeline_stream = do
-          (non_cafs, (used_info, lf_infos)) <-
+          (non_cafs,  lf_infos) <-
             {-# SCC "cmmPipeline" #-}
             Stream.mapAccumL_ (cmmPipeline hsc_env) (emptySRT this_mod) ppr_stream1
               <&> first (srtMapNonCAFs . moduleSRTMap)
 
-          return CgInfos{ cgNonCafs = non_cafs, cgLFInfos = lf_infos, cgIPEStub = used_info }
+          return (non_cafs, lf_infos)
 
         dump2 a = do
           unless (null a) $
             putDumpFileMaybe logger Opt_D_dump_cmm "Output Cmm" FormatCMM (pdoc platform a)
           return a
 
-    return (Stream.mapM dump2 pipeline_stream)
+    return $ Stream.mapM dump2 $ generateCgIPEStub hsc_env this_mod denv pipeline_stream
 
 myCoreToStgExpr :: Logger -> DynFlags -> InteractiveContext
                 -> Bool
diff --git a/compiler/GHC/Runtime/Heap/Layout.hs b/compiler/GHC/Runtime/Heap/Layout.hs
index 1195f83937..3e6f6ed405 100644
--- a/compiler/GHC/Runtime/Heap/Layout.hs
+++ b/compiler/GHC/Runtime/Heap/Layout.hs
@@ -174,6 +174,7 @@ data SMRep
   | RTSRep              -- The RTS needs to declare info tables with specific
         Int             -- type tags, so this form lets us override the default
         SMRep           -- tag for an SMRep.
+  deriving Eq
 
 -- | True \<=> This is a static closure.  Affects how we garbage-collect it.
 -- Static closure have an extra static link field at the end.
@@ -191,6 +192,7 @@ data ClosureTypeInfo
   | ThunkSelector SelectorOffset
   | BlackHole
   | IndStatic
+  deriving Eq
 
 type ConstrDescription = ByteString -- result of dataConIdentity
 type FunArity          = Int
@@ -445,6 +447,8 @@ rtsClosureType rep
       HeapRep False _ _ BlackHole -> BLACKHOLE
       HeapRep False _ _ IndStatic -> IND_STATIC
 
+      StackRep _ -> STACK
+
       _ -> panic "rtsClosureType"
 
 -- We export these ones
diff --git a/compiler/GHC/Stg/Debug.hs b/compiler/GHC/Stg/Debug.hs
index 77ef7910ec..bea6fe5c8e 100644
--- a/compiler/GHC/Stg/Debug.hs
+++ b/compiler/GHC/Stg/Debug.hs
@@ -33,7 +33,7 @@ data R = R { rDynFlags :: DynFlags, rModLocation :: ModLocation, rSpan :: Maybe
 
 type M a = ReaderT R (State InfoTableProvMap) a
 
-withSpan :: (RealSrcSpan, String) -> M a -> M a
+withSpan :: IpeSourceLocation -> M a -> M a
 withSpan (new_s, new_l) act = local maybe_replace act
   where
     maybe_replace r@R{ rModLocation = cur_mod, rSpan = Just (SpanWithLabel old_s _old_l) }
@@ -171,16 +171,21 @@ to a position in the source. The prime example is being able to map a THUNK to
 a specific place in the source program, the mapping is usually quite precise because
 a fresh info table is created for each distinct THUNK.
 
+The info table map is also used to generate stacktraces.
+See Note [Stacktraces from Info Table Provenance Entries (IPE based stack unwinding)]
+for details.
+
 There are three parts to the implementation
 
-1. In GHC.Stg.Debug, the SourceNote information is used in order to give a source location to
-some specific closures.
-2. In StgToCmm, the actually used info tables are recorded in an IORef, this
-is important as it's hard to predict beforehand what code generation will do
-and which ids will end up in the generated program.
-3. During code generation, a mapping from the info table to the statically
-determined location is emitted which can then be queried at runtime by
-various tools.
+1. In GHC.Stg.Debug, the SourceNote information is used in order to give a source location
+   to some specific closures.
+2. In GHC.Driver.GenerateCgIPEStub, the actually used info tables are collected after the
+   Cmm pipeline. This is important as it's hard to predict beforehand what code generation
+   will do and which ids will end up in the generated program. Additionally, info tables of
+   return frames (used to create stacktraces) are generated in the Cmm pipeline and aren't
+   available before.
+3. During code generation, a mapping from the info table to the statically determined location
+   is emitted which can then be queried at runtime by various tools.
 
 -- Giving Source Locations to Closures
 
@@ -189,6 +194,8 @@ is collected in the `InfoTableProvMap` which provides a mapping from:
 
 1. Data constructors to a list of where they are used.
 2. `Name`s and where they originate from.
+3. Stack represented info tables (return frames) to an approximated source location
+   of the call that pushed a contiunation on the stacks.
 
 During the CoreToStg phase, this map is populated whenever something is turned into
 a StgRhsClosure or an StgConApp. The current source position is recorded
@@ -197,28 +204,27 @@ depending on the location indicated by the surrounding SourceNote.
 The functions which add information to the map are `recordStgIdPosition` and
 `numberDataCon`.
 
-When the -fdistinct-constructor-tables` flag is turned on then every
+When the `-fdistinct-constructor-tables` flag is turned on then every
 usage of a data constructor gets its own distinct info table. This is orchestrated
 in `collectExpr` where an incrementing number is used to distinguish each
 occurrence of a data constructor.
 
--- StgToCmm
+-- GenerateCgIPEStub
+
+The info tables which are actually used in the generated program are collected after
+the Cmm pipeline. `initInfoTableProv` is used to create a CStub, that initializes the
+map in C code.
 
-The info tables which are actually used in the generated program are recorded during the
-conversion from STG to Cmm. The used info tables are recorded in the `emitProc` function.
-All the used info tables are recorded in the `cgs_used_info` field. This step
-is necessary because when the information about names is collected in the previous
-phase it's unpredictable about which names will end up needing info tables. If
-you don't record which ones are actually used then you end up generating code
-which references info tables which don't exist.
+This step has to be done after the Cmm pipeline to make sure that all info tables are
+really used and, even more importantly, return frame info tables are generated by the
+pipeline.
 
 -- Code Generation
 
 The output of these two phases is combined together during code generation.
-A C stub is generated which
-creates the static map from info table pointer to the information about where that
-info table was created from. This is created by `ipInitCode` in the same manner as a
-C stub is generated for cost centres.
+A C stub is generated which creates the static map from info table pointer to the
+information about where that info table was created from. This is created by
+`ipInitCode` in the same manner as a C stub is generated for cost centres.
 
 This information can be consumed in two ways.
 
diff --git a/compiler/GHC/StgToCmm.hs b/compiler/GHC/StgToCmm.hs
index 546c270f76..ee297e4220 100644
--- a/compiler/GHC/StgToCmm.hs
+++ b/compiler/GHC/StgToCmm.hs
@@ -18,7 +18,7 @@ import GHC.Prelude as Prelude
 import GHC.Driver.Backend
 import GHC.Driver.Session
 
-import GHC.StgToCmm.Prof (initInfoTableProv, initCostCentres, ldvEnter)
+import GHC.StgToCmm.Prof (initCostCentres, ldvEnter)
 import GHC.StgToCmm.Monad
 import GHC.StgToCmm.Env
 import GHC.StgToCmm.Bind
@@ -47,7 +47,6 @@ import GHC.Types.Basic
 import GHC.Types.Var.Set ( isEmptyDVarSet )
 import GHC.Types.Unique.FM
 import GHC.Types.Name.Env
-import GHC.Types.ForeignStubs
 
 import GHC.Core.DataCon
 import GHC.Core.TyCon
@@ -70,13 +69,8 @@ import Control.Monad (when,void, forM_)
 import GHC.Utils.Misc
 import System.IO.Unsafe
 import qualified Data.ByteString as BS
-import Data.Maybe
 import Data.IORef
 
-data CodeGenState = CodeGenState { codegen_used_info :: !(OrdList CmmInfoTable)
-                                 , codegen_state :: !CgState }
-
-
 codeGen :: Logger
         -> TmpFs
         -> DynFlags
@@ -86,34 +80,26 @@ codeGen :: Logger
         -> CollectedCCs                -- (Local/global) cost-centres needing declaring/registering.
         -> [CgStgTopBinding]           -- Bindings to convert
         -> HpcInfo
-        -> Stream IO CmmGroup (CStub, ModuleLFInfos)       -- Output as a stream, so codegen can
+        -> Stream IO CmmGroup ModuleLFInfos       -- Output as a stream, so codegen can
                                        -- be interleaved with output
 
-codeGen logger tmpfs dflags this_mod ip_map@(InfoTableProvMap (UniqMap denv) _) data_tycons
+codeGen logger tmpfs dflags this_mod (InfoTableProvMap (UniqMap denv) _ _) data_tycons
         cost_centre_info stg_binds hpc_info
   = do  {     -- cg: run the code generator, and yield the resulting CmmGroup
               -- Using an IORef to store the state is a bit crude, but otherwise
               -- we would need to add a state monad layer which regresses
               -- allocations by 0.5-2%.
-        ; cgref <- liftIO $ initC >>= \s -> newIORef (CodeGenState mempty s)
+        ; cgref <- liftIO $ initC >>= \s -> newIORef s
         ; let cg :: FCode a -> Stream IO CmmGroup a
               cg fcode = do
                 (a, cmm) <- liftIO . withTimingSilent logger (text "STG -> Cmm") (`seq` ()) $ do
-                         CodeGenState ts st <- readIORef cgref
+                         st <- readIORef cgref
                          let (a,st') = runC dflags this_mod st (getCmm fcode)
 
                          -- NB. stub-out cgs_tops and cgs_stmts.  This fixes
                          -- a big space leak.  DO NOT REMOVE!
                          -- This is observed by the #3294 test
-                         let !used_info
-                                | gopt Opt_InfoTableMap dflags = toOL (mapMaybe topInfoTable (snd a)) `mappend` ts
-                                | otherwise = mempty
-                         writeIORef cgref $!
-                                    CodeGenState used_info
-                                      (st'{ cgs_tops = nilOL,
-                                            cgs_stmts = mkNop
-                                          })
-
+                         writeIORef cgref $! (st'{ cgs_tops = nilOL, cgs_stmts = mkNop })
                          return a
                 yield cmm
                 return a
@@ -144,10 +130,7 @@ codeGen logger tmpfs dflags this_mod ip_map@(InfoTableProvMap (UniqMap denv) _)
         ; mapM_ (\(dc, ns) -> forM_ ns $ \(k, _ss) -> cg (cgDataCon (UsageSite this_mod k) dc)) (nonDetEltsUFM denv)
 
         ; final_state <- liftIO (readIORef cgref)
-        ; let cg_id_infos = cgs_binds . codegen_state $ final_state
-              used_info = fromOL . codegen_used_info $ final_state
-
-        ; !foreign_stub <- cg (initInfoTableProv used_info ip_map this_mod)
+        ; let cg_id_infos = cgs_binds final_state
 
           -- See Note [Conveying CAF-info and LFInfo between modules] in
           -- GHC.StgToCmm.Types
@@ -162,7 +145,7 @@ codeGen logger tmpfs dflags this_mod ip_map@(InfoTableProvMap (UniqMap denv) _)
                 | otherwise
                 = mkNameEnv (Prelude.map extractInfo (eltsUFM cg_id_infos))
 
-        ; return (foreign_stub, generatedInfo)
+        ; return generatedInfo
         }
 
 ---------------------------------------------------------------
diff --git a/compiler/GHC/StgToCmm/Prim.hs b/compiler/GHC/StgToCmm/Prim.hs
index 0f5943bf48..dff86341b1 100644
--- a/compiler/GHC/StgToCmm/Prim.hs
+++ b/compiler/GHC/StgToCmm/Prim.hs
@@ -1669,7 +1669,6 @@ emitPrimOp dflags primop = case primop of
   TraceEventBinaryOp -> alwaysExternal
   TraceMarkerOp -> alwaysExternal
   SetThreadAllocationCounter -> alwaysExternal
-  CloneMyStack -> alwaysExternal
 
   -- See Note [keepAlive# magic] in GHC.CoreToStg.Prep.
   KeepAliveOp -> panic "keepAlive# should have been eliminated in CorePrep"
diff --git a/compiler/GHC/StgToCmm/Prof.hs b/compiler/GHC/StgToCmm/Prof.hs
index 58fdfeafd9..852b77ef2b 100644
--- a/compiler/GHC/StgToCmm/Prof.hs
+++ b/compiler/GHC/StgToCmm/Prof.hs
@@ -284,8 +284,6 @@ initInfoTableProv infos itmap this_mod
   = do
        dflags <- getDynFlags
        let ents = convertInfoProvMap dflags infos this_mod itmap
-       --pprTraceM "UsedInfo" (ppr (length infos))
-       --pprTraceM "initInfoTable" (ppr (length ents))
        -- Output the actual IPE data
        mapM_ emitInfoTableProv ents
        -- Create the C stub which initialises the IPE map
diff --git a/compiler/GHC/StgToCmm/Utils.hs b/compiler/GHC/StgToCmm/Utils.hs
index 934fd6d726..8c6a40c69d 100644
--- a/compiler/GHC/StgToCmm/Utils.hs
+++ b/compiler/GHC/StgToCmm/Utils.hs
@@ -90,6 +90,7 @@ import GHC.Core.DataCon
 import GHC.Types.Unique.FM
 import GHC.Data.Maybe
 import Control.Monad
+import qualified Data.Map.Strict as Map
 
 --------------------------------------------------------------------------
 --
@@ -600,7 +601,7 @@ cmmInfoTableToInfoProvEnt this_mod cmit =
 -- | Convert source information collected about identifiers in 'GHC.STG.Debug'
 -- to entries suitable for placing into the info table provenenance table.
 convertInfoProvMap :: DynFlags -> [CmmInfoTable] -> Module -> InfoTableProvMap -> [InfoProvEnt]
-convertInfoProvMap dflags defns this_mod (InfoTableProvMap (UniqMap dcenv) denv) =
+convertInfoProvMap dflags defns this_mod (InfoTableProvMap (UniqMap dcenv) denv infoTableToSourceLocationMap) =
   map (\cmit ->
     let cl = cit_lbl cmit
         cn  = rtsClosureType (cit_rep cmit)
@@ -620,8 +621,16 @@ convertInfoProvMap dflags defns this_mod (InfoTableProvMap (UniqMap dcenv) denv)
             -- Lookup is linear but lists will be small (< 100)
             return $ InfoProvEnt cl cn (tyString (dataConTyCon dc)) this_mod (join $ lookup n (NE.toList ns))
 
+        lookupInfoTableToSourceLocation = do
+            sourceNote <- Map.lookup (cit_lbl cmit) infoTableToSourceLocationMap
+            return $ InfoProvEnt cl cn "" this_mod sourceNote
+
         -- This catches things like prim closure types and anything else which doesn't have a
         -- source location
         simpleFallback = cmmInfoTableToInfoProvEnt this_mod cmit
 
-    in fromMaybe simpleFallback (lookupDataConMap `firstJust` lookupClosureMap)) defns
+  in
+    if (isStackRep . cit_rep) cmit then
+      fromMaybe simpleFallback lookupInfoTableToSourceLocation
+    else
+      fromMaybe simpleFallback (lookupDataConMap `firstJust` lookupClosureMap)) defns
diff --git a/compiler/GHC/Types/IPE.hs b/compiler/GHC/Types/IPE.hs
index c69aeb004a..461bae6a55 100644
--- a/compiler/GHC/Types/IPE.hs
+++ b/compiler/GHC/Types/IPE.hs
@@ -1,5 +1,10 @@
-module GHC.Types.IPE(DCMap, ClosureMap, InfoTableProvMap(..)
-                    , emptyInfoTableProvMap) where
+module GHC.Types.IPE (
+    DCMap,
+    ClosureMap,
+    InfoTableProvMap(..),
+    emptyInfoTableProvMap,
+    IpeSourceLocation
+) where
 
 import GHC.Prelude
 
@@ -10,11 +15,17 @@ import GHC.Core.DataCon
 import GHC.Types.Unique.Map
 import GHC.Core.Type
 import Data.List.NonEmpty
+import GHC.Cmm.CLabel (CLabel)
+import qualified Data.Map.Strict as Map
+
+-- | Position and information about an info table.
+-- For return frames these are the contents of a 'CoreSyn.SourceNote'.
+type IpeSourceLocation = (RealSrcSpan, String)
 
 -- | A map from a 'Name' to the best approximate source position that
 -- name arose from.
 type ClosureMap = UniqMap Name  -- The binding
-                          (Type, Maybe (RealSrcSpan, String))
+                          (Type, Maybe IpeSourceLocation)
                           -- The best approximate source position.
                           -- (rendered type, source position, source note
                           -- label)
@@ -26,11 +37,15 @@ type ClosureMap = UniqMap Name  -- The binding
 -- the constructor was used at, if possible and a string which names
 -- the source location. This is the same information as is the payload
 -- for the 'GHC.Core.SourceNote' constructor.
-type DCMap = UniqMap DataCon (NonEmpty (Int, Maybe (RealSrcSpan, String)))
+type DCMap = UniqMap DataCon (NonEmpty (Int, Maybe IpeSourceLocation))
+
+type InfoTableToSourceLocationMap = Map.Map CLabel (Maybe IpeSourceLocation)
 
 data InfoTableProvMap = InfoTableProvMap
                           { provDC :: DCMap
-                          , provClosure :: ClosureMap }
+                          , provClosure :: ClosureMap
+                          , provInfoTables :: InfoTableToSourceLocationMap
+                          }
 
 emptyInfoTableProvMap :: InfoTableProvMap
-emptyInfoTableProvMap = InfoTableProvMap emptyUniqMap emptyUniqMap
+emptyInfoTableProvMap = InfoTableProvMap emptyUniqMap emptyUniqMap Map.empty
diff --git a/compiler/ghc.cabal.in b/compiler/ghc.cabal.in
index b1acf005da..25833017d1 100644
--- a/compiler/ghc.cabal.in
+++ b/compiler/ghc.cabal.in
@@ -426,6 +426,7 @@ Library
         GHC.Driver.Errors.Ppr
         GHC.Driver.Errors.Types
         GHC.Driver.Flags
+        GHC.Driver.GenerateCgIPEStub
         GHC.Driver.Hooks
         GHC.Driver.Main
         GHC.Driver.Make