Move unit DBs in UnitEnv

Also make the HomeUnit optional to keep the field strict and prepare for
UnitEnvs without a HomeUnit (e.g. in Plugins envs, cf #14335).

9 files changed, 56 insertions, 38 deletions

diff --git a/compiler/GHC.hs b/compiler/GHC.hs
index d7b43caa84..65716d0e95 100644
--- a/compiler/GHC.hs
+++ b/compiler/GHC.hs
@@ -641,7 +641,8 @@ setSessionDynFlags dflags0 = do
   logger <- getLogger
   dflags <- checkNewDynFlags logger dflags0
   hsc_env <- getSession
-  (dbs,unit_state,home_unit) <- liftIO $ initUnits logger dflags (hsc_unit_dbs hsc_env)
+  let cached_unit_dbs = ue_unit_dbs (hsc_unit_env hsc_env)
+  (dbs,unit_state,home_unit) <- liftIO $ initUnits logger dflags cached_unit_dbs
 
   -- Interpreter
   interp <- if gopt Opt_ExternalInterpreter dflags
@@ -682,8 +683,9 @@ setSessionDynFlags dflags0 = do
   let unit_env = UnitEnv
         { ue_platform  = targetPlatform dflags
         , ue_namever   = ghcNameVersion dflags
-        , ue_home_unit = home_unit
+        , ue_home_unit = Just home_unit
         , ue_units     = unit_state
+        , ue_unit_dbs  = Just dbs
         }
   modifySession $ \h -> h{ hsc_dflags = dflags
                          , hsc_IC = (hsc_IC h){ ic_dflags = dflags }
@@ -691,7 +693,6 @@ setSessionDynFlags dflags0 = do
                            -- we only update the interpreter if there wasn't
                            -- already one set up
                          , hsc_unit_env = unit_env
-                         , hsc_unit_dbs = Just dbs
                          }
   invalidateModSummaryCache
 
@@ -713,15 +714,16 @@ setProgramDynFlags_ invalidate_needed dflags = do
   if changed
     then do
         hsc_env <- getSession
-        (dbs,unit_state,home_unit) <- liftIO $ initUnits logger dflags' (hsc_unit_dbs hsc_env)
+        let cached_unit_dbs = ue_unit_dbs (hsc_unit_env hsc_env)
+        (dbs,unit_state,home_unit) <- liftIO $ initUnits logger dflags' cached_unit_dbs
         let unit_env = UnitEnv
               { ue_platform  = targetPlatform dflags'
               , ue_namever   = ghcNameVersion dflags'
-              , ue_home_unit = home_unit
+              , ue_home_unit = Just home_unit
               , ue_units     = unit_state
+              , ue_unit_dbs  = Just dbs
               }
         modifySession $ \h -> h{ hsc_dflags   = dflags'
-                               , hsc_unit_dbs = Just dbs
                                , hsc_unit_env = unit_env
                                }
     else modifySession $ \h -> h{ hsc_dflags = dflags' }
@@ -991,7 +993,7 @@ guessTarget str mUnitId Nothing
 -- of the current 'HomeUnit'.
 unitIdOrHomeUnit :: GhcMonad m => Maybe UnitId -> m UnitId
 unitIdOrHomeUnit mUnitId = do
-  currentHomeUnitId <- homeUnitId . ue_home_unit . hsc_unit_env <$> getSession
+  currentHomeUnitId <- homeUnitId . hsc_home_unit <$> getSession
   pure (fromMaybe currentHomeUnitId mUnitId)
 
 -- | Inform GHC that the working directory has changed.  GHC will flush
diff --git a/compiler/GHC/Driver/Backpack.hs b/compiler/GHC/Driver/Backpack.hs
index 16a7d58448..205ecccc40 100644
--- a/compiler/GHC/Driver/Backpack.hs
+++ b/compiler/GHC/Driver/Backpack.hs
@@ -417,7 +417,7 @@ addUnit :: GhcMonad m => UnitInfo -> m ()
 addUnit u = do
     hsc_env <- getSession
     logger <- getLogger
-    newdbs <- case hsc_unit_dbs hsc_env of
+    newdbs <- case ue_unit_dbs (hsc_unit_env hsc_env) of
         Nothing  -> panic "addUnit: called too early"
         Just dbs ->
          let newdb = UnitDatabase
@@ -429,12 +429,12 @@ addUnit u = do
     let unit_env = UnitEnv
           { ue_platform  = targetPlatform (hsc_dflags hsc_env)
           , ue_namever   = ghcNameVersion (hsc_dflags hsc_env)
-          , ue_home_unit = home_unit
+          , ue_home_unit = Just home_unit
           , ue_units     = unit_state
+          , ue_unit_dbs  = Just dbs
           }
     setSession $ hsc_env
-        { hsc_unit_dbs = Just dbs
-        , hsc_unit_env = unit_env
+        { hsc_unit_env = unit_env
         }
 
 compileInclude :: Int -> (Int, Unit) -> BkpM ()
diff --git a/compiler/GHC/Driver/Env.hs b/compiler/GHC/Driver/Env.hs
index 4bf2d9b72e..105c0a64a0 100644
--- a/compiler/GHC/Driver/Env.hs
+++ b/compiler/GHC/Driver/Env.hs
@@ -85,7 +85,7 @@ runInteractiveHsc :: HscEnv -> Hsc a -> IO a
 runInteractiveHsc hsc_env = runHsc (mkInteractiveHscEnv hsc_env)
 
 hsc_home_unit :: HscEnv -> HomeUnit
-hsc_home_unit = ue_home_unit . hsc_unit_env
+hsc_home_unit = unsafeGetHomeUnit . hsc_unit_env
 
 hsc_units :: HscEnv -> UnitState
 hsc_units = ue_units . hsc_unit_env
diff --git a/compiler/GHC/Driver/Env/Types.hs b/compiler/GHC/Driver/Env/Types.hs
index 402366894d..ff387f1d1e 100644
--- a/compiler/GHC/Driver/Env/Types.hs
+++ b/compiler/GHC/Driver/Env/Types.hs
@@ -18,7 +18,6 @@ import GHC.Unit.Finder.Types
 import GHC.Unit.Home.ModInfo
 import GHC.Unit.Module.Graph
 import GHC.Unit.Env
-import GHC.Unit.State
 import GHC.Unit.Types
 import GHC.Utils.Logger
 import GHC.Utils.TmpFs
@@ -134,16 +133,6 @@ data HscEnv
                 -- To add dynamically loaded plugins through the GHC API see
                 -- 'addPluginModuleName' instead.
 
-        , hsc_unit_dbs :: !(Maybe [UnitDatabase UnitId])
-                -- ^ Stack of unit databases for the target platform.
-                --
-                -- This field is populated with the result of `initUnits`.
-                --
-                -- 'Nothing' means the databases have never been read from disk.
-                --
-                -- Usually we don't reload the databases from disk if they are
-                -- cached, even if the database flags changed!
-
         , hsc_unit_env :: UnitEnv
                 -- ^ Unit environment (unit state, home unit, etc.).
                 --
diff --git a/compiler/GHC/Driver/Main.hs b/compiler/GHC/Driver/Main.hs
index 1650de05a7..7ea0cd6ae0 100644
--- a/compiler/GHC/Driver/Main.hs
+++ b/compiler/GHC/Driver/Main.hs
@@ -266,7 +266,6 @@ newHscEnv dflags = do
                   ,  hsc_unit_env       = unit_env
                   ,  hsc_plugins        = []
                   ,  hsc_static_plugins = []
-                  ,  hsc_unit_dbs       = Nothing
                   ,  hsc_hooks          = emptyHooks
                   ,  hsc_tmpfs          = tmpfs
                   }
diff --git a/compiler/GHC/Driver/Pipeline.hs b/compiler/GHC/Driver/Pipeline.hs
index fc6fe68281..191e802e02 100644
--- a/compiler/GHC/Driver/Pipeline.hs
+++ b/compiler/GHC/Driver/Pipeline.hs
@@ -822,12 +822,12 @@ runPipeline stop_phase hsc_env0 (input_fn, mb_input_buf, mb_phase)
                        let unit_env = UnitEnv
                              { ue_platform  = targetPlatform dflags'
                              , ue_namever   = ghcNameVersion dflags'
-                             , ue_home_unit = home_unit
+                             , ue_home_unit = Just home_unit
                              , ue_units     = unit_state
+                             , ue_unit_dbs  = Just dbs
                              }
                        let hsc_env'' = hsc_env'
                             { hsc_unit_env = unit_env
-                            , hsc_unit_dbs = Just dbs
                             }
                        _ <- runPipeline' start_phase hsc_env'' env input_fn'
                                          maybe_loc foreign_os
diff --git a/compiler/GHC/Iface/Load.hs b/compiler/GHC/Iface/Load.hs
index 8cd93e058f..2cd2e15819 100644
--- a/compiler/GHC/Iface/Load.hs
+++ b/compiler/GHC/Iface/Load.hs
@@ -1414,7 +1414,7 @@ cantFindErr using_cabal cannot_find _ unit_env profile tried_these mod_name find
   = ptext cannot_find <+> quotes (ppr mod_name)
     $$ more_info
   where
-    home_unit  = ue_home_unit unit_env
+    mhome_unit = ue_home_unit unit_env
     more_info
       = case find_result of
             NoPackage pkg
@@ -1424,7 +1424,13 @@ cantFindErr using_cabal cannot_find _ unit_env profile tried_these mod_name find
             NotFound { fr_paths = files, fr_pkg = mb_pkg
                      , fr_mods_hidden = mod_hiddens, fr_pkgs_hidden = pkg_hiddens
                      , fr_unusables = unusables, fr_suggestions = suggest }
-                | Just pkg <- mb_pkg, not (isHomeUnit home_unit pkg)
+                | Just pkg <- mb_pkg
+                , Nothing <- mhome_unit           -- no home-unit
+                -> not_found_in_package pkg files
+
+                | Just pkg <- mb_pkg
+                , Just home_unit <- mhome_unit    -- there is a home-unit but the
+                , not (isHomeUnit home_unit pkg)  -- module isn't from it
                 -> not_found_in_package pkg files
 
                 | not (null suggest)
diff --git a/compiler/GHC/Types/Name/Ppr.hs b/compiler/GHC/Types/Name/Ppr.hs
index e48f39576e..14fb5670e1 100644
--- a/compiler/GHC/Types/Name/Ppr.hs
+++ b/compiler/GHC/Types/Name/Ppr.hs
@@ -132,9 +132,10 @@ mkPrintUnqualified unit_env env
 -- | Creates a function for formatting modules based on two heuristics:
 -- (1) if the module is the current module, don't qualify, and (2) if there
 -- is only one exposed package which exports this module, don't qualify.
-mkQualModule :: UnitState -> HomeUnit -> QueryQualifyModule
-mkQualModule unit_state home_unit mod
-     | isHomeModule home_unit mod = False
+mkQualModule :: UnitState -> Maybe HomeUnit -> QueryQualifyModule
+mkQualModule unit_state mhome_unit mod
+     | Just home_unit <- mhome_unit
+     , isHomeModule home_unit mod = False
 
      | [(_, pkgconfig)] <- lookup,
        mkUnit pkgconfig == moduleUnit mod
diff --git a/compiler/GHC/Unit/Env.hs b/compiler/GHC/Unit/Env.hs
index 565c6a8a8e..89e8d77586 100644
--- a/compiler/GHC/Unit/Env.hs
+++ b/compiler/GHC/Unit/Env.hs
@@ -1,6 +1,7 @@
 module GHC.Unit.Env
     ( UnitEnv (..)
     , initUnitEnv
+    , unsafeGetHomeUnit
     , preloadUnitsInfo
     , preloadUnitsInfo'
     )
@@ -21,7 +22,17 @@ data UnitEnv = UnitEnv
     { ue_units     :: !UnitState
         -- ^ External units
 
-    , ue_home_unit :: !HomeUnit
+    , ue_unit_dbs :: !(Maybe [UnitDatabase UnitId])
+        -- ^ Stack of unit databases for the target platform.
+        --
+        -- This field is populated with the result of `initUnits`.
+        --
+        -- 'Nothing' means the databases have never been read from disk.
+        --
+        -- Usually we don't reload the databases from disk if they are
+        -- cached, even if the database flags changed!
+
+    , ue_home_unit :: !(Maybe HomeUnit)
         -- ^ Home unit
 
     , ue_platform  :: !Platform
@@ -35,11 +46,20 @@ initUnitEnv :: GhcNameVersion -> Platform -> IO UnitEnv
 initUnitEnv namever platform = do
   return $ UnitEnv
     { ue_units     = emptyUnitState
-    , ue_home_unit = panic "No home unit"
+    , ue_unit_dbs  = Nothing
+    , ue_home_unit = Nothing
     , ue_platform  = platform
     , ue_namever   = namever
     }
 
+-- | Get home-unit
+--
+-- Unsafe because the home-unit may not be set
+unsafeGetHomeUnit :: UnitEnv -> HomeUnit
+unsafeGetHomeUnit ue = case ue_home_unit ue of
+  Nothing -> panic "unsafeGetHomeUnit: No home unit"
+  Just h  -> h
+
 -- -----------------------------------------------------------------------------
 -- Extracting information from the packages in scope
 
@@ -57,15 +77,16 @@ initUnitEnv namever platform = do
 preloadUnitsInfo' :: UnitEnv -> [UnitId] -> MaybeErr UnitErr [UnitInfo]
 preloadUnitsInfo' unit_env ids0 = all_infos
   where
-    home_unit  = ue_home_unit unit_env
-    unit_state = ue_units     unit_env
+    unit_state = ue_units unit_env
     ids      = ids0 ++ inst_ids
-    inst_ids
+    inst_ids = case ue_home_unit unit_env of
+      Nothing -> []
+      Just home_unit
        -- An indefinite package will have insts to HOLE,
        -- which is not a real package. Don't look it up.
        -- Fixes #14525
-       | isHomeUnitIndefinite home_unit = []
-       | otherwise = map (toUnitId . moduleUnit . snd) (homeUnitInstantiations home_unit)
+       | isHomeUnitIndefinite home_unit -> []
+       | otherwise -> map (toUnitId . moduleUnit . snd) (homeUnitInstantiations home_unit)
     pkg_map = unitInfoMap unit_state
     preload = preloadUnits unit_state