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).

8 files changed, 218 insertions, 266 deletions

diff --git a/compiler/GHC/Driver/Main.hs b/compiler/GHC/Driver/Main.hs
index 26647df369..1f5cab3b97 100644
--- a/compiler/GHC/Driver/Main.hs
+++ b/compiler/GHC/Driver/Main.hs
@@ -159,7 +159,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
@@ -1766,24 +1765,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
@@ -1828,7 +1825,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
@@ -1851,7 +1848,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
@@ -1859,11 +1856,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..1f3d7f4125 100644
--- a/compiler/GHC/Stg/FVs.hs
+++ b/compiler/GHC/Stg/FVs.hs
@@ -40,7 +40,7 @@ Top-level closure bindings never capture variables as all of their free
 variables are global.
 -}
 module GHC.Stg.FVs (
-    annTopBindingsFreeVars,
+    depSortWithAnnotStgPgm,
     annBindingFreeVars
   ) where
 
@@ -48,39 +48,63 @@ import GHC.Prelude
 
 import GHC.Stg.Syntax
 import GHC.Types.Id
-import GHC.Types.Var.Set
+import GHC.Types.Name (Name, nameIsLocalOrFrom)
+import GHC.Types.Name.Env
+import GHC.Utils.Outputable
+import GHC.Utils.Panic
 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 )
-
-newtype Env
-  = Env
-  { locals :: IdSet
-  }
+import Data.Graph (SCC (..))
+import Data.Bifunctor (first)
 
-emptyEnv :: Env
-emptyEnv = Env emptyVarSet
+--------------------------------------------------------------------------------
+-- * Dependency sorting
 
-addLocals :: [Id] -> Env -> Env
-addLocals bndrs env
-  = env { locals = extendVarSetList (locals env) bndrs }
+-- | Dependency sort a STG program so that dependencies come before uses; also
+--  perform non-global free variable analysis by annotating non-top-level
+-- closure bindings with captured variables.
+depSortWithAnnotStgPgm :: Module -> [StgTopBinding] -> [CgStgTopBinding]
+depSortWithAnnotStgPgm this_mod =
+    {-# SCC "STG.depSortWithAnnotStgPgm" #-}
+    map fst . depSort . annAndDepAnalStgFVs this_mod
 
--- | Annotates a top-level STG binding group with its free variables.
-annTopBindingsFreeVars :: [StgTopBinding] -> [CgStgTopBinding]
-annTopBindingsFreeVars = map go
+-- | Sort free-variable-annotated STG bindings so that dependencies come before
+-- uses.
+depSort :: [(CgStgTopBinding, FVs)] -> [(CgStgTopBinding, FVs)]
+depSort = concatMap get_binds . depAnal defs uses
   where
-    go (StgTopStringLit id bs) = StgTopStringLit id bs
-    go (StgTopLifted bind)
-      = StgTopLifted (annBindingFreeVars bind)
+    uses, defs :: (CgStgTopBinding, FVs) -> [Name]
 
--- | Annotates an STG binding with its free variables.
-annBindingFreeVars :: StgBinding -> CgStgBinding
-annBindingFreeVars = fst . binding emptyEnv emptyDVarSet
+    -- 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.
 
-boundIds :: StgBinding -> [Id]
-boundIds (StgNonRec b _) = [b]
-boundIds (StgRec pairs)  = map fst pairs
+    uses (StgTopStringLit{}, _) = []
+    uses (StgTopLifted{}, fvs)  = map idName (nonDetEltsUniqSet fvs)
+
+    defs (bind, _) = map idName (bindersOfTop bind)
+
+    get_binds :: SCC (CgStgTopBinding, FVs) -> [(CgStgTopBinding, FVs)]
+    get_binds (AcyclicSCC bind) =
+      [bind]
+    get_binds (CyclicSCC binds) =
+      pprPanic "depSortStgBinds"
+               (text "Found cyclic SCC:"
+               $$ ppr (map (first (pprStgTopBinding panicStgPprOpts)) binds))
+
+--------------------------------------------------------------------------------
+-- * Non-global free variable analysis
+
+newtype Env
+  = Env
+  { locals :: IdSet
+  }
 
 -- Note [Tracking local binders]
 -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -89,80 +113,155 @@ boundIds (StgRec pairs)  = map fst pairs
 -- places where new local binders are introduced.
 -- Why do it there rather than in 'binding'? Two reasons:
 --
---   1. We call 'binding' from 'annTopBindingsFreeVars', which would
+--   1. We call 'binding' from 'annAndDepAnalStgFVs', 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.
+emptyEnv :: Env
+emptyEnv = Env emptyVarSet
+
+addLocals :: [Id] -> Env -> Env
+addLocals bndrs env
+  = env { locals = extendVarSetList (locals env) bndrs }
+
 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
+boundIds :: StgBinding -> [Id]
+boundIds (StgNonRec b _) = [b]
+boundIds (StgRec pairs)  = map fst pairs
 
-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
+--------------------------------------------------------------------------------
+-- * 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.
+--
+
+-- | Perform Dependency sorting as well as annotate each binding with non-global free variables
+annAndDepAnalStgFVs :: Module -> [StgTopBinding] -> [(CgStgTopBinding,FVs)]
+annAndDepAnalStgFVs this_mod bs = map go bs
   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)
+    go (StgTopStringLit id bs) = (StgTopStringLit id bs, emptyVarSet)
+    go (StgTopLifted bind)
+      | (bind', fvs, _) <- binding this_mod emptyVarSet emptyEnv emptyDVarSet bind
+      = (StgTopLifted bind', fvs)
+
+annBindingFreeVars :: Module -> StgBinding -> CgStgBinding
+annBindingFreeVars this_mod = fstOf3 . binding this_mod emptyVarSet emptyEnv emptyDVarSet
+
+binding :: Module -> BVs -> Env -> DIdSet -> StgBinding -> (CgStgBinding, FVs, DIdSet)
+binding mod bounds env body_fv b =
+  case b of
+    StgNonRec bndr r -> (StgNonRec bndr r', fvs, id_set)
       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')
+        (r', fvs, rhs_id_set) = rhs bounds env r
+        id_set = delDVarSet body_fv bndr `unionDVarSet` rhs_id_set
+    StgRec pairs -> (StgRec pairs', fvs, id_sets)
       where
-        (e', fvs) = go e
-        fvs' = unionDVarSet (tickish tick) fvs
-        tickish (Breakpoint _ _ ids) = mkDVarSet ids
-        tickish _                    = emptyDVarSet
+        -- See Note [Tracking local binders]
+        bndrs = map fst pairs
+        bounds' = extendVarSetList bounds bndrs
+        (rhss, rhs_fvss, rhs_id_sets) = mapAndUnzip3 (rhs bounds' env . snd) pairs
+        fvs = unionVarSets rhs_fvss
+        pairs' = zip bndrs rhss
+        id_sets = delDVarSetList (unionDVarSets (body_fv:rhs_id_sets)) bndrs
+  where
+    var_fvs :: BVs -> Var -> FVs
+    var_fvs bounds v
+      | not (elemVarSet v bounds)
+      , nameIsLocalOrFrom mod (idName v)
+      = unitVarSet v
+      | otherwise
+      = emptyVarSet
 
-    go_bind dc bind body = (dc bind' body', fvs)
+    expr :: BVs -> Env -> StgExpr -> (CgStgExpr, FVs, DIdSet)
+    expr bounds env = go
       where
-        -- See Note [Tracking local binders]
-        env' = addLocals (boundIds bind) env
-        (body', body_fvs) = expr env' body
-        (bind', fvs) = binding env' body_fvs bind
+        go (StgApp f as)
+          | (args_fvs, id_set) <- args bounds env as
+          = ( StgApp f as
+            , var_fvs bounds f `unionVarSet` args_fvs
+            , unionDVarSet id_set (mkFreeVarSet env [f]))
+        go (StgLit lit) = (StgLit lit, emptyVarSet, emptyDVarSet)
+        go (StgConApp dc n as tys)
+          | (args_fvs, id_set) <- args bounds env as
+          = (StgConApp dc n as tys, args_fvs, id_set)
+        go (StgOpApp op as ty)
+          | (fvs, id_set) <- args bounds env as
+          = (StgOpApp op as ty, fvs, id_set)
+        go (StgCase scrut bndr ty alts)
+          | (scrut',scrut_fvs,scrut_id_set) <- expr bounds env scrut
+          -- See Note [Tracking local binders]
+          , (alts',alts_fvss,alts_id_sets)
+              <- mapAndUnzip3 (alt (extendVarSet bounds bndr) (addLocals [bndr] env)) alts
+          , fvs' <- scrut_fvs `unionVarSet` (unionVarSets alts_fvss)
+          , alts_id_set <- unionDVarSets alts_id_sets
+          , id_set' <- delDVarSet (unionDVarSet scrut_id_set alts_id_set) bndr
+          = (StgCase scrut' bndr ty alts', fvs',id_set')
+        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',fvs, id_set) <- expr bounds env e
+          , id_set' <- unionDVarSet (tickish tick) id_set
+          = (StgTick tick e', fvs, id_set')
+            where
+              tickish (Breakpoint _ _ ids) = mkDVarSet ids
+              tickish _                    = emptyDVarSet
 
-rhs :: Env -> StgRhs -> (CgStgRhs, DIdSet)
-rhs env (StgRhsClosure _ ccs uf bndrs body)
-  = (StgRhsClosure fvs ccs uf bndrs body', fvs)
-  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_bind dc bind body = (dc bind' body', fvs, id_set)
+          where
+            -- See Note [Tracking local binders]
+            env' = addLocals (boundIds bind) env
+            (body', body_fvs, body_set_ids)
+                = expr (extendVarSetList bounds (bindersOf bind)) env' body
+            (bind', bind_fvs, id_set)
+                = binding mod bounds env' body_set_ids bind
+            fvs = bind_fvs `unionVarSet` body_fvs
 
-alt :: Env -> StgAlt -> (CgStgAlt, DIdSet)
-alt env (con, bndrs, e) = ((con, bndrs, e'), fvs)
-  where
-    -- See Note [Tracking local binders]
-    (e', rhs_fvs) = expr (addLocals bndrs env) e
-    fvs = delDVarSetList rhs_fvs bndrs
+
+    rhs :: BVs -> Env -> StgRhs -> (CgStgRhs, FVs, DIdSet)
+    rhs bounds env (StgRhsClosure _ ccs uf bs body)
+      | (body', fvss, id_set)
+        <- expr (extendVarSetList bounds bs) (addLocals bs env) body
+      , id_set' <- delDVarSetList id_set bs
+      = (StgRhsClosure id_set' ccs uf bs body', fvss, id_set')
+    rhs bounds env (StgRhsCon ccs dc mu ts bs)
+      | (fvs, id_set) <- args bounds env bs
+      = (StgRhsCon ccs dc mu ts bs, fvs, id_set)
+
+    args :: BVs -> Env -> [StgArg] -> (FVs, DIdSet)
+    args bounds env = foldl' f (emptyVarSet, emptyDVarSet)
+      where
+        f (fvs,ids) StgLitArg{}  = (fvs, ids )
+        f (fvs,ids) (StgVarArg v)  = (fvs', ids')
+          where
+            !fvs' = var_fvs bounds v `unionVarSet` fvs
+            !ids' | v `elemVarSet` locals env
+                  = extendDVarSet ids v
+                  | otherwise = ids
+
+    alt :: BVs -> Env -> StgAlt -> (CgStgAlt, FVs, DIdSet)
+    alt bounds env (con,bndrs,e)
+      | (e', fvs, id_set)
+        <- expr (extendVarSetList bounds bndrs) (addLocals bndrs env) e
+      , id_set' <- delDVarSetList id_set bndrs
+      = ((con,bndrs, e'), fvs, id_set')
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.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