Combine STG free variable traversals (#17978)

Previously we would traverse the STG AST twice looking for free variables.
 * Once in `annTopBindingsDeps` which considers top level and imported ids free.
   Its output is used to put bindings in dependency order. The pass happens
   in STG pipeline.
 * Once in `annTopBindingsFreeVars` which only considers non-top level ids free.
   Its output is used by the code generator to compute offsets into closures.
   This happens in Cmm (CodeGen) pipeline.

Now these two traversal operations are merged into one - `FVs.depSortWithAnnotStgPgm`.
The pass happens right at the end of STG pipeline. Some type signatures had to be
updated due to slight shifts of StgPass boundaries (for example, top-level CodeGen
handler now directly works with CodeGen flavoured Stg AST instead of Vanilla).

Due to changed order of bindings, a few debugger type reconstruction bugs
have resurfaced again (see tests break018, break021) - work item #18004 tracks this
investigation.

authors: simonpj, nineonine

13 files changed, 288 insertions, 319 deletions

diff --git a/compiler/GHC/Driver/Main.hs b/compiler/GHC/Driver/Main.hs
index ffc0cef89d..c403b3e85a 100644
--- a/compiler/GHC/Driver/Main.hs
+++ b/compiler/GHC/Driver/Main.hs
@@ -160,7 +160,6 @@ import GHC.Tc.Utils.Monad
 import GHC.Tc.Utils.Zonk    ( ZonkFlexi (DefaultFlexi) )
 
 import GHC.Stg.Syntax
-import GHC.Stg.FVs      ( annTopBindingsFreeVars )
 import GHC.Stg.Pipeline ( stg2stg )
 
 import GHC.Builtin.Utils
@@ -1767,24 +1766,22 @@ This reduces residency towards the end of the CodeGen phase significantly
 (5-10%).
 -}
 
-doCodeGen   :: HscEnv -> Module -> InfoTableProvMap -> [TyCon]
-            -> CollectedCCs
-            -> [StgTopBinding]
-            -> HpcInfo
-            -> IO (Stream IO CmmGroupSRTs CgInfos)
+doCodeGen :: HscEnv -> Module -> InfoTableProvMap -> [TyCon]
+          -> CollectedCCs
+          -> [CgStgTopBinding] -- ^ Bindings come already annotated with fvs
+          -> HpcInfo
+          -> IO (Stream IO CmmGroupSRTs CgInfos)
          -- Note we produce a 'Stream' of CmmGroups, so that the
          -- backend can be run incrementally.  Otherwise it generates all
          -- the C-- up front, which has a significant space cost.
 doCodeGen hsc_env this_mod denv data_tycons
-              cost_centre_info stg_binds hpc_info = do
+              cost_centre_info stg_binds_w_fvs hpc_info = do
     let dflags = hsc_dflags hsc_env
     let logger = hsc_logger hsc_env
     let hooks  = hsc_hooks hsc_env
     let tmpfs  = hsc_tmpfs hsc_env
     let platform = targetPlatform dflags
 
-    let stg_binds_w_fvs = annTopBindingsFreeVars stg_binds
-
     putDumpFileMaybe logger Opt_D_dump_stg_final "Final STG:" FormatSTG (pprGenStgTopBindings (initStgPprOpts dflags) stg_binds_w_fvs)
 
     let stg_to_cmm = case stgToCmmHook hooks of
@@ -1829,7 +1826,7 @@ myCoreToStgExpr :: Logger -> DynFlags -> InteractiveContext
                 -> Bool
                 -> Module -> ModLocation -> CoreExpr
                 -> IO ( Id
-                      , [StgTopBinding]
+                      , [CgStgTopBinding]
                       , InfoTableProvMap
                       , CollectedCCs )
 myCoreToStgExpr logger dflags ictxt for_bytecode this_mod ml prepd_expr = do
@@ -1852,7 +1849,7 @@ myCoreToStgExpr logger dflags ictxt for_bytecode this_mod ml prepd_expr = do
 myCoreToStg :: Logger -> DynFlags -> InteractiveContext
             -> Bool
             -> Module -> ModLocation -> CoreProgram
-            -> IO ( [StgTopBinding] -- output program
+            -> IO ( [CgStgTopBinding] -- output program
                   , InfoTableProvMap
                   , CollectedCCs )  -- CAF cost centre info (declared and used)
 myCoreToStg logger dflags ictxt for_bytecode this_mod ml prepd_binds = do
@@ -1860,11 +1857,11 @@ myCoreToStg logger dflags ictxt for_bytecode this_mod ml prepd_binds = do
          = {-# SCC "Core2Stg" #-}
            coreToStg dflags this_mod ml prepd_binds
 
-    stg_binds2
+    stg_binds_with_fvs
         <- {-# SCC "Stg2Stg" #-}
            stg2stg logger dflags ictxt for_bytecode this_mod stg_binds
 
-    return (stg_binds2, denv, cost_centre_info)
+    return (stg_binds_with_fvs, denv, cost_centre_info)
 
 {- **********************************************************************
 %*                                                                      *
diff --git a/compiler/GHC/Stg/DepAnal.hs b/compiler/GHC/Stg/DepAnal.hs
deleted file mode 100644
index 74b490969a..0000000000
--- a/compiler/GHC/Stg/DepAnal.hs
+++ /dev/null
@@ -1,146 +0,0 @@
-
-
-module GHC.Stg.DepAnal (depSortStgPgm) where
-
-import GHC.Prelude
-
-import GHC.Stg.Syntax
-import GHC.Types.Id
-import GHC.Types.Name (Name, nameIsLocalOrFrom)
-import GHC.Types.Name.Env
-import GHC.Utils.Outputable
-import GHC.Utils.Panic
-import GHC.Types.Unique.Set (nonDetEltsUniqSet)
-import GHC.Types.Var.Set
-import GHC.Unit.Module (Module)
-
-import Data.Graph (SCC (..))
-import Data.Bifunctor (first)
-
---------------------------------------------------------------------------------
--- * Dependency analysis
-
--- | Set of bound variables
-type BVs = VarSet
-
--- | Set of free variables
-type FVs = VarSet
-
--- | Dependency analysis on STG terms.
---
--- Dependencies of a binding are just free variables in the binding. This
--- includes imported ids and ids in the current module. For recursive groups we
--- just return one set of free variables which is just the union of dependencies
--- of all bindings in the group.
---
--- Implementation: pass bound variables (BVs) to recursive calls, get free
--- variables (FVs) back. We ignore imported FVs as they do not change the
--- ordering but it improves performance.
---
-annTopBindingsDeps :: Module -> [StgTopBinding] -> [(StgTopBinding, FVs)]
-annTopBindingsDeps this_mod bs = zip bs (map top_bind bs)
-  where
-    top_bind :: StgTopBinding -> FVs
-    top_bind StgTopStringLit{} =
-      emptyVarSet
-
-    top_bind (StgTopLifted bs) =
-      binding emptyVarSet bs
-
-    binding :: BVs -> StgBinding -> FVs
-    binding bounds (StgNonRec _ r) =
-      rhs bounds r
-    binding bounds (StgRec bndrs) =
-      unionVarSets $
-        map (bind_non_rec (extendVarSetList bounds (map fst bndrs))) bndrs
-
-    bind_non_rec :: BVs -> (Id, StgRhs) -> FVs
-    bind_non_rec bounds (_, r) =
-        rhs bounds r
-
-    rhs :: BVs -> StgRhs -> FVs
-    rhs bounds (StgRhsClosure _ _ _ as e) =
-      expr (extendVarSetList bounds as) e
-
-    rhs bounds (StgRhsCon _ _ _ _ as) =
-      args bounds as
-
-    var :: BVs -> Var -> FVs
-    var bounds v
-      | not (elemVarSet v bounds)
-      , nameIsLocalOrFrom this_mod (idName v)
-      = unitVarSet v
-      | otherwise
-      = emptyVarSet
-
-    arg :: BVs -> StgArg -> FVs
-    arg bounds (StgVarArg v) = var bounds v
-    arg _ StgLitArg{} = emptyVarSet
-
-    args :: BVs -> [StgArg] -> FVs
-    args bounds as = unionVarSets (map (arg bounds) as)
-
-    expr :: BVs -> StgExpr -> FVs
-    expr bounds (StgApp f as) =
-      var bounds f `unionVarSet` args bounds as
-
-    expr _ StgLit{} =
-      emptyVarSet
-
-    expr bounds (StgConApp _ _ as _) =
-      args bounds as
-    expr bounds (StgOpApp _ as _) =
-      args bounds as
-    expr bounds (StgCase scrut scrut_bndr _ as) =
-      expr bounds scrut `unionVarSet`
-        alts (extendVarSet bounds scrut_bndr) as
-    expr bounds (StgLet _ bs e) =
-      binding bounds bs `unionVarSet`
-        expr (extendVarSetList bounds (bindersOf bs)) e
-    expr bounds (StgLetNoEscape _ bs e) =
-      binding bounds bs `unionVarSet`
-        expr (extendVarSetList bounds (bindersOf bs)) e
-
-    expr bounds (StgTick _ e) =
-      expr bounds e
-
-    alts :: BVs -> [StgAlt] -> FVs
-    alts bounds = unionVarSets . map (alt bounds)
-
-    alt :: BVs -> StgAlt -> FVs
-    alt bounds (_, bndrs, e) =
-      expr (extendVarSetList bounds bndrs) e
-
---------------------------------------------------------------------------------
--- * Dependency sorting
-
--- | Dependency sort a STG program so that dependencies come before uses.
-depSortStgPgm :: Module -> [StgTopBinding] -> [StgTopBinding]
-depSortStgPgm this_mod =
-    {-# SCC "STG.depSort" #-}
-    map fst . depSort . annTopBindingsDeps this_mod
-
--- | Sort free-variable-annotated STG bindings so that dependencies come before
--- uses.
-depSort :: [(StgTopBinding, FVs)] -> [(StgTopBinding, FVs)]
-depSort = concatMap get_binds . depAnal defs uses
-  where
-    uses, defs :: (StgTopBinding, FVs) -> [Name]
-
-    -- TODO (osa): I'm unhappy about two things in this code:
-    --
-    --     * Why do we need Name instead of Id for uses and dependencies?
-    --     * Why do we need a [Name] instead of `Set Name`? Surely depAnal
-    --       doesn't need any ordering.
-
-    uses (StgTopStringLit{}, _) = []
-    uses (StgTopLifted{}, fvs)  = map idName (nonDetEltsUniqSet fvs)
-
-    defs (bind, _) = map idName (bindersOfTop bind)
-
-    get_binds (AcyclicSCC bind) =
-      [bind]
-    get_binds (CyclicSCC binds) =
-      pprPanic "depSortStgBinds"
-               (text "Found cyclic SCC:"
-               $$ ppr (map (first (pprStgTopBinding panicStgPprOpts)) binds))
diff --git a/compiler/GHC/Stg/FVs.hs b/compiler/GHC/Stg/FVs.hs
index 62053001a6..00b4bdcc0b 100644
--- a/compiler/GHC/Stg/FVs.hs
+++ b/compiler/GHC/Stg/FVs.hs
@@ -40,129 +40,236 @@ Top-level closure bindings never capture variables as all of their free
 variables are global.
 -}
 module GHC.Stg.FVs (
-    annTopBindingsFreeVars,
+    depSortWithAnnotStgPgm,
     annBindingFreeVars
   ) where
 
-import GHC.Prelude
+import GHC.Prelude hiding (mod)
 
 import GHC.Stg.Syntax
 import GHC.Types.Id
-import GHC.Types.Var.Set
+import GHC.Types.Name (Name, nameIsLocalOrFrom)
 import GHC.Types.Tickish ( GenTickish(Breakpoint) )
+import GHC.Types.Unique.Set (nonDetEltsUniqSet)
+import GHC.Types.Var.Set
+import GHC.Unit.Module (Module)
 import GHC.Utils.Misc
 
-import Data.Maybe ( mapMaybe )
+import Data.Graph (SCC (..))
+import GHC.Data.Graph.Directed( Node(..), stronglyConnCompFromEdgedVerticesUniq )
+
+{- Note [Why do we need dependency analysis?]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The program needs to be in dependency order for the SRT algorithm to
+work (see CmmBuildInfoTables, which also includes a detailed
+description of the algorithm).
+
+But isn't it in correct dependency order already?  No:
+
+* The simplifier does not guarantee to produce programs in dependency
+  order (see #16192 and Note [Glomming] in GHC.Core.Opt.OccurAnal).
+  This could be solved by a final run of the occurrence analyser, but
+  that's more work
+
+* We also don't guarantee that StgLiftLams will preserve the order or
+  only create minimal recursive groups.
+-}
+
+--------------------------------------------------------------------------------
+-- | Dependency sort a STG program, and annotate it with free variables
+-- The returned bindings:
+--   * Are in dependency order
+--   * Each StgRhsClosure is correctly annotated (in its extension field)
+--     with the free variables needed in the closure
+--   * Each StgCase is correctly annotated (in its extension field) with
+--     the variables that must be saved across the case
+depSortWithAnnotStgPgm :: Module -> [StgTopBinding] -> [CgStgTopBinding]
+depSortWithAnnotStgPgm this_mod binds
+  = {-# SCC "STG.depSortWithAnnotStgPgm" #-}
+    lit_binds ++ map from_scc sccs
+  where
+    lit_binds :: [CgStgTopBinding]
+    pairs     :: [(Id, StgRhs)]
+    (lit_binds, pairs) = flattenTopStgBindings binds
+
+    nodes :: [Node Name (Id, CgStgRhs)]
+    nodes = map (annotateTopPair env0) pairs
+    env0 = Env { locals = emptyVarSet, mod = this_mod }
+
+    -- Do strongly connected component analysis.  Why?
+    -- See Note [Why do we need dependency analysis?]
+    sccs :: [SCC (Id,CgStgRhs)]
+    sccs  = stronglyConnCompFromEdgedVerticesUniq nodes
+
+    from_scc (CyclicSCC pairs)       = StgTopLifted (StgRec pairs)
+    from_scc (AcyclicSCC (bndr,rhs)) = StgTopLifted (StgNonRec bndr rhs)
 
-newtype Env
+
+flattenTopStgBindings :: [StgTopBinding] -> ([CgStgTopBinding], [(Id,StgRhs)])
+flattenTopStgBindings binds
+  = go [] [] binds
+  where
+    go lits pairs [] = (lits, pairs)
+    go lits pairs (bind:binds)
+      = case bind of
+          StgTopStringLit bndr rhs -> go (StgTopStringLit bndr rhs:lits) pairs binds
+          StgTopLifted stg_bind -> go lits (flatten_one stg_bind ++ pairs) binds
+
+    flatten_one (StgNonRec b r) = [(b,r)]
+    flatten_one (StgRec pairs)  = pairs
+
+annotateTopPair :: Env -> (Id, StgRhs) -> Node Name (Id, CgStgRhs)
+annotateTopPair env0 (bndr, rhs)
+  = DigraphNode { node_key          = idName bndr
+                , node_payload      = (bndr, rhs')
+                , node_dependencies = map idName (nonDetEltsUniqSet top_fvs) }
+  where
+    (rhs', top_fvs, _) = rhsFVs env0 rhs
+
+--------------------------------------------------------------------------------
+-- * Non-global free variable analysis
+
+data Env
   = Env
-  { locals :: IdSet
+  { -- | Set of locally-bound, not-top-level binders in scope.
+    -- That is, variables bound by a let (but not let-no-escape), a lambda
+    -- (in a StgRhsClsoure), a case binder, or a case alternative.  These
+    -- are the variables that must be captured in a function closure, if they
+    -- are free in the RHS. Example
+    --   f = \x. let g = \y. x+1
+    --           let h = \z. g z + 1
+    --           in h x
+    -- In the body of h we have locals = {x, g, z}.  Note that f is top level
+    -- and does not appear in locals.
+    locals :: IdSet
+  , mod    :: Module
   }
 
-emptyEnv :: Env
-emptyEnv = Env emptyVarSet
-
 addLocals :: [Id] -> Env -> Env
 addLocals bndrs env
   = env { locals = extendVarSetList (locals env) bndrs }
 
--- | Annotates a top-level STG binding group with its free variables.
-annTopBindingsFreeVars :: [StgTopBinding] -> [CgStgTopBinding]
-annTopBindingsFreeVars = map go
-  where
-    go (StgTopStringLit id bs) = StgTopStringLit id bs
-    go (StgTopLifted bind)
-      = StgTopLifted (annBindingFreeVars bind)
-
--- | Annotates an STG binding with its free variables.
-annBindingFreeVars :: StgBinding -> CgStgBinding
-annBindingFreeVars = fst . binding emptyEnv emptyDVarSet
-
-boundIds :: StgBinding -> [Id]
-boundIds (StgNonRec b _) = [b]
-boundIds (StgRec pairs)  = map fst pairs
-
--- Note [Tracking local binders]
--- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
--- 'locals' contains non-toplevel, non-imported binders.
--- We maintain the set in 'expr', 'alt' and 'rhs', which are the only
--- places where new local binders are introduced.
--- Why do it there rather than in 'binding'? Two reasons:
+--------------------------------------------------------------------------------
+-- | TopFVs: set of variables that are:
+--    (a) bound at the top level of this module, and
+--    (b) appear free in the expression
+-- It is a /non-deterministic/ set because we use it only to perform dependency
+-- analysis on the top-level bindings.
+type TopFVs   = IdSet
+
+-- | LocalFVs: set of variable that are:
+--     (a) bound locally (by a lambda, non-top-level let, or case); that is,
+--         it appears in the 'locals' field of 'Env'
+--     (b) appear free in the expression
+-- It is a /deterministic/ set because it is used to annotate closures with
+-- their free variables, and we want closure layout to be deterministic.
 --
---   1. We call 'binding' from 'annTopBindingsFreeVars', which would
---      add top-level bindings to the 'locals' set.
---   2. In the let(-no-escape) case, we need to extend the environment
---      prior to analysing the body, but we also need the fvs from the
---      body to analyse the RHSs. No way to do this without some
---      knot-tying.
-
--- | This makes sure that only local, non-global free vars make it into the set.
-mkFreeVarSet :: Env -> [Id] -> DIdSet
-mkFreeVarSet env = mkDVarSet . filter (`elemVarSet` locals env)
-
-args :: Env -> [StgArg] -> DIdSet
-args env = mkFreeVarSet env . mapMaybe f
-  where
-    f (StgVarArg occ) = Just occ
-    f _               = Nothing
+-- Invariant: the LocalFVs returned is a subset of the 'locals' field of Env
+type LocalFVs = DIdSet
 
-binding :: Env -> DIdSet -> StgBinding -> (CgStgBinding, DIdSet)
-binding env body_fv (StgNonRec bndr r) = (StgNonRec bndr r', fvs)
-  where
-    -- See Note [Tracking local binders]
-    (r', rhs_fvs) = rhs env r
-    fvs = delDVarSet body_fv bndr `unionDVarSet` rhs_fvs
-binding env body_fv (StgRec pairs) = (StgRec pairs', fvs)
-  where
-    -- See Note [Tracking local binders]
-    bndrs = map fst pairs
-    (rhss, rhs_fvss) = mapAndUnzip (rhs env . snd) pairs
-    pairs' = zip bndrs rhss
-    fvs = delDVarSetList (unionDVarSets (body_fv:rhs_fvss)) bndrs
-
-expr :: Env -> StgExpr -> (CgStgExpr, DIdSet)
-expr env = go
-  where
-    go (StgApp occ as)
-      = (StgApp occ as, unionDVarSet (args env as) (mkFreeVarSet env [occ]))
-    go (StgLit lit) = (StgLit lit, emptyDVarSet)
-    go (StgConApp dc n as tys) = (StgConApp dc n as tys, args env as)
-    go (StgOpApp op as ty) = (StgOpApp op as ty, args env as)
-    go (StgCase scrut bndr ty alts) = (StgCase scrut' bndr ty alts', fvs)
-      where
-        (scrut', scrut_fvs) = go scrut
-        -- See Note [Tracking local binders]
-        (alts', alt_fvss) = mapAndUnzip (alt (addLocals [bndr] env)) alts
-        alt_fvs = unionDVarSets alt_fvss
-        fvs = delDVarSet (unionDVarSet scrut_fvs alt_fvs) bndr
-    go (StgLet ext bind body) = go_bind (StgLet ext) bind body
-    go (StgLetNoEscape ext bind body) = go_bind (StgLetNoEscape ext) bind body
-    go (StgTick tick e) = (StgTick tick e', fvs')
+-- | Dependency analysis on STG terms.
+--
+-- Dependencies of a binding are just free variables in the binding. This
+-- includes imported ids and ids in the current module. For recursive groups we
+-- just return one set of free variables which is just the union of dependencies
+-- of all bindings in the group.
+--
+-- Implementation: pass bound variables (NestedIds) to recursive calls, get free
+-- variables (TopFVs) back. We ignore imported TopFVs as they do not change the
+-- ordering but it improves performance (see `nameIsExternalFrom` call in `vars_fvs`).
+--
+
+annBindingFreeVars :: Module -> StgBinding -> CgStgBinding
+annBindingFreeVars this_mod = fstOf3 . bindingFVs (Env emptyVarSet this_mod) emptyDVarSet
+
+bindingFVs :: Env -> LocalFVs -> StgBinding -> (CgStgBinding, TopFVs, LocalFVs)
+bindingFVs env body_fv b =
+  case b of
+    StgNonRec bndr r -> (StgNonRec bndr r', fvs, lcl_fvs)
       where
-        (e', fvs) = go e
-        fvs' = unionDVarSet (tickish tick) fvs
-        tickish (Breakpoint _ _ ids) = mkDVarSet ids
-        tickish _                    = emptyDVarSet
+        (r', fvs, rhs_lcl_fvs) = rhsFVs env r
+        lcl_fvs = delDVarSet body_fv bndr `unionDVarSet` rhs_lcl_fvs
 
-    go_bind dc bind body = (dc bind' body', fvs)
+    StgRec pairs -> (StgRec pairs', fvs, lcl_fvss)
       where
-        -- See Note [Tracking local binders]
-        env' = addLocals (boundIds bind) env
-        (body', body_fvs) = expr env' body
-        (bind', fvs) = binding env' body_fvs bind
-
-rhs :: Env -> StgRhs -> (CgStgRhs, DIdSet)
-rhs env (StgRhsClosure _ ccs uf bndrs body)
-  = (StgRhsClosure fvs ccs uf bndrs body', fvs)
+        bndrs = map fst pairs
+        env' = addLocals bndrs env
+        (rhss, rhs_fvss, rhs_lcl_fvss) = mapAndUnzip3 (rhsFVs env' . snd) pairs
+        fvs = unionVarSets rhs_fvss
+        pairs' = zip bndrs rhss
+        lcl_fvss = delDVarSetList (unionDVarSets (body_fv:rhs_lcl_fvss)) bndrs
+
+varFVs :: Env -> Id -> (TopFVs, LocalFVs) -> (TopFVs, LocalFVs)
+varFVs env v (top_fvs, lcl_fvs)
+  | v `elemVarSet` locals env                -- v is locally bound
+  = (top_fvs, lcl_fvs `extendDVarSet` v)
+  | nameIsLocalOrFrom (mod env) (idName v)   -- v is bound at top level
+  = (top_fvs `extendVarSet` v, lcl_fvs)
+  | otherwise                                -- v is imported
+  = (top_fvs, lcl_fvs)
+
+exprFVs :: Env -> StgExpr -> (CgStgExpr, TopFVs, LocalFVs)
+exprFVs env = go
   where
-    -- See Note [Tracking local binders]
-    (body', body_fvs) = expr (addLocals bndrs env) body
-    fvs = delDVarSetList body_fvs bndrs
-rhs env (StgRhsCon ccs dc mu ts as) = (StgRhsCon ccs dc mu ts as, args env as)
+    go (StgApp f as)
+      | (top_fvs, lcl_fvs) <- varFVs env f (argsFVs env as)
+      = (StgApp f as, top_fvs, lcl_fvs)
+
+    go (StgLit lit) = (StgLit lit, emptyVarSet, emptyDVarSet)
 
-alt :: Env -> StgAlt -> (CgStgAlt, DIdSet)
-alt env (con, bndrs, e) = ((con, bndrs, e'), fvs)
+    go (StgConApp dc n as tys)
+      | (top_fvs, lcl_fvs) <- argsFVs env as
+      = (StgConApp dc n as tys, top_fvs, lcl_fvs)
+
+    go (StgOpApp op as ty)
+      | (top_fvs, lcl_fvs) <- argsFVs env as
+      = (StgOpApp op as ty, top_fvs, lcl_fvs)
+
+    go (StgCase scrut bndr ty alts)
+      | (scrut',scrut_top_fvs,scrut_lcl_fvs) <- exprFVs env scrut
+      , (alts',alts_top_fvss,alts_lcl_fvss)
+          <- mapAndUnzip3 (altFVs (addLocals [bndr] env)) alts
+      , let top_fvs = scrut_top_fvs `unionVarSet` unionVarSets alts_top_fvss
+            alts_lcl_fvs = unionDVarSets alts_lcl_fvss
+            lcl_fvs = delDVarSet (unionDVarSet scrut_lcl_fvs alts_lcl_fvs) bndr
+      = (StgCase scrut' bndr ty alts', top_fvs,lcl_fvs)
+
+    go (StgLet ext         bind body) = go_bind (StgLet ext) bind body
+    go (StgLetNoEscape ext bind body) = go_bind (StgLetNoEscape ext) bind body
+
+    go (StgTick tick e)
+      | (e', top_fvs, lcl_fvs) <- exprFVs env e
+      , let lcl_fvs' = unionDVarSet (tickish tick) lcl_fvs
+      = (StgTick tick e', top_fvs, lcl_fvs')
+        where
+          tickish (Breakpoint _ _ ids) = mkDVarSet ids
+          tickish _                    = emptyDVarSet
+
+    go_bind dc bind body = (dc bind' body', top_fvs, lcl_fvs)
+      where
+        env' = addLocals (bindersOf bind) env
+        (body', body_top_fvs, body_lcl_fvs) = exprFVs env' body
+        (bind', bind_top_fvs, lcl_fvs)      = bindingFVs env' body_lcl_fvs bind
+        top_fvs = bind_top_fvs `unionVarSet` body_top_fvs
+
+
+rhsFVs :: Env -> StgRhs -> (CgStgRhs, TopFVs, LocalFVs)
+rhsFVs env (StgRhsClosure _ ccs uf bs body)
+  | (body', top_fvs, lcl_fvs) <- exprFVs (addLocals bs env) body
+  , let lcl_fvs' = delDVarSetList lcl_fvs bs
+  = (StgRhsClosure lcl_fvs' ccs uf bs body', top_fvs, lcl_fvs')
+rhsFVs env (StgRhsCon ccs dc mu ts bs)
+  | (top_fvs, lcl_fvs) <- argsFVs env bs
+  = (StgRhsCon ccs dc mu ts bs, top_fvs, lcl_fvs)
+
+argsFVs :: Env -> [StgArg] -> (TopFVs, LocalFVs)
+argsFVs env = foldl' f (emptyVarSet, emptyDVarSet)
   where
-    -- See Note [Tracking local binders]
-    (e', rhs_fvs) = expr (addLocals bndrs env) e
-    fvs = delDVarSetList rhs_fvs bndrs
+    f (fvs,ids) StgLitArg{}   = (fvs, ids)
+    f (fvs,ids) (StgVarArg v) = varFVs env v (fvs, ids)
+
+altFVs :: Env -> StgAlt -> (CgStgAlt, TopFVs, LocalFVs)
+altFVs env (con,bndrs,e)
+  | (e', top_fvs, lcl_fvs) <- exprFVs (addLocals bndrs env) e
+  , let lcl_fvs' = delDVarSetList lcl_fvs bndrs
+  = ((con,bndrs, e'), top_fvs, lcl_fvs')
diff --git a/compiler/GHC/Stg/Lift.hs b/compiler/GHC/Stg/Lift.hs
index f83ccd388f..b4c473b50d 100644
--- a/compiler/GHC/Stg/Lift.hs
+++ b/compiler/GHC/Stg/Lift.hs
@@ -24,8 +24,9 @@ import GHC.Stg.FVs ( annBindingFreeVars )
 import GHC.Stg.Lift.Analysis
 import GHC.Stg.Lift.Monad
 import GHC.Stg.Syntax
-import GHC.Utils.Outputable
+import GHC.Unit.Module (Module)
 import GHC.Types.Unique.Supply
+import GHC.Utils.Outputable
 import GHC.Utils.Panic
 import GHC.Types.Var.Set
 import Control.Monad ( when )
@@ -124,17 +125,17 @@ import Data.Maybe ( isNothing )
 --
 -- (Mostly) textbook instance of the lambda lifting transformation, selecting
 -- which bindings to lambda lift by consulting 'goodToLift'.
-stgLiftLams :: DynFlags -> UniqSupply -> [InStgTopBinding] -> [OutStgTopBinding]
-stgLiftLams dflags us = runLiftM dflags us . foldr liftTopLvl (pure ())
+stgLiftLams :: Module -> DynFlags -> UniqSupply -> [InStgTopBinding] -> [OutStgTopBinding]
+stgLiftLams this_mod dflags us = runLiftM dflags us . foldr (liftTopLvl this_mod) (pure ())
 
-liftTopLvl :: InStgTopBinding -> LiftM () -> LiftM ()
-liftTopLvl (StgTopStringLit bndr lit) rest = withSubstBndr bndr $ \bndr' -> do
+liftTopLvl :: Module -> InStgTopBinding -> LiftM () -> LiftM ()
+liftTopLvl _ (StgTopStringLit bndr lit) rest = withSubstBndr bndr $ \bndr' -> do
   addTopStringLit bndr' lit
   rest
-liftTopLvl (StgTopLifted bind) rest = do
+liftTopLvl this_mod (StgTopLifted bind) rest = do
   let is_rec = isRec $ fst $ decomposeStgBinding bind
   when is_rec startBindingGroup
-  let bind_w_fvs = annBindingFreeVars bind
+  let bind_w_fvs = annBindingFreeVars this_mod bind
   withLiftedBind TopLevel (tagSkeletonTopBind bind_w_fvs) NilSk $ \mb_bind' -> do
     -- We signal lifting of a binding through returning Nothing.
     -- Should never happen for a top-level binding, though, since we are already
diff --git a/compiler/GHC/Stg/Pipeline.hs b/compiler/GHC/Stg/Pipeline.hs
index b0e1848f19..afd27cb67f 100644
--- a/compiler/GHC/Stg/Pipeline.hs
+++ b/compiler/GHC/Stg/Pipeline.hs
@@ -17,7 +17,7 @@ import GHC.Stg.Syntax
 
 import GHC.Stg.Lint     ( lintStgTopBindings )
 import GHC.Stg.Stats    ( showStgStats )
-import GHC.Stg.DepAnal  ( depSortStgPgm )
+import GHC.Stg.FVs      ( depSortWithAnnotStgPgm )
 import GHC.Stg.Unarise  ( unarise )
 import GHC.Stg.BcPrep   ( bcPrep )
 import GHC.Stg.CSE      ( stgCse )
@@ -52,13 +52,13 @@ stg2stg :: Logger
         -> Bool                      -- prepare for bytecode?
         -> Module                    -- module being compiled
         -> [StgTopBinding]           -- input program
-        -> IO [StgTopBinding]        -- output program
+        -> IO [CgStgTopBinding]        -- output program
 stg2stg logger dflags ictxt for_bytecode this_mod binds
   = do  { dump_when Opt_D_dump_stg_from_core "Initial STG:" binds
         ; showPass logger "Stg2Stg"
         -- Do the main business!
         ; binds' <- runStgM 'g' $
-            foldM do_stg_pass binds (getStgToDo for_bytecode dflags)
+            foldM (do_stg_pass this_mod) binds (getStgToDo for_bytecode dflags)
 
           -- Dependency sort the program as last thing. The program needs to be
           -- in dependency order for the SRT algorithm to work (see
@@ -68,8 +68,10 @@ stg2stg logger dflags ictxt for_bytecode this_mod binds
           -- dependency order. We also don't guarantee that StgLiftLams will
           -- preserve the order or only create minimal recursive groups, so a
           -- sorting pass is necessary.
-        ; let binds_sorted = depSortStgPgm this_mod binds'
-        ; return binds_sorted
+          -- This pass will also augment each closure with non-global free variables
+          -- annotations (which is used by code generator to compute offsets into closures)
+        ; let binds_sorted_with_fvs = depSortWithAnnotStgPgm this_mod binds'
+        ; return binds_sorted_with_fvs
    }
 
   where
@@ -80,8 +82,8 @@ stg2stg logger dflags ictxt for_bytecode this_mod binds
       = \ _whodunnit _binds -> return ()
 
     -------------------------------------------
-    do_stg_pass :: [StgTopBinding] -> StgToDo -> StgM [StgTopBinding]
-    do_stg_pass binds to_do
+    do_stg_pass :: Module -> [StgTopBinding] -> StgToDo -> StgM [StgTopBinding]
+    do_stg_pass this_mod binds to_do
       = case to_do of
           StgDoNothing ->
             return binds
@@ -95,7 +97,8 @@ stg2stg logger dflags ictxt for_bytecode this_mod binds
 
           StgLiftLams -> do
             us <- getUniqueSupplyM
-            let binds' = {-# SCC "StgLiftLams" #-} stgLiftLams dflags us binds
+            --
+            let binds' = {-# SCC "StgLiftLams" #-} stgLiftLams this_mod dflags us binds
             end_pass "StgLiftLams" binds'
 
           StgBcPrep -> do
diff --git a/compiler/GHC/Stg/Syntax.hs b/compiler/GHC/Stg/Syntax.hs
index 910101eceb..cc93157126 100644
--- a/compiler/GHC/Stg/Syntax.hs
+++ b/compiler/GHC/Stg/Syntax.hs
@@ -718,13 +718,13 @@ pprGenStgTopBindings :: (OutputablePass pass) => StgPprOpts -> [GenStgTopBinding
 pprGenStgTopBindings opts binds
   = vcat $ intersperse blankLine (map (pprGenStgTopBinding opts) binds)
 
-pprStgBinding :: StgPprOpts -> StgBinding -> SDoc
+pprStgBinding :: OutputablePass pass => StgPprOpts -> GenStgBinding pass -> SDoc
 pprStgBinding = pprGenStgBinding
 
-pprStgTopBinding :: StgPprOpts -> StgTopBinding -> SDoc
+pprStgTopBinding :: OutputablePass pass => StgPprOpts -> GenStgTopBinding pass -> SDoc
 pprStgTopBinding = pprGenStgTopBinding
 
-pprStgTopBindings :: StgPprOpts -> [StgTopBinding] -> SDoc
+pprStgTopBindings :: OutputablePass pass => StgPprOpts -> [GenStgTopBinding pass] -> SDoc
 pprStgTopBindings = pprGenStgTopBindings
 
 instance Outputable StgArg where
diff --git a/compiler/GHC/StgToByteCode.hs b/compiler/GHC/StgToByteCode.hs
index e7d4df472d..3e027d07b5 100644
--- a/compiler/GHC/StgToByteCode.hs
+++ b/compiler/GHC/StgToByteCode.hs
@@ -2,6 +2,7 @@
 {-# LANGUAGE DeriveFunctor              #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE RecordWildCards            #-}
+{-# LANGUAGE FlexibleContexts           #-}
 
 {-# OPTIONS_GHC -fprof-auto-top #-}
 {-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
@@ -88,14 +89,13 @@ import GHC.Stack.CCS
 import Data.Either ( partitionEithers )
 
 import GHC.Stg.Syntax
-import GHC.Stg.FVs
 
 -- -----------------------------------------------------------------------------
 -- Generating byte code for a complete module
 
 byteCodeGen :: HscEnv
             -> Module
-            -> [StgTopBinding]
+            -> [CgStgTopBinding]
             -> [TyCon]
             -> Maybe ModBreaks
             -> IO CompiledByteCode
@@ -116,8 +116,7 @@ byteCodeGen hsc_env this_mod binds tycs mb_modBreaks
 
         (BcM_State{..}, proto_bcos) <-
            runBc hsc_env this_mod mb_modBreaks (mkVarEnv stringPtrs) $ do
-             let flattened_binds =
-                   concatMap (flattenBind . annBindingFreeVars) (reverse lifted_binds)
+             let flattened_binds = concatMap flattenBind (reverse lifted_binds)
              mapM schemeTopBind flattened_binds
 
         when (notNull ffis)
diff --git a/compiler/GHC/Types/Name.hs b/compiler/GHC/Types/Name.hs
index 527f8da0fe..2c3cfac8c5 100644
--- a/compiler/GHC/Types/Name.hs
+++ b/compiler/GHC/Types/Name.hs
@@ -64,7 +64,7 @@ module GHC.Types.Name (
         isWiredInName, isWiredIn, isBuiltInSyntax,
         isHoleName,
         wiredInNameTyThing_maybe,
-        nameIsLocalOrFrom, nameIsHomePackage,
+        nameIsLocalOrFrom, nameIsExternalOrFrom, nameIsHomePackage,
         nameIsHomePackageImport, nameIsFromExternalPackage,
         stableNameCmp,
 
@@ -327,6 +327,9 @@ nameModule_maybe (Name { n_sort = External mod})    = Just mod
 nameModule_maybe (Name { n_sort = WiredIn mod _ _}) = Just mod
 nameModule_maybe _                                  = Nothing
 
+is_interactive_or_from :: Module -> Module -> Bool
+is_interactive_or_from from mod = from == mod || isInteractiveModule mod
+
 nameIsLocalOrFrom :: Module -> Name -> Bool
 -- ^ Returns True if the name is
 --   (a) Internal
@@ -351,9 +354,16 @@ nameIsLocalOrFrom :: Module -> Name -> Bool
 -- See Note [The interactive package] in "GHC.Runtime.Context"
 
 nameIsLocalOrFrom from name
-  | Just mod <- nameModule_maybe name = from == mod || isInteractiveModule mod
+  | Just mod <- nameModule_maybe name = is_interactive_or_from from mod
   | otherwise                         = True
 
+nameIsExternalOrFrom :: Module -> Name -> Bool
+-- ^ Returns True if the name is external or from the 'interactive' package
+-- See documentation of `nameIsLocalOrFrom` function
+nameIsExternalOrFrom from name
+  | Just mod <- nameModule_maybe name = is_interactive_or_from from mod
+  | otherwise                         = False
+
 nameIsHomePackage :: Module -> Name -> Bool
 -- True if the Name is defined in module of this package
 nameIsHomePackage this_mod
diff --git a/compiler/ghc.cabal.in b/compiler/ghc.cabal.in
index 750dc0389d..232d89c89f 100644
--- a/compiler/ghc.cabal.in
+++ b/compiler/ghc.cabal.in
@@ -546,7 +546,6 @@ Library
         GHC.Stg.BcPrep
         GHC.Stg.CSE
         GHC.Stg.Debug
-        GHC.Stg.DepAnal
         GHC.Stg.FVs
         GHC.Stg.Lift
         GHC.Stg.Lift.Analysis
diff --git a/testsuite/tests/codeGen/should_compile/T14373d.stderr-ws-32 b/testsuite/tests/codeGen/should_compile/T14373d.stderr-ws-32
index 884e8abcca..2ff5530f61 100644
--- a/testsuite/tests/codeGen/should_compile/T14373d.stderr-ws-32
+++ b/testsuite/tests/codeGen/should_compile/T14373d.stderr-ws-32
@@ -1,19 +1,19 @@
+[T14373d.lateDefault_entry() { //
+           switch [0 .. 15]
+               case 15 : goto 
+               default: {goto 
+           R1 = XYZ_closure+2;
 [T14373d.earlyDefault_entry() { //
            switch [1 .. 3]
-               case 2 : goto
-               default: {goto
+               case 2 : goto 
+               default: {goto 
            R1 = XYZ_closure+2;
 [T14373d.mixedDefault_entry() { //
            switch [1 .. 3]
-               case 2 : goto
-               case 3 : goto
-               default: {goto
+               case 2 : goto 
+               case 3 : goto 
+               default: {goto 
            switch [2 .. 15]
-               case 15 : goto
-               default: {goto
-           R1 = XYZ_closure+2;
-[T14373d.lateDefault_entry() { //
-           switch [0 .. 15]
-               case 15 : goto
-               default: {goto
+               case 15 : goto 
+               default: {goto 
            R1 = XYZ_closure+2;
diff --git a/testsuite/tests/codeGen/should_compile/T14373d.stderr-ws-64 b/testsuite/tests/codeGen/should_compile/T14373d.stderr-ws-64
index 6258d38e4a..0850a479cc 100644
--- a/testsuite/tests/codeGen/should_compile/T14373d.stderr-ws-64
+++ b/testsuite/tests/codeGen/should_compile/T14373d.stderr-ws-64
@@ -1,19 +1,19 @@
+[T14373d.lateDefault_entry() { //
+           switch [0 .. 15]
+               case 15 : goto 
+               default: {goto 
+           R1 = XYZ_closure+2;
 [T14373d.earlyDefault_entry() { //
            switch [1 .. 7]
-               case 2 : goto
-               default: {goto
+               case 2 : goto 
+               default: {goto 
            R1 = XYZ_closure+2;
 [T14373d.mixedDefault_entry() { //
            switch [1 .. 7]
-               case 2 : goto
-               case 7 : goto
-               default: {goto
+               case 2 : goto 
+               case 7 : goto 
+               default: {goto 
            switch [6 .. 15]
-               case 15 : goto
-               default: {goto
-           R1 = XYZ_closure+2;
-[T14373d.lateDefault_entry() { //
-           switch [0 .. 15]
-               case 15 : goto
-               default: {goto
+               case 15 : goto 
+               default: {goto 
            R1 = XYZ_closure+2;
diff --git a/testsuite/tests/ghci.debugger/scripts/all.T b/testsuite/tests/ghci.debugger/scripts/all.T
index dd68cb3c9f..661c75e4ad 100644
--- a/testsuite/tests/ghci.debugger/scripts/all.T
+++ b/testsuite/tests/ghci.debugger/scripts/all.T
@@ -72,8 +72,7 @@ test('break015', expect_broken(1532), ghci_script, ['break015.script'])
 test('break016', combined_output, ghci_script, ['break016.script'])
 test('break017', [extra_files(['../QSort.hs']),
                   combined_output], ghci_script, ['break017.script'])
-test('break018', [expect_broken(18004), extra_files(['../mdo.hs'])],
-     ghci_script, ['break018.script'])
+test('break018', extra_files(['../mdo.hs']), ghci_script, ['break018.script'])
 test('break019', extra_files(['../Test2.hs']), ghci_script, ['break019.script'])
 test('break020', extra_files(['Break020b.hs']), ghci_script, ['break020.script'])
 test('break021', extra_files(['Break020b.hs', 'break020.hs']), ghci_script, ['break021.script'])
diff --git a/testsuite/tests/ghci.debugger/scripts/break021.stdout b/testsuite/tests/ghci.debugger/scripts/break021.stdout
index bf64680b1a..199b3cdf05 100644
--- a/testsuite/tests/ghci.debugger/scripts/break021.stdout
+++ b/testsuite/tests/ghci.debugger/scripts/break021.stdout
@@ -17,7 +17,7 @@ _result :: IO () = _
       ^^^^^^^
 11    line2 0
 Stopped in Main.line1, break020.hs:3:11-19
-_result :: IO () = _
+_result :: m () = _
 2  
 3  line1 _ = return ()
              ^^^^^^^^^
@@ -29,7 +29,7 @@ _result :: IO () = _
       ^^^^^^^
 12    in_another_decl 0
 Stopped in Main.line2, break020.hs:4:11-19
-_result :: IO () = _
+_result :: m () = _
 3  line1 _ = return ()
 4  line2 _ = return ()
              ^^^^^^^^^
@@ -41,7 +41,7 @@ _result :: IO () = _
       ^^^^^^^^^^^^^^^^^
 13    in_another_module 0
 Stopped in Main.in_another_decl, break020.hs:(6,21)-(7,30)
-_result :: IO () = _
+_result :: m () = _
 5  
                      vv
 6  in_another_decl _ = do line1 0
@@ -49,25 +49,25 @@ _result :: IO () = _
                                  ^^
 8  
 Stopped in Main.in_another_decl, break020.hs:6:24-30
-_result :: IO () = _
+_result :: m () = _
 5  
 6  in_another_decl _ = do line1 0
                           ^^^^^^^
 7                         line2 0
 Stopped in Main.line1, break020.hs:3:11-19
-_result :: IO () = _
+_result :: m () = _
 2  
 3  line1 _ = return ()
              ^^^^^^^^^
 4  line2 _ = return ()
 Stopped in Main.in_another_decl, break020.hs:7:24-30
-_result :: IO () = _
+_result :: m () = _
 6  in_another_decl _ = do line1 0
 7                         line2 0
                           ^^^^^^^
 8  
 Stopped in Main.line2, break020.hs:4:11-19
-_result :: IO () = _
+_result :: m () = _
 3  line1 _ = return ()
 4  line2 _ = return ()
              ^^^^^^^^^
@@ -85,7 +85,7 @@ _result :: IO () = _
       ^^^^^^^
 15    return ()
 Stopped in Main.line2, break020.hs:4:11-19
-_result :: IO () = _
+_result :: m () = _
 3  line1 _ = return ()
 4  line2 _ = return ()
              ^^^^^^^^^