From f84ff16139501ebd4bd822ff25d0f4a8d61f2d05 Mon Sep 17 00:00:00 2001
From: Sylvain Henry <sylvain@haskus.fr>
Date: Wed, 31 May 2023 13:03:48 +0200
Subject: [PATCH] Stg: return imported FVs

This is used to determine what to link when using the interpreter.
For now it's only used by the JS interpreter but it could easily be used
by the native interpreter too (instead of extracting names from compiled BCOs).
---
 compiler/GHC/Driver/Main.hs                   |  25 ++--
 compiler/GHC/Stg/FVs.hs                       | 115 ++++++++++--------
 compiler/GHC/Stg/Pipeline.hs                  |   8 +-
 compiler/GHC/Utils/Misc.hs                    |  17 ++-
 .../should_run/control-flow/LoadCmmGroup.hs   |   2 +-
 5 files changed, 104 insertions(+), 63 deletions(-)

diff --git a/compiler/GHC/Driver/Main.hs b/compiler/GHC/Driver/Main.hs
index ba176b6ce568..7233ffff6122 100644
--- a/compiler/GHC/Driver/Main.hs
+++ b/compiler/GHC/Driver/Main.hs
@@ -234,6 +234,7 @@ import GHC.Types.SafeHaskell
 import GHC.Types.ForeignStubs
 import GHC.Types.Name.Env      ( mkNameEnv )
 import GHC.Types.Var.Env       ( mkEmptyTidyEnv )
+import GHC.Types.Var.Set
 import GHC.Types.Error
 import GHC.Types.Fixity.Env
 import GHC.Types.CostCentre
@@ -248,6 +249,7 @@ import GHC.Types.Name.Set (NonCaffySet)
 import GHC.Types.TyThing
 import GHC.Types.HpcInfo
 import GHC.Types.Unique.Supply (uniqFromMask)
+import GHC.Types.Unique.Set
 
 import GHC.Utils.Fingerprint ( Fingerprint )
 import GHC.Utils.Panic
@@ -278,7 +280,7 @@ import System.FilePath as FilePath
 import System.Directory
 import qualified Data.Set as S
 import Data.Set (Set)
-import Data.Functor
+import Data.Functor ((<&>))
 import Control.DeepSeq (force)
 import Data.Bifunctor (first)
 import Data.List.NonEmpty (NonEmpty ((:|)))
@@ -1847,7 +1849,7 @@ hscGenHardCode hsc_env cgguts location output_filename = do
             this_mod location late_cc_binds data_tycons
 
         -----------------  Convert to STG ------------------
-        (stg_binds, denv, (caf_ccs, caf_cc_stacks), stg_cg_infos)
+        (stg_binds_with_deps, denv, (caf_ccs, caf_cc_stacks), stg_cg_infos)
             <- {-# SCC "CoreToStg" #-}
                withTiming logger
                    (text "CoreToStg"<+>brackets (ppr this_mod))
@@ -1859,6 +1861,8 @@ hscGenHardCode hsc_env cgguts location output_filename = do
                         (seqEltsUFM (seqTagSig) tag_env))
                    (myCoreToStg logger dflags (interactiveInScope (hsc_IC hsc_env)) False this_mod location prepd_binds)
 
+        let (stg_binds,_stg_deps) = unzip stg_binds_with_deps
+
         let cost_centre_info =
               (late_local_ccs ++ caf_ccs, caf_cc_stacks)
             platform = targetPlatform dflags
@@ -1977,9 +1981,12 @@ hscInteractive hsc_env cgguts location = do
 
     -- The stg cg info only provides a runtime benfit, but is not requires so we just
     -- omit it here
-    (stg_binds, _infotable_prov, _caf_ccs__caf_cc_stacks, _ignore_stg_cg_infos)
+    (stg_binds_with_deps, _infotable_prov, _caf_ccs__caf_cc_stacks, _ignore_stg_cg_infos)
       <- {-# SCC "CoreToStg" #-}
           myCoreToStg logger dflags (interactiveInScope (hsc_IC hsc_env)) True this_mod location prepd_binds
+
+    let (stg_binds,_stg_deps) = unzip stg_binds_with_deps
+
     -----------------  Generate byte code ------------------
     comp_bc <- byteCodeGen hsc_env this_mod stg_binds data_tycons mod_breaks
     ------------------ Create f-x-dynamic C-side stuff -----
@@ -2157,7 +2164,7 @@ doCodeGen hsc_env this_mod denv data_tycons
 myCoreToStg :: Logger -> DynFlags -> [Var]
             -> Bool
             -> Module -> ModLocation -> CoreProgram
-            -> IO ( [CgStgTopBinding] -- output program
+            -> IO ( [(CgStgTopBinding,IdSet)] -- output program and its dependencies
                   , InfoTableProvMap
                   , CollectedCCs -- CAF cost centre info (declared and used)
                   , StgCgInfos )
@@ -2172,7 +2179,7 @@ myCoreToStg logger dflags ic_inscope for_bytecode this_mod ml prepd_binds = do
                    this_mod stg_binds
 
     putDumpFileMaybe logger Opt_D_dump_stg_cg "CodeGenInput STG:" FormatSTG
-        (pprGenStgTopBindings (initStgPprOpts dflags) stg_binds_with_fvs)
+        (pprGenStgTopBindings (initStgPprOpts dflags) (fmap fst stg_binds_with_fvs))
 
     return (stg_binds_with_fvs, denv, cost_centre_info, stg_cg_info)
 
@@ -2325,7 +2332,7 @@ hscParsedDecls hsc_env decls = runInteractiveHsc hsc_env $ do
         (initCorePrepPgmConfig (hsc_dflags hsc_env) (interactiveInScope $ hsc_IC hsc_env))
         this_mod iNTERACTIVELoc core_binds data_tycons
 
-    (stg_binds, _infotable_prov, _caf_ccs__caf_cc_stacks, _stg_cg_info)
+    (stg_binds_with_deps, _infotable_prov, _caf_ccs__caf_cc_stacks, _stg_cg_info)
         <- {-# SCC "CoreToStg" #-}
            liftIO $ myCoreToStg (hsc_logger hsc_env)
                                 (hsc_dflags hsc_env)
@@ -2335,6 +2342,8 @@ hscParsedDecls hsc_env decls = runInteractiveHsc hsc_env $ do
                                 iNTERACTIVELoc
                                 prepd_binds
 
+    let (stg_binds,_stg_deps) = unzip stg_binds_with_deps
+
     {- Generate byte code -}
     cbc <- liftIO $ byteCodeGen hsc_env this_mod
                                 stg_binds data_tycons mod_breaks
@@ -2590,7 +2599,7 @@ hscCompileCoreExpr' hsc_env srcspan ds_expr = do
   let this_mod = mkInteractiveModule (show u)
   let for_bytecode = True
 
-  (stg_binds, _prov_map, _collected_ccs, _stg_cg_infos) <-
+  (stg_binds_with_deps, _prov_map, _collected_ccs, _stg_cg_infos) <-
        myCoreToStg logger
                    dflags
                    (interactiveInScope (hsc_IC hsc_env))
@@ -2599,6 +2608,8 @@ hscCompileCoreExpr' hsc_env srcspan ds_expr = do
                    this_loc
                    [NonRec binding_id prepd_expr]
 
+  let (stg_binds, _stg_deps) = unzip stg_binds_with_deps
+
   let interp = hscInterp hsc_env
   let tmpfs = hsc_tmpfs hsc_env
   let tmp_dir = tmpDir dflags
diff --git a/compiler/GHC/Stg/FVs.hs b/compiler/GHC/Stg/FVs.hs
index 801ac1fed291..72f25f7bcb6c 100644
--- a/compiler/GHC/Stg/FVs.hs
+++ b/compiler/GHC/Stg/FVs.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE LambdaCase #-}
 
 {- |
 Non-global free variable analysis on STG terms. This pass annotates
@@ -84,26 +85,31 @@ But isn't it in correct dependency order already?  No:
 --     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 :: Module -> [StgTopBinding] -> [(CgStgTopBinding,ImpFVs)]
 depSortWithAnnotStgPgm this_mod binds
   = {-# SCC "STG.depSortWithAnnotStgPgm" #-}
-    lit_binds ++ map from_scc sccs
+    zip lit_binds (repeat emptyVarSet) ++ map from_scc sccs
   where
     lit_binds :: [CgStgTopBinding]
     pairs     :: [(Id, StgRhs)]
     (lit_binds, pairs) = flattenTopStgBindings binds
 
-    nodes :: [Node Name (Id, CgStgRhs)]
+    nodes :: [Node Name (Id, CgStgRhs, ImpFVs)]
     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 :: [SCC (Id,CgStgRhs,ImpFVs)]
     sccs  = stronglyConnCompFromEdgedVerticesUniq nodes
 
-    from_scc (CyclicSCC pairs)       = StgTopLifted (StgRec pairs)
-    from_scc (AcyclicSCC (bndr,rhs)) = StgTopLifted (StgNonRec bndr rhs)
+    from_scc = \case
+      AcyclicSCC (bndr,rhs,imp_fvs) -> (StgTopLifted (StgNonRec bndr rhs), imp_fvs)
+      CyclicSCC triples             -> (StgTopLifted (StgRec pairs), imp_fvs)
+        where
+          (ids,rhss,imp_fvss) = unzip3 triples
+          pairs = zip ids rhss
+          imp_fvs = unionVarSets imp_fvss
 
 
 flattenTopStgBindings :: [StgTopBinding] -> ([CgStgTopBinding], [(Id,StgRhs)])
@@ -119,13 +125,13 @@ flattenTopStgBindings binds
     flatten_one (StgNonRec b r) = [(b,r)]
     flatten_one (StgRec pairs)  = pairs
 
-annotateTopPair :: Env -> (Id, StgRhs) -> Node Name (Id, CgStgRhs)
+annotateTopPair :: Env -> (Id, StgRhs) -> Node Name (Id, CgStgRhs, ImpFVs)
 annotateTopPair env0 (bndr, rhs)
   = DigraphNode { node_key          = idName bndr
-                , node_payload      = (bndr, rhs')
+                , node_payload      = (bndr, rhs', imp_fvs)
                 , node_dependencies = map idName (nonDetEltsUniqSet top_fvs) }
   where
-    (rhs', top_fvs, _) = rhsFVs env0 rhs
+    (rhs', imp_fvs, top_fvs, _) = rhsFVs env0 rhs
 
 --------------------------------------------------------------------------------
 -- * Non-global free variable analysis
@@ -158,6 +164,12 @@ addLocals bndrs env
 -- analysis on the top-level bindings.
 type TopFVs   = IdSet
 
+-- | ImpFVs: set of variables that are imported
+--
+-- It is a /non-deterministic/ set because we use it only to perform module
+-- dependency analysis.
+type ImpFVs   = 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'
@@ -181,97 +193,100 @@ type LocalFVs = DIdSet
 --
 
 annBindingFreeVars :: Module -> StgBinding -> CgStgBinding
-annBindingFreeVars this_mod = fstOf3 . bindingFVs (Env emptyVarSet this_mod) emptyDVarSet
+annBindingFreeVars this_mod = fstOf4 . bindingFVs (Env emptyVarSet this_mod) emptyDVarSet
 
-bindingFVs :: Env -> LocalFVs -> StgBinding -> (CgStgBinding, TopFVs, LocalFVs)
+bindingFVs :: Env -> LocalFVs -> StgBinding -> (CgStgBinding, ImpFVs, TopFVs, LocalFVs)
 bindingFVs env body_fv b =
   case b of
-    StgNonRec bndr r -> (StgNonRec bndr r', fvs, lcl_fvs)
+    StgNonRec bndr r -> (StgNonRec bndr r', imp_fvs, top_fvs, lcl_fvs)
       where
-        (r', fvs, rhs_lcl_fvs) = rhsFVs env r
+        (r', imp_fvs, top_fvs, rhs_lcl_fvs) = rhsFVs env r
         lcl_fvs = delDVarSet body_fv bndr `unionDVarSet` rhs_lcl_fvs
 
-    StgRec pairs -> (StgRec pairs', fvs, lcl_fvss)
+    StgRec pairs -> (StgRec pairs', imp_fvs, top_fvs, lcl_fvss)
       where
         bndrs = map fst pairs
         env' = addLocals bndrs env
-        (rhss, rhs_fvss, rhs_lcl_fvss) = mapAndUnzip3 (rhsFVs env' . snd) pairs
-        fvs = unionVarSets rhs_fvss
+        (rhss, rhs_imp_fvss, rhs_top_fvss, rhs_lcl_fvss) = mapAndUnzip4 (rhsFVs env' . snd) pairs
+        top_fvs = unionVarSets rhs_top_fvss
+        imp_fvs = unionVarSets rhs_imp_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)
+varFVs :: Env -> Id -> (ImpFVs, TopFVs, LocalFVs) -> (ImpFVs, TopFVs, LocalFVs)
+varFVs env v (imp_fvs, top_fvs, lcl_fvs)
   | v `elemVarSet` locals env                -- v is locally bound
-  = (top_fvs, lcl_fvs `extendDVarSet` v)
+  = (imp_fvs, top_fvs, lcl_fvs `extendDVarSet` v)
   | nameIsLocalOrFrom (mod env) (idName v)   -- v is bound at top level
-  = (top_fvs `extendVarSet` v, lcl_fvs)
+  = (imp_fvs, top_fvs `extendVarSet` v, lcl_fvs)
   | otherwise                                -- v is imported
-  = (top_fvs, lcl_fvs)
+  = (imp_fvs `extendVarSet` v, top_fvs, lcl_fvs)
 
-exprFVs :: Env -> StgExpr -> (CgStgExpr, TopFVs, LocalFVs)
+exprFVs :: Env -> StgExpr -> (CgStgExpr, ImpFVs, TopFVs, LocalFVs)
 exprFVs env = go
   where
     go (StgApp f as)
-      | (top_fvs, lcl_fvs) <- varFVs env f (argsFVs env as)
-      = (StgApp f as, top_fvs, lcl_fvs)
+      | (imp_fvs, top_fvs, lcl_fvs) <- varFVs env f (argsFVs env as)
+      = (StgApp f as, imp_fvs, top_fvs, lcl_fvs)
 
-    go (StgLit lit) = (StgLit lit, emptyVarSet, emptyDVarSet)
+    go (StgLit lit) = (StgLit lit, emptyVarSet, emptyVarSet, emptyDVarSet)
 
     go (StgConApp dc n as tys)
-      | (top_fvs, lcl_fvs) <- argsFVs env as
-      = (StgConApp dc n as tys, top_fvs, lcl_fvs)
+      | (imp_fvs, top_fvs, lcl_fvs) <- argsFVs env as
+      = (StgConApp dc n as tys, imp_fvs, top_fvs, lcl_fvs)
 
     go (StgOpApp op as ty)
-      | (top_fvs, lcl_fvs) <- argsFVs env as
-      = (StgOpApp op as ty, top_fvs, lcl_fvs)
+      | (imp_fvs, top_fvs, lcl_fvs) <- argsFVs env as
+      = (StgOpApp op as ty, imp_fvs, 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
+      | (scrut',scrut_imp_fvs,scrut_top_fvs,scrut_lcl_fvs) <- exprFVs env scrut
+      , (alts',alts_imp_fvss,alts_top_fvss,alts_lcl_fvss)
+          <- mapAndUnzip4 (altFVs (addLocals [bndr] env)) alts
       , let top_fvs = scrut_top_fvs `unionVarSet` unionVarSets alts_top_fvss
+            imp_fvs = scrut_imp_fvs `unionVarSet` unionVarSets alts_imp_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)
+      = (StgCase scrut' bndr ty alts', imp_fvs, 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
+      | (e', imp_fvs, top_fvs, lcl_fvs) <- exprFVs env e
       , let lcl_fvs' = unionDVarSet (tickish tick) lcl_fvs
-      = (StgTick tick e', top_fvs, lcl_fvs')
+      = (StgTick tick e', imp_fvs, top_fvs, lcl_fvs')
         where
           tickish (Breakpoint _ _ ids) = mkDVarSet ids
           tickish _                    = emptyDVarSet
 
-    go_bind dc bind body = (dc bind' body', top_fvs, lcl_fvs)
+    go_bind dc bind body = (dc bind' body', imp_fvs, 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
+        (body', body_imp_fvs, body_top_fvs, body_lcl_fvs) = exprFVs env' body
+        (bind', bind_imp_fvs, bind_top_fvs, lcl_fvs)      = bindingFVs env' body_lcl_fvs bind
         top_fvs = bind_top_fvs `unionVarSet` body_top_fvs
+        imp_fvs = bind_imp_fvs `unionVarSet` body_imp_fvs
 
 
-rhsFVs :: Env -> StgRhs -> (CgStgRhs, TopFVs, LocalFVs)
+rhsFVs :: Env -> StgRhs -> (CgStgRhs, ImpFVs, TopFVs, LocalFVs)
 rhsFVs env (StgRhsClosure _ ccs uf bs body typ)
-  | (body', top_fvs, lcl_fvs) <- exprFVs (addLocals bs env) body
+  | (body', imp_fvs, top_fvs, lcl_fvs) <- exprFVs (addLocals bs env) body
   , let lcl_fvs' = delDVarSetList lcl_fvs bs
-  = (StgRhsClosure lcl_fvs' ccs uf bs body' typ, top_fvs, lcl_fvs')
+  = (StgRhsClosure lcl_fvs' ccs uf bs body' typ, imp_fvs, top_fvs, lcl_fvs')
 rhsFVs env (StgRhsCon ccs dc mu ts bs typ)
-  | (top_fvs, lcl_fvs) <- argsFVs env bs
-  = (StgRhsCon ccs dc mu ts bs typ, top_fvs, lcl_fvs)
+  | (imp_fvs, top_fvs, lcl_fvs) <- argsFVs env bs
+  = (StgRhsCon ccs dc mu ts bs typ, imp_fvs, top_fvs, lcl_fvs)
 
-argsFVs :: Env -> [StgArg] -> (TopFVs, LocalFVs)
-argsFVs env = foldl' f (emptyVarSet, emptyDVarSet)
+argsFVs :: Env -> [StgArg] -> (ImpFVs, TopFVs, LocalFVs)
+argsFVs env = foldl' f (emptyVarSet, emptyVarSet, emptyDVarSet)
   where
-    f (fvs,ids) StgLitArg{}   = (fvs, ids)
-    f (fvs,ids) (StgVarArg v) = varFVs env v (fvs, ids)
+    f (imp_fvs,fvs,ids) StgLitArg{}   = (imp_fvs, fvs, ids)
+    f (imp_fvs,fvs,ids) (StgVarArg v) = varFVs env v (imp_fvs, fvs, ids)
 
-altFVs :: Env -> StgAlt -> (CgStgAlt, TopFVs, LocalFVs)
+altFVs :: Env -> StgAlt -> (CgStgAlt, ImpFVs, TopFVs, LocalFVs)
 altFVs env GenStgAlt{alt_con=con, alt_bndrs=bndrs, alt_rhs=e}
-  | (e', top_fvs, lcl_fvs) <- exprFVs (addLocals bndrs env) e
+  | (e', imp_fvs, top_fvs, lcl_fvs) <- exprFVs (addLocals bndrs env) e
   , let lcl_fvs' = delDVarSetList lcl_fvs bndrs
   , let newAlt   = GenStgAlt{alt_con=con, alt_bndrs=bndrs, alt_rhs=e'}
-  = (newAlt, top_fvs, lcl_fvs')
+  = (newAlt, imp_fvs, top_fvs, lcl_fvs')
diff --git a/compiler/GHC/Stg/Pipeline.hs b/compiler/GHC/Stg/Pipeline.hs
index 348d8e0e8f02..25dadc6dcb3f 100644
--- a/compiler/GHC/Stg/Pipeline.hs
+++ b/compiler/GHC/Stg/Pipeline.hs
@@ -33,6 +33,7 @@ import GHC.Unit.Module ( Module )
 
 import GHC.Utils.Error
 import GHC.Types.Var
+import GHC.Types.Var.Set
 import GHC.Types.Unique.Supply
 import GHC.Utils.Outputable
 import GHC.Utils.Logger
@@ -70,7 +71,7 @@ stg2stg :: Logger
         -> StgPipelineOpts
         -> Module                    -- ^ module being compiled
         -> [StgTopBinding]           -- ^ input program
-        -> IO ([CgStgTopBinding], StgCgInfos) -- output program
+        -> IO ([(CgStgTopBinding,IdSet)], StgCgInfos) -- output program
 stg2stg logger extra_vars opts this_mod binds
   = do  { dump_when Opt_D_dump_stg_from_core "Initial STG:" binds
         ; showPass logger "Stg2Stg"
@@ -88,9 +89,10 @@ stg2stg logger extra_vars opts this_mod binds
           -- sorting pass is necessary.
           -- 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'
+        ; let (binds_sorted_with_fvs, imp_fvs) = unzip (depSortWithAnnotStgPgm this_mod binds')
         -- See Note [Tag inference for interactive contexts]
-        ; inferTags (stgPipeline_pprOpts opts) (stgPipeline_forBytecode opts) logger this_mod binds_sorted_with_fvs
+        ; (cg_binds, cg_infos) <- inferTags (stgPipeline_pprOpts opts) (stgPipeline_forBytecode opts) logger this_mod binds_sorted_with_fvs
+        ; pure (zip cg_binds imp_fvs, cg_infos)
    }
 
   where
diff --git a/compiler/GHC/Utils/Misc.hs b/compiler/GHC/Utils/Misc.hs
index af32a679bb77..e6ae6db1e1c0 100644
--- a/compiler/GHC/Utils/Misc.hs
+++ b/compiler/GHC/Utils/Misc.hs
@@ -22,7 +22,7 @@ module GHC.Utils.Misc (
         unzipWith,
 
         mapFst, mapSnd, chkAppend,
-        mapAndUnzip, mapAndUnzip3,
+        mapAndUnzip, mapAndUnzip3, mapAndUnzip4,
         filterOut, partitionWith,
 
         dropWhileEndLE, spanEnd, last2, lastMaybe, onJust,
@@ -55,6 +55,7 @@ module GHC.Utils.Misc (
 
         -- * Tuples
         fstOf3, sndOf3, thdOf3,
+        fstOf4, sndOf4,
         fst3, snd3, third3,
         uncurry3,
 
@@ -183,6 +184,11 @@ fstOf3      (a,_,_) =  a
 sndOf3      (_,b,_) =  b
 thdOf3      (_,_,c) =  c
 
+fstOf4   :: (a,b,c,d) -> a
+sndOf4   :: (a,b,c,d) -> b
+fstOf4      (a,_,_,_) =  a
+sndOf4      (_,b,_,_) =  b
+
 fst3 :: (a -> d) -> (a, b, c) -> (d, b, c)
 fst3 f (a, b, c) = (f a, b, c)
 
@@ -324,7 +330,6 @@ mapAndUnzip f (x:xs)
     (r1:rs1, r2:rs2)
 
 mapAndUnzip3 :: (a -> (b, c, d)) -> [a] -> ([b], [c], [d])
-
 mapAndUnzip3 _ [] = ([], [], [])
 mapAndUnzip3 f (x:xs)
   = let (r1,  r2,  r3)  = f x
@@ -332,6 +337,14 @@ mapAndUnzip3 f (x:xs)
     in
     (r1:rs1, r2:rs2, r3:rs3)
 
+mapAndUnzip4 :: (a -> (b, c, d, e)) -> [a] -> ([b], [c], [d], [e])
+mapAndUnzip4 _ [] = ([], [], [], [])
+mapAndUnzip4 f (x:xs)
+  = let (r1,  r2,  r3, r4)  = f x
+        (rs1, rs2, rs3, rs4) = mapAndUnzip4 f xs
+    in
+    (r1:rs1, r2:rs2, r3:rs3, r4:rs4)
+
 zipWithAndUnzip :: (a -> b -> (c,d)) -> [a] -> [b] -> ([c],[d])
 zipWithAndUnzip f (a:as) (b:bs)
   = let (r1,  r2)  = f a b
diff --git a/testsuite/tests/wasm/should_run/control-flow/LoadCmmGroup.hs b/testsuite/tests/wasm/should_run/control-flow/LoadCmmGroup.hs
index f0836762efcb..ba4cbfb50cd1 100644
--- a/testsuite/tests/wasm/should_run/control-flow/LoadCmmGroup.hs
+++ b/testsuite/tests/wasm/should_run/control-flow/LoadCmmGroup.hs
@@ -68,7 +68,7 @@ cmmOfSummary summ = do
   let infotable = emptyInfoTableProvMap
       tycons = []
       ccs = emptyCollectedCCs
-      stg' = depSortWithAnnotStgPgm (ms_mod summ) stg
+      stg' = fmap fst (depSortWithAnnotStgPgm (ms_mod summ) stg)
       hpcinfo = emptyHpcInfo False
       tmpfs = hsc_tmpfs env
       stg_to_cmm dflags mod = codeGen logger tmpfs (initStgToCmmConfig dflags mod)
-- 
GitLab