Add allowVirtualUnits field in PackageState

Instead of always querying DynFlags to know whether we are allowed to
use virtual units (i.e. instantiated on-the-fly, cf Note [About units]
in GHC.Unit), we store it once for all in
`PackageState.allowVirtualUnits`.

This avoids using DynFlags too much (cf #17957) and is preliminary work
for #14335.

10 files changed, 89 insertions, 72 deletions

diff --git a/compiler/GHC/Driver/Backpack.hs b/compiler/GHC/Driver/Backpack.hs
index 177a601425..2c04fb8b37 100644
--- a/compiler/GHC/Driver/Backpack.hs
+++ b/compiler/GHC/Driver/Backpack.hs
@@ -396,13 +396,13 @@ addPackage pkg = do
 compileInclude :: Int -> (Int, Unit) -> BkpM ()
 compileInclude n (i, uid) = do
     hsc_env <- getSession
-    let dflags = hsc_dflags hsc_env
+    let pkgs = pkgState (hsc_dflags hsc_env)
     msgInclude (i, n) uid
     -- Check if we've compiled it already
     case uid of
       HoleUnit   -> return ()
       RealUnit _ -> return ()
-      VirtUnit i -> case lookupUnit dflags uid of
+      VirtUnit i -> case lookupUnit pkgs uid of
         Nothing -> innerBkpM $ compileUnit (instUnitInstanceOf i) (instUnitInsts i)
         Just _  -> return ()
 
diff --git a/compiler/GHC/Driver/CodeOutput.hs b/compiler/GHC/Driver/CodeOutput.hs
index c614606186..db9b331d34 100644
--- a/compiler/GHC/Driver/CodeOutput.hs
+++ b/compiler/GHC/Driver/CodeOutput.hs
@@ -131,7 +131,7 @@ outputC dflags filenm cmm_stream packages
          --   * -#include options from the cmdline and OPTIONS pragmas
          --   * the _stub.h file, if there is one.
          --
-         let rts = unsafeLookupUnit dflags rtsUnitId
+         let rts = unsafeLookupUnit (pkgState dflags) rtsUnitId
 
          let cc_injects = unlines (map mk_include (unitIncludes rts))
              mk_include h_file =
@@ -223,7 +223,7 @@ outputForeignStubs dflags mod location stubs
 
         -- we need the #includes from the rts package for the stub files
         let rts_includes =
-               let rts_pkg = unsafeLookupUnit dflags rtsUnitId in
+               let rts_pkg = unsafeLookupUnit (pkgState dflags) rtsUnitId in
                concatMap mk_include (unitIncludes rts_pkg)
             mk_include i = "#include \"" ++ i ++ "\"\n"
 
diff --git a/compiler/GHC/Driver/Finder.hs b/compiler/GHC/Driver/Finder.hs
index 6a7b9eb3ee..09ef8e2d25 100644
--- a/compiler/GHC/Driver/Finder.hs
+++ b/compiler/GHC/Driver/Finder.hs
@@ -182,14 +182,14 @@ findExposedPackageModule :: HscEnv -> ModuleName -> Maybe FastString
 findExposedPackageModule hsc_env mod_name mb_pkg
   = findLookupResult hsc_env
   $ lookupModuleWithSuggestions
-        (hsc_dflags hsc_env) mod_name mb_pkg
+        (pkgState (hsc_dflags hsc_env)) mod_name mb_pkg
 
 findExposedPluginPackageModule :: HscEnv -> ModuleName
                                -> IO FindResult
 findExposedPluginPackageModule hsc_env mod_name
   = findLookupResult hsc_env
   $ lookupPluginModuleWithSuggestions
-        (hsc_dflags hsc_env) mod_name Nothing
+        (pkgState (hsc_dflags hsc_env)) mod_name Nothing
 
 findLookupResult :: HscEnv -> LookupResult -> IO FindResult
 findLookupResult hsc_env r = case r of
@@ -226,12 +226,15 @@ findLookupResult hsc_env r = case r of
                           , fr_mods_hidden = []
                           , fr_unusables = unusables'
                           , fr_suggestions = [] })
-     LookupNotFound suggest ->
+     LookupNotFound suggest -> do
+       let suggest'
+             | gopt Opt_HelpfulErrors (hsc_dflags hsc_env) = suggest
+             | otherwise = []
        return (NotFound{ fr_paths = [], fr_pkg = Nothing
                        , fr_pkgs_hidden = []
                        , fr_mods_hidden = []
                        , fr_unusables = []
-                       , fr_suggestions = suggest })
+                       , fr_suggestions = suggest' })
 
 modLocationCache :: HscEnv -> InstalledModule -> IO InstalledFindResult -> IO InstalledFindResult
 modLocationCache hsc_env mod do_this = do
@@ -669,6 +672,7 @@ cantFindErr cannot_find _ dflags mod_name find_result
   = ptext cannot_find <+> quotes (ppr mod_name)
     $$ more_info
   where
+    pkgs = pkgState dflags
     more_info
       = case find_result of
             NoPackage pkg
@@ -723,11 +727,11 @@ cantFindErr cannot_find _ dflags mod_name find_result
         <> dot $$ pkg_hidden_hint uid
     pkg_hidden_hint uid
      | gopt Opt_BuildingCabalPackage dflags
-        = let pkg = expectJust "pkg_hidden" (lookupUnit dflags uid)
+        = let pkg = expectJust "pkg_hidden" (lookupUnit pkgs uid)
            in text "Perhaps you need to add" <+>
               quotes (ppr (unitPackageName pkg)) <+>
               text "to the build-depends in your .cabal file."
-     | Just pkg <- lookupUnit dflags uid
+     | Just pkg <- lookupUnit pkgs uid
          = text "You can run" <+>
            quotes (text ":set -package " <> ppr (unitPackageName pkg)) <+>
            text "to expose it." $$
diff --git a/compiler/GHC/Driver/Main.hs b/compiler/GHC/Driver/Main.hs
index 95bad1e615..e5381e188f 100644
--- a/compiler/GHC/Driver/Main.hs
+++ b/compiler/GHC/Driver/Main.hs
@@ -1159,21 +1159,22 @@ hscCheckSafe' m l = do
                     return (trust == Sf_Trustworthy, pkgRs)
 
                 where
+                    state = pkgState dflags
                     inferredImportWarn = unitBag
                         $ makeIntoWarning (Reason Opt_WarnInferredSafeImports)
-                        $ mkWarnMsg dflags l (pkgQual dflags)
+                        $ mkWarnMsg dflags l (pkgQual state)
                         $ sep
                             [ text "Importing Safe-Inferred module "
                                 <> ppr (moduleName m)
                                 <> text " from explicitly Safe module"
                             ]
-                    pkgTrustErr = unitBag $ mkErrMsg dflags l (pkgQual dflags) $
+                    pkgTrustErr = unitBag $ mkErrMsg dflags l (pkgQual state) $
                         sep [ ppr (moduleName m)
                                 <> text ": Can't be safely imported!"
                             , text "The package (" <> ppr (moduleUnit m)
                                 <> text ") the module resides in isn't trusted."
                             ]
-                    modTrustErr = unitBag $ mkErrMsg dflags l (pkgQual dflags) $
+                    modTrustErr = unitBag $ mkErrMsg dflags l (pkgQual state) $
                         sep [ ppr (moduleName m)
                                 <> text ": Can't be safely imported!"
                             , text "The module itself isn't safe." ]
@@ -1192,7 +1193,7 @@ hscCheckSafe' m l = do
     packageTrusted _ Sf_SafeInferred False _ = True
     packageTrusted dflags _ _ m
         | isHomeModule dflags m = True
-        | otherwise = unitIsTrusted $ unsafeLookupUnit dflags (moduleUnit m)
+        | otherwise = unitIsTrusted $ unsafeLookupUnit (pkgState dflags) (moduleUnit m)
 
     lookup' :: Module -> Hsc (Maybe ModIface)
     lookup' m = do
@@ -1215,11 +1216,12 @@ checkPkgTrust :: Set UnitId -> Hsc ()
 checkPkgTrust pkgs = do
     dflags <- getDynFlags
     let errors = S.foldr go [] pkgs
+        state  = pkgState dflags
         go pkg acc
-            | unitIsTrusted $ getInstalledPackageDetails (pkgState dflags) pkg
+            | unitIsTrusted $ getInstalledPackageDetails state pkg
             = acc
             | otherwise
-            = (:acc) $ mkErrMsg dflags noSrcSpan (pkgQual dflags)
+            = (:acc) $ mkErrMsg dflags noSrcSpan (pkgQual state)
                      $ text "The package (" <> ppr pkg <> text ") is required" <>
                        text " to be trusted but it isn't!"
     case errors of
diff --git a/compiler/GHC/Driver/Make.hs b/compiler/GHC/Driver/Make.hs
index 7aceafdd0b..714619d7b2 100644
--- a/compiler/GHC/Driver/Make.hs
+++ b/compiler/GHC/Driver/Make.hs
@@ -307,10 +307,11 @@ warnUnusedPackages = do
     eps <- liftIO $ hscEPS hsc_env
 
     let dflags = hsc_dflags hsc_env
+        state  = pkgState dflags
         pit = eps_PIT eps
 
     let loadedPackages
-          = map (unsafeLookupUnit dflags)
+          = map (unsafeLookupUnit state)
           . nub . sort
           . map moduleUnit
           . moduleEnvKeys
diff --git a/compiler/GHC/Driver/Pipeline.hs b/compiler/GHC/Driver/Pipeline.hs
index fa9527b74e..bbc44a4653 100644
--- a/compiler/GHC/Driver/Pipeline.hs
+++ b/compiler/GHC/Driver/Pipeline.hs
@@ -2002,8 +2002,9 @@ doCpp dflags raw input_fn output_fn = do
     let hsSourceCppOpts = [ "-include", ghcVersionH ]
 
     -- MIN_VERSION macros
-    let uids = explicitPackages (pkgState dflags)
-        pkgs = catMaybes (map (lookupUnit dflags) uids)
+    let state = pkgState dflags
+        uids = explicitPackages state
+        pkgs = catMaybes (map (lookupUnit state) uids)
     mb_macro_include <-
         if not (null pkgs) && gopt Opt_VersionMacros dflags
             then do macro_stub <- newTempName dflags TFL_CurrentModule "h"
diff --git a/compiler/GHC/Driver/Types.hs b/compiler/GHC/Driver/Types.hs
index 5ae44bca21..1b5591793b 100644
--- a/compiler/GHC/Driver/Types.hs
+++ b/compiler/GHC/Driver/Types.hs
@@ -1955,8 +1955,9 @@ with some holes, we should try to give the user some more useful information.
 mkPrintUnqualified :: DynFlags -> GlobalRdrEnv -> PrintUnqualified
 mkPrintUnqualified dflags env = QueryQualify qual_name
                                              (mkQualModule dflags)
-                                             (mkQualPackage dflags)
+                                             (mkQualPackage pkgs)
   where
+  pkgs = pkgState dflags
   qual_name mod occ
         | [gre] <- unqual_gres
         , right_name gre
@@ -2022,32 +2023,30 @@ mkQualModule dflags mod
      = False
 
      | otherwise = True
-     where lookup = lookupModuleInAllPackages dflags (moduleName mod)
+     where lookup = lookupModuleInAllPackages (pkgState dflags) (moduleName mod)
 
 -- | Creates a function for formatting packages based on two heuristics:
 -- (1) don't qualify if the package in question is "main", and (2) only qualify
 -- with a unit id if the package ID would be ambiguous.
-mkQualPackage :: DynFlags -> QueryQualifyPackage
-mkQualPackage dflags uid
+mkQualPackage :: PackageState -> QueryQualifyPackage
+mkQualPackage pkgs uid
      | uid == mainUnitId || uid == interactiveUnitId
         -- Skip the lookup if it's main, since it won't be in the package
         -- database!
      = False
      | Just pkgid <- mb_pkgid
-     , searchPackageId (pkgState dflags) pkgid `lengthIs` 1
+     , searchPackageId pkgs pkgid `lengthIs` 1
         -- this says: we are given a package pkg-0.1@MMM, are there only one
         -- exposed packages whose package ID is pkg-0.1?
      = False
      | otherwise
      = True
-     where mb_pkgid = fmap unitPackageId (lookupUnit dflags uid)
+     where mb_pkgid = fmap unitPackageId (lookupUnit pkgs uid)
 
 -- | A function which only qualifies package names if necessary; but
 -- qualifies all other identifiers.
-pkgQual :: DynFlags -> PrintUnqualified
-pkgQual dflags = alwaysQualify {
-        queryQualifyPackage = mkQualPackage dflags
-    }
+pkgQual :: PackageState -> PrintUnqualified
+pkgQual pkgs = alwaysQualify { queryQualifyPackage = mkQualPackage pkgs }
 
 {-
 ************************************************************************
diff --git a/compiler/GHC/HsToCore/Usage.hs b/compiler/GHC/HsToCore/Usage.hs
index 7a00d75b23..97ab4ba89a 100644
--- a/compiler/GHC/HsToCore/Usage.hs
+++ b/compiler/GHC/HsToCore/Usage.hs
@@ -169,7 +169,7 @@ One way to improve this is to either:
 -}
 mkPluginUsage :: HscEnv -> ModIface -> IO [Usage]
 mkPluginUsage hsc_env pluginModule
-  = case lookupPluginModuleWithSuggestions dflags pNm Nothing of
+  = case lookupPluginModuleWithSuggestions pkgs pNm Nothing of
     LookupFound _ pkg -> do
     -- The plugin is from an external package:
     -- search for the library files containing the plugin.
@@ -215,6 +215,7 @@ mkPluginUsage hsc_env pluginModule
   where
     dflags   = hsc_dflags hsc_env
     platform = targetPlatform dflags
+    pkgs     = pkgState dflags
     pNm      = moduleName $ mi_module pluginModule
     pPkg     = moduleUnit $ mi_module pluginModule
     deps     = map gwib_mod $
diff --git a/compiler/GHC/SysTools/ExtraObj.hs b/compiler/GHC/SysTools/ExtraObj.hs
index b0949ee889..ef04468ebd 100644
--- a/compiler/GHC/SysTools/ExtraObj.hs
+++ b/compiler/GHC/SysTools/ExtraObj.hs
@@ -50,12 +50,14 @@ mkExtraObj dflags extn xs
                     else asmOpts ccInfo)
       return oFile
     where
+      pkgs = pkgState dflags
+
       -- Pass a different set of options to the C compiler depending one whether
       -- we're compiling C or assembler. When compiling C, we pass the usual
       -- set of include directories and PIC flags.
       cOpts = map Option (picCCOpts dflags)
                     ++ map (FileOption "-I")
-                            (unitIncludeDirs $ unsafeLookupUnit dflags rtsUnitId)
+                            (unitIncludeDirs $ unsafeLookupUnit pkgs rtsUnitId)
 
       -- When compiling assembler code, we drop the usual C options, and if the
       -- compiler is Clang, we add an extra argument to tell Clang to ignore
diff --git a/compiler/GHC/Unit/State.hs b/compiler/GHC/Unit/State.hs
index 3ebce53b2a..64c4fdaee2 100644
--- a/compiler/GHC/Unit/State.hs
+++ b/compiler/GHC/Unit/State.hs
@@ -359,7 +359,13 @@ data PackageState = PackageState {
   -- and @r[C=<A>]:C@.
   --
   -- There's an entry in this map for each hole in our home library.
-  requirementContext :: Map ModuleName [InstantiatedModule]
+  requirementContext :: Map ModuleName [InstantiatedModule],
+
+  -- | Indicate if we can instantiate units on-the-fly.
+  --
+  -- This should only be true when we are type-checking an indefinite unit.
+  -- See Note [About units] in GHC.Unit.
+  allowVirtualUnits :: !Bool
   }
 
 emptyPackageState :: PackageState
@@ -371,7 +377,8 @@ emptyPackageState = PackageState {
     explicitPackages = [],
     moduleNameProvidersMap = Map.empty,
     pluginModuleNameProvidersMap = Map.empty,
-    requirementContext = Map.empty
+    requirementContext = Map.empty,
+    allowVirtualUnits = False
     }
 
 -- | Package database
@@ -387,12 +394,12 @@ emptyUnitInfoMap :: UnitInfoMap
 emptyUnitInfoMap = UnitInfoMap emptyUDFM emptyUniqSet
 
 -- | Find the unit we know about with the given unit id, if any
-lookupUnit :: DynFlags -> Unit -> Maybe UnitInfo
-lookupUnit dflags = lookupUnit' (homeUnitIsIndefinite dflags) (unitInfoMap (pkgState dflags))
+lookupUnit :: PackageState -> Unit -> Maybe UnitInfo
+lookupUnit pkgs = lookupUnit' (allowVirtualUnits pkgs) (unitInfoMap pkgs)
 
 -- | A more specialized interface, which takes a boolean specifying
 -- whether or not to look for on-the-fly renamed interfaces, and
--- just a 'UnitInfoMap' rather than a 'DynFlags' (so it can
+-- just a 'UnitInfoMap' rather than a 'PackageState' (so it can
 -- be used while we're initializing 'DynFlags'
 lookupUnit' :: Bool -> UnitInfoMap -> Unit -> Maybe UnitInfo
 lookupUnit' False (UnitInfoMap pkg_map _) uid  = lookupUDFM pkg_map uid
@@ -424,11 +431,11 @@ extendUnitInfoMap (UnitInfoMap pkg_map closure) new_pkgs
 
 -- | Looks up the package with the given id in the package state, panicing if it is
 -- not found
-unsafeLookupUnit :: HasDebugCallStack => DynFlags -> Unit -> UnitInfo
-unsafeLookupUnit dflags pid =
-    case lookupUnit dflags pid of
-      Just config -> config
-      Nothing -> pprPanic "unsafeLookupUnit" (ppr pid)
+unsafeLookupUnit :: HasDebugCallStack => PackageState -> Unit -> UnitInfo
+unsafeLookupUnit pkgs pid =
+    case lookupUnit pkgs pid of
+      Just info -> info
+      Nothing   -> pprPanic "unsafeLookupUnit" (ppr pid)
 
 lookupInstalledPackage :: PackageState -> UnitId -> Maybe UnitInfo
 lookupInstalledPackage pkgstate uid = lookupInstalledPackage' (unitInfoMap pkgstate) uid
@@ -1559,17 +1566,22 @@ mkPackageState dflags dbs preload0 = do
     FormatText
     (pprModuleMap mod_map)
 
-  -- Force pstate to avoid leaking the dflags0 passed to mkPackageState
-  let !pstate = PackageState{
-    preloadPackages     = dep_preload,
-    explicitPackages    = explicit_pkgs,
-    unitInfoMap            = pkg_db,
-    moduleNameProvidersMap  = mod_map,
-    pluginModuleNameProvidersMap = mkModuleNameProvidersMap dflags pkg_db plugin_vis_map,
-    packageNameMap          = pkgname_map,
-    unwireMap = Map.fromList [ (v,k) | (k,v) <- Map.toList wired_map ],
-    requirementContext = req_ctx
-    }
+  -- Force pstate to avoid leaking the dflags passed to mkPackageState
+  let !pstate = PackageState
+         { preloadPackages              = dep_preload
+         , explicitPackages             = explicit_pkgs
+         , unitInfoMap                  = pkg_db
+         , moduleNameProvidersMap       = mod_map
+         , pluginModuleNameProvidersMap = mkModuleNameProvidersMap dflags pkg_db plugin_vis_map
+         , packageNameMap               = pkgname_map
+         , unwireMap                    = Map.fromList [ (v,k) | (k,v) <- Map.toList wired_map ]
+         , requirementContext           = req_ctx
+
+           -- when the home unit is indefinite, it means we are type-checking it
+           -- only (not producing any code). Hence we can use virtual units
+           -- instantiated on-the-fly (see Note [About units] in GHC.Unit)
+         , allowVirtualUnits            = homeUnitIsIndefinite dflags
+         }
   let new_insts = map (fmap (upd_wired_in_mod wired_map)) (homeUnitInstantiations dflags)
   return (pstate, new_dep_preload, new_insts)
 
@@ -1842,14 +1854,14 @@ getPackageFrameworks dflags pkgs = do
 
 -- | Takes a 'ModuleName', and if the module is in any package returns
 -- list of modules which take that name.
-lookupModuleInAllPackages :: DynFlags
+lookupModuleInAllPackages :: PackageState
                           -> ModuleName
                           -> [(Module, UnitInfo)]
-lookupModuleInAllPackages dflags m
-  = case lookupModuleWithSuggestions dflags m Nothing of
+lookupModuleInAllPackages pkgs m
+  = case lookupModuleWithSuggestions pkgs m Nothing of
       LookupFound a b -> [(a,b)]
       LookupMultiple rs -> map f rs
-        where f (m,_) = (m, expectJust "lookupModule" (lookupUnit dflags
+        where f (m,_) = (m, expectJust "lookupModule" (lookupUnit pkgs
                                                          (moduleUnit m)))
       _ -> []
 
@@ -1872,28 +1884,26 @@ data LookupResult =
 data ModuleSuggestion = SuggestVisible ModuleName Module ModuleOrigin
                       | SuggestHidden ModuleName Module ModuleOrigin
 
-lookupModuleWithSuggestions :: DynFlags
+lookupModuleWithSuggestions :: PackageState
                             -> ModuleName
                             -> Maybe FastString
                             -> LookupResult
-lookupModuleWithSuggestions dflags
-  = lookupModuleWithSuggestions' dflags
-        (moduleNameProvidersMap (pkgState dflags))
+lookupModuleWithSuggestions pkgs
+  = lookupModuleWithSuggestions' pkgs (moduleNameProvidersMap pkgs)
 
-lookupPluginModuleWithSuggestions :: DynFlags
+lookupPluginModuleWithSuggestions :: PackageState
                                   -> ModuleName
                                   -> Maybe FastString
                                   -> LookupResult
-lookupPluginModuleWithSuggestions dflags
-  = lookupModuleWithSuggestions' dflags
-        (pluginModuleNameProvidersMap (pkgState dflags))
+lookupPluginModuleWithSuggestions pkgs
+  = lookupModuleWithSuggestions' pkgs (pluginModuleNameProvidersMap pkgs)
 
-lookupModuleWithSuggestions' :: DynFlags
+lookupModuleWithSuggestions' :: PackageState
                             -> ModuleNameProvidersMap
                             -> ModuleName
                             -> Maybe FastString
                             -> LookupResult
-lookupModuleWithSuggestions' dflags mod_map m mb_pn
+lookupModuleWithSuggestions' pkgs mod_map m mb_pn
   = case Map.lookup m mod_map of
         Nothing -> LookupNotFound suggestions
         Just xs ->
@@ -1920,7 +1930,7 @@ lookupModuleWithSuggestions' dflags mod_map m mb_pn
             | otherwise
             -> (x:hidden_pkg, hidden_mod, unusable, exposed)
 
-    unit_lookup p = lookupUnit dflags p `orElse` pprPanic "lookupModuleWithSuggestions" (ppr p <+> ppr m)
+    unit_lookup p = lookupUnit pkgs p `orElse` pprPanic "lookupModuleWithSuggestions" (ppr p <+> ppr m)
     mod_unit = unit_lookup . moduleUnit
 
     -- Filters out origins which are not associated with the given package
@@ -1945,15 +1955,12 @@ lookupModuleWithSuggestions' dflags mod_map m mb_pn
                 }
       where go pkg = pn == fsPackageName pkg
 
-    suggestions
-      | gopt Opt_HelpfulErrors dflags =
-           fuzzyLookup (moduleNameString m) all_mods
-      | otherwise = []
+    suggestions = fuzzyLookup (moduleNameString m) all_mods
 
     all_mods :: [(String, ModuleSuggestion)]     -- All modules
     all_mods = sortBy (comparing fst) $
         [ (moduleNameString m, suggestion)
-        | (m, e) <- Map.toList (moduleNameProvidersMap (pkgState dflags))
+        | (m, e) <- Map.toList (moduleNameProvidersMap pkgs)
         , suggestion <- map (getSuggestion m) (Map.toList e)
         ]
     getSuggestion name (mod, origin) =