gistfile1.txt

commit a27803fa7ab18582f61558aaf3799d94726a0a1f
Author: Matthew Pickering <[email protected]>
Date:   Wed Apr 2 14:16:16 2025 +0100

    Move ModuleGraph into UnitEnv
    
    The ModuleGraph is a piece of information associated with the
    ExternalPackageState and HomeUnitGraph. Therefore we should store it
    inside the HomeUnitEnv.

diff --git a/compiler/GHC.hs b/compiler/GHC.hs
index 95a75846f89..9ab326bf408 100644
--- a/compiler/GHC.hs
+++ b/compiler/GHC.hs
@@ -859,6 +859,7 @@ setProgramDynFlags_ invalidate_needed dflags = do
               , ue_namever         = ghcNameVersion dflags1
               , ue_home_unit_graph = home_unit_graph
               , ue_current_unit    = ue_currentUnit old_unit_env
+              , ue_module_graph    = ue_module_graph old_unit_env
               , ue_eps             = ue_eps old_unit_env
               }
         modifySession $ \h -> hscSetFlags dflags1 h{ hsc_unit_env = unit_env }
@@ -916,6 +917,7 @@ setProgramHUG_ invalidate_needed new_hug0 = do
             , ue_home_unit_graph = home_unit_graph
             , ue_current_unit    = ue_currentUnit unit_env0
             , ue_eps             = ue_eps unit_env0
+            , ue_module_graph    = ue_module_graph unit_env0
             }
       modifySession $ \h ->
         -- hscSetFlags takes care of updating the logger as well.
@@ -996,7 +998,7 @@ setProgramHUG_ invalidate_needed new_hug0 = do
 --
 invalidateModSummaryCache :: GhcMonad m => m ()
 invalidateModSummaryCache =
-  modifySession $ \h -> h { hsc_mod_graph = mapMG inval (hsc_mod_graph h) }
+  modifySession $ \hsc_env -> setModuleGraph (mapMG inval (hsc_mod_graph hsc_env)) hsc_env
  where
   inval ms = ms { ms_hs_hash = fingerprint0 }
 
diff --git a/compiler/GHC/Core/Opt/Pipeline.hs b/compiler/GHC/Core/Opt/Pipeline.hs
index 314931b5821..038c7ab1ab7 100644
--- a/compiler/GHC/Core/Opt/Pipeline.hs
+++ b/compiler/GHC/Core/Opt/Pipeline.hs
@@ -97,10 +97,11 @@ core2core hsc_env guts@(ModGuts { mg_module  = mod
   where
     dflags         = hsc_dflags hsc_env
     logger         = hsc_logger hsc_env
+    unit_env       = hsc_unit_env hsc_env
     extra_vars     = interactiveInScope (hsc_IC hsc_env)
     home_pkg_rules = hugRulesBelow hsc_env (moduleUnitId mod)
                       (GWIB { gwib_mod = moduleName mod, gwib_isBoot = NotBoot })
-    name_ppr_ctx   = mkNamePprCtx ptc (hsc_unit_env hsc_env) rdr_env
+    name_ppr_ctx   = mkNamePprCtx ptc unit_env rdr_env
     ptc            = initPromotionTickContext dflags
     -- mod: get the module out of the current HscEnv so we can retrieve it from the monad.
     -- This is very convienent for the users of the monad (e.g. plugins do not have to
diff --git a/compiler/GHC/Driver/Backpack.hs b/compiler/GHC/Driver/Backpack.hs
index ddbd21fd19d..9ca6ee734d7 100644
--- a/compiler/GHC/Driver/Backpack.hs
+++ b/compiler/GHC/Driver/Backpack.hs
@@ -457,6 +457,7 @@ addUnit u = do
                     (homeUnitId home_unit)
                     (HUG.mkHomeUnitEnv unit_state (Just dbs) dflags (ue_hpt old_unit_env) (Just home_unit))
           , ue_eps       = ue_eps old_unit_env
+          , ue_module_graph = ue_module_graph old_unit_env
           }
     setSession $ hscSetFlags dflags $ hsc_env { hsc_unit_env = unit_env }
 
diff --git a/compiler/GHC/Driver/Env.hs b/compiler/GHC/Driver/Env.hs
index a3ec1e646f4..1357d86b0da 100644
--- a/compiler/GHC/Driver/Env.hs
+++ b/compiler/GHC/Driver/Env.hs
@@ -2,6 +2,8 @@
 module GHC.Driver.Env
    ( Hsc(..)
    , HscEnv (..)
+   , hsc_mod_graph
+   , setModuleGraph
    , hscUpdateFlags
    , hscSetFlags
    , hsc_home_unit
@@ -130,6 +132,9 @@ hsc_HUE = ue_currentHomeUnitEnv . hsc_unit_env
 hsc_HUG :: HscEnv -> HomeUnitGraph
 hsc_HUG = ue_home_unit_graph . hsc_unit_env
 
+hsc_mod_graph :: HscEnv -> ModuleGraph
+hsc_mod_graph = ue_module_graph . hsc_unit_env
+
 hsc_all_home_unit_ids :: HscEnv -> Set.Set UnitId
 hsc_all_home_unit_ids = HUG.allUnits . hsc_HUG
 
@@ -139,6 +144,9 @@ hscInsertHPT hmi hsc_env = UnitEnv.insertHpt hmi (hsc_unit_env hsc_env)
 hscUpdateHUG :: (HomeUnitGraph -> HomeUnitGraph) -> HscEnv -> HscEnv
 hscUpdateHUG f hsc_env = hsc_env { hsc_unit_env = updateHug f (hsc_unit_env hsc_env) }
 
+setModuleGraph :: ModuleGraph -> HscEnv -> HscEnv
+setModuleGraph mod_graph hsc_env = hsc_env { hsc_unit_env = (hsc_unit_env hsc_env) { ue_module_graph = mod_graph } }
+
 {-
 
 Note [Target code interpreter]
@@ -220,15 +228,15 @@ hscEPS hsc_env = readIORef (euc_eps (ue_eps (hsc_unit_env hsc_env)))
 -- | Find all rules in modules that are in the transitive closure of the given
 -- module.
 hugRulesBelow :: HscEnv -> UnitId -> ModuleNameWithIsBoot -> IO RuleBase
-hugRulesBelow hsc uid mn = foldr (flip extendRuleBaseList) emptyRuleBase <$>
-  hugSomeThingsBelowUs (md_rules . hm_details) False hsc uid mn
+hugRulesBelow hsc_env uid mn = foldr (flip extendRuleBaseList) emptyRuleBase <$>
+  hugSomeThingsBelowUs (md_rules . hm_details) False hsc_env uid mn
 
 -- | Get annotations from all modules "below" this one (in the dependency
 -- sense) within the home units. If the module is @Nothing@, returns /all/
 -- annotations in the home units.
 hugAnnsBelow :: HscEnv -> UnitId -> ModuleNameWithIsBoot -> IO AnnEnv
-hugAnnsBelow hsc uid mn = foldr (flip extendAnnEnvList) emptyAnnEnv <$>
-  hugSomeThingsBelowUs (md_anns . hm_details) False hsc uid mn
+hugAnnsBelow hsc_env uid mn = foldr (flip extendAnnEnvList) emptyAnnEnv <$>
+  hugSomeThingsBelowUs (md_anns . hm_details) False hsc_env uid mn
 
 -- | Find all COMPLETE pragmas in modules that are in the transitive closure of the
 -- given module.
@@ -260,7 +268,8 @@ hugInstancesBelow hsc_env uid mnwib = do
 hugSomeThingsBelowUs :: (HomeModInfo -> [a]) -> Bool -> HscEnv -> UnitId -> ModuleNameWithIsBoot -> IO [[a]]
 -- An explicit check to see if we are in one-shot mode to avoid poking the ModuleGraph thunk
 -- These things are currently stored in the EPS for home packages. (See #25795 for
--- progress in removing these kind of checks)
+-- progress in removing these kind of checks; and making these functions of
+-- `UnitEnv` rather than `HscEnv`)
 -- See Note [Downsweep and the ModuleGraph]
 hugSomeThingsBelowUs _ _ hsc_env _ _ | isOneShot (ghcMode (hsc_dflags hsc_env)) = return []
 hugSomeThingsBelowUs extract include_hi_boot hsc_env uid mn
diff --git a/compiler/GHC/Driver/Env/Types.hs b/compiler/GHC/Driver/Env/Types.hs
index 2bcdb833c3d..d70b56f529f 100644
--- a/compiler/GHC/Driver/Env/Types.hs
+++ b/compiler/GHC/Driver/Env/Types.hs
@@ -18,7 +18,6 @@ import GHC.Types.Name.Cache
 import GHC.Types.Target
 import GHC.Types.TypeEnv
 import GHC.Unit.Finder.Types
-import GHC.Unit.Module.Graph
 import GHC.Unit.Env
 import GHC.Utils.Logger
 import GHC.Utils.TmpFs
@@ -65,10 +64,6 @@ data HscEnv
         hsc_targets :: [Target],
                 -- ^ The targets (or roots) of the current session
 
-        hsc_mod_graph :: ModuleGraph,
-                -- ^ The module graph of the current session
-                -- See Note [Downsweep and the ModuleGraph] for when this is constructed.
-
         hsc_IC :: InteractiveContext,
                 -- ^ The context for evaluating interactive statements
 
@@ -113,3 +108,4 @@ data HscEnv
         , hsc_llvm_config :: !LlvmConfigCache
                 -- ^ LLVM configuration cache.
  }
+
diff --git a/compiler/GHC/Driver/Main.hs b/compiler/GHC/Driver/Main.hs
index a81afa0803b..6b1244234ad 100644
--- a/compiler/GHC/Driver/Main.hs
+++ b/compiler/GHC/Driver/Main.hs
@@ -332,7 +332,6 @@ newHscEnvWithHUG top_dir top_dynflags cur_unit home_unit_graph = do
     return HscEnv { hsc_dflags         = top_dynflags
                   , hsc_logger         = setLogFlags logger (initLogFlags top_dynflags)
                   , hsc_targets        = []
-                  , hsc_mod_graph      = emptyMG
                   , hsc_IC             = emptyInteractiveContext dflags
                   , hsc_NC             = nc_var
                   , hsc_FC             = fc_var
diff --git a/compiler/GHC/Driver/Make.hs b/compiler/GHC/Driver/Make.hs
index da026b01400..6f886719c56 100644
--- a/compiler/GHC/Driver/Make.hs
+++ b/compiler/GHC/Driver/Make.hs
@@ -190,12 +190,12 @@ depanalE diag_wrapper msg excluded_mods allow_dup_roots = do
 
         all_errs <- liftIO $ HUG.unitEnv_foldWithKey one_unit_messages (return emptyMessages) (hsc_HUG hsc_env)
         logDiagnostics (GhcDriverMessage <$> all_errs)
-        setSession hsc_env { hsc_mod_graph = mod_graph }
+        setSession (setModuleGraph mod_graph hsc_env)
         pure (emptyMessages, mod_graph)
       else do
         -- We don't have a complete module dependency graph,
         -- The graph may be disconnected and is unusable.
-        setSession hsc_env { hsc_mod_graph = emptyMG }
+        setSession (setModuleGraph emptyMG hsc_env)
         pure (errs, emptyMG)
 
 
@@ -616,7 +616,7 @@ load' mhmi_cache how_much diag_wrapper mHscMessage mod_graph = do
     -- for any client who might interact with GHC via load'.
     -- See Note [Timing of plugin initialization]
     initializeSessionPlugins
-    modifySession $ \hsc_env -> hsc_env { hsc_mod_graph = mod_graph }
+    modifySession (setModuleGraph mod_graph)
     guessOutputFile
     hsc_env <- getSession
 
diff --git a/compiler/GHC/Driver/Pipeline/Execute.hs b/compiler/GHC/Driver/Pipeline/Execute.hs
index ec04de51c7d..ce1634512b1 100644
--- a/compiler/GHC/Driver/Pipeline/Execute.hs
+++ b/compiler/GHC/Driver/Pipeline/Execute.hs
@@ -768,8 +768,9 @@ runHscPhase pipe_env hsc_env0 input_fn src_flavour = do
   -- files. See GHC.Tc.Utils.TcGblEnv.tcg_type_env_var.
   -- See also Note [hsc_type_env_var hack]
   type_env_var <- newIORef emptyNameEnv
-  let hsc_env' = hsc_env { hsc_type_env_vars = knotVarsFromModuleEnv (mkModuleEnv [(mod, type_env_var)])
-                         , hsc_mod_graph = mg }
+  let hsc_env' =
+        setModuleGraph mg
+          hsc_env { hsc_type_env_vars = knotVarsFromModuleEnv (mkModuleEnv [(mod, type_env_var)]) }
 
 
 
diff --git a/compiler/GHC/Iface/Load.hs b/compiler/GHC/Iface/Load.hs
index 3101bb66463..a940af4d513 100644
--- a/compiler/GHC/Iface/Load.hs
+++ b/compiler/GHC/Iface/Load.hs
@@ -671,7 +671,7 @@ dontLeakTheHUG thing_inside = do
          -- oneshot mode does not support backpack
          -- and we want to avoid prodding the hsc_mod_graph thunk
          | isOneShot (ghcMode (hsc_dflags hsc_env)) = False
-         | mgHasHoles (hsc_mod_graph hsc_env) = True
+         | mgHasHoles (ue_module_graph old_unit_env) = True
          | otherwise = False
         pruneHomeUnitEnv hme = do
           -- NB: These are empty HPTs because Iface/Load first consults the HPT
@@ -683,19 +683,19 @@ dontLeakTheHUG thing_inside = do
           | otherwise
           = do
             hug' <- traverse pruneHomeUnitEnv (ue_home_unit_graph old_unit_env)
+            let !new_mod_graph = emptyMG { mg_mss = panic "cleanTopEnv: mg_mss"
+                                         , mg_graph = panic "cleanTopEnv: mg_graph"
+                                         , mg_has_holes = keepFor20509 }
             return old_unit_env
               { ue_home_unit_graph = hug'
+              , ue_module_graph    = new_mod_graph
               }
       in do
         !unit_env <- unit_env_io
         -- mg_has_holes will be checked again, but nothing else about the module graph
-        let !new_mod_graph = emptyMG { mg_mss = panic "cleanTopEnv: mg_mss"
-                                     , mg_graph = panic "cleanTopEnv: mg_graph"
-                                     , mg_has_holes = keepFor20509 }
         pure $
           hsc_env
                 {  hsc_targets      = panic "cleanTopEnv: hsc_targets"
-                ,  hsc_mod_graph    = new_mod_graph
                 ,  hsc_IC           = panic "cleanTopEnv: hsc_IC"
                 ,  hsc_type_env_vars = case maybe_type_vars of
                                            Just vars -> vars
diff --git a/compiler/GHC/Tc/Module.hs b/compiler/GHC/Tc/Module.hs
index a25c4236c76..13eaa6a9dae 100644
--- a/compiler/GHC/Tc/Module.hs
+++ b/compiler/GHC/Tc/Module.hs
@@ -2109,7 +2109,7 @@ for the unit portion of the graph, if it's not already been performed.
 withInteractiveModuleNode :: HscEnv -> TcM a -> TcM a
 withInteractiveModuleNode hsc_env thing_inside = do
   mg <- liftIO $ downsweepInteractiveImports hsc_env (hsc_IC hsc_env)
-  updTopEnv (\env -> env { hsc_mod_graph = mg }) thing_inside
+  updTopEnv (setModuleGraph mg) thing_inside
 
 
 runTcInteractive :: HscEnv -> TcRn a -> IO (Messages TcRnMessage, Maybe a)
diff --git a/compiler/GHC/Unit/Env.hs b/compiler/GHC/Unit/Env.hs
index 0aa09d795a5..8d96a1ad9c0 100644
--- a/compiler/GHC/Unit/Env.hs
+++ b/compiler/GHC/Unit/Env.hs
@@ -23,21 +23,22 @@
 -- ┌▽────────────┐    │                     │
 -- │HomeUnitGraph│    │                     │
 -- └┬────────────┘    │                     │
--- ┌▽─────────────────▽┐                    │
--- │UnitEnv            │                    │
--- └┬──────────────────┘                    │
--- ┌▽───────────────────────────────────────▽┐
--- │HscEnv                                   │
--- └─────────────────────────────────────────┘
+-- ┌▽─────────────────▽─────────────────────▽┐
+-- │UnitEnv                                  │
+-- └┬─────────────-──────────────────────────┘
+--  │
+--  │
+-- ┌▽──────────────────────────────────────▽┐
+-- │HscEnv                                  │
+-- └────────────────────────────────────────┘
 -- @
 --
--- The 'UnitEnv' references both the 'HomeUnitGraph' (with all the home unit
--- modules) and the 'ExternalPackageState' (information about all
--- non-home/external units). The 'HscEnv' references this 'UnitEnv' and the
--- 'ModuleGraph' (which describes the relationship between the modules being
--- compiled). The 'HomeUnitGraph' has one 'HomePackageTable' for every unit.
---
--- TODO: Arguably, the 'ModuleGraph' should be part of 'UnitEnv' rather than being in the 'HscEnv'.
+-- The 'UnitEnv' references the 'HomeUnitGraph' (with all the home unit
+-- modules), the 'ExternalPackageState' (information about all
+-- non-home/external units), and the 'ModuleGraph' (which describes the
+-- relationship between the modules being compiled).
+-- The 'HscEnv' references this 'UnitEnv'.
+-- The 'HomeUnitGraph' has one 'HomePackageTable' for every unit.
 module GHC.Unit.Env
     ( UnitEnv (..)
     , initUnitEnv
@@ -119,6 +120,7 @@ import GHC.Unit.Home.ModInfo
 import GHC.Unit.Home.PackageTable
 import GHC.Unit.Home.Graph (HomeUnitGraph, HomeUnitEnv)
 import qualified GHC.Unit.Home.Graph as HUG
+import GHC.Unit.Module.Graph
 
 import GHC.Platform
 import GHC.Settings
@@ -163,6 +165,10 @@ data UnitEnv = UnitEnv
 
     , ue_current_unit    :: UnitId
 
+    , ue_module_graph    :: !ModuleGraph
+        -- ^ The module graph of the current session
+        -- See Note [Downsweep and the ModuleGraph] for when this is constructed.
+
     , ue_home_unit_graph :: !HomeUnitGraph
         -- See Note [Multiple Home Units]
 
@@ -182,6 +188,7 @@ initUnitEnv cur_unit hug namever platform = do
   return $ UnitEnv
     { ue_eps             = eps
     , ue_home_unit_graph = hug
+    , ue_module_graph    = emptyMG
     , ue_current_unit    = cur_unit
     , ue_platform        = platform
     , ue_namever         = namever
diff --git a/ghc/GHCi/UI.hs b/ghc/GHCi/UI.hs
index e609260d55f..3357b6476e8 100644
--- a/ghc/GHCi/UI.hs
+++ b/ghc/GHCi/UI.hs
@@ -4680,7 +4680,7 @@ clearHPTs = do
   let pruneHomeUnitEnv hme = liftIO $ do
         emptyHpt <- emptyHomePackageTable
         pure  hme{ homeUnitEnv_hpt = emptyHpt }
-      discardMG hsc = hsc { hsc_mod_graph = GHC.emptyMG }
+      discardMG hsc = setModuleGraph GHC.emptyMG hsc
   modifySessionM $ \hsc_env -> do
     hug' <- traverse pruneHomeUnitEnv $ hsc_HUG hsc_env
     pure $ discardMG $ discardIC $ hscUpdateHUG (const hug') hsc_env