package imports: Take into account package visibility when renaming

In 806e49ae the package imports refactoring code was modified to rename
package imports. There was a small oversight which meant the code didn't
account for module visibility. This patch fixes that oversight.

In general the "lookupPackageName" function is unsafe to use as it
doesn't account for package visiblity/thinning/renaming etc, there is
just one use in the compiler which would be good to audit.

Fixes #20779

7 files changed, 36 insertions, 19 deletions

diff --git a/compiler/GHC.hs b/compiler/GHC.hs
index a8e02e60c0..d70ca74d25 100644
--- a/compiler/GHC.hs
+++ b/compiler/GHC.hs
@@ -1637,7 +1637,7 @@ showRichTokenStream ts = go startLoc ts ""
 -- using the algorithm that is used for an @import@ declaration.
 findModule :: GhcMonad m => ModuleName -> Maybe FastString -> m Module
 findModule mod_name maybe_pkg = do
-  pkg_qual <- renamePkgQualM maybe_pkg
+  pkg_qual <- renamePkgQualM mod_name maybe_pkg
   findQualifiedModule pkg_qual mod_name
 
 
@@ -1673,11 +1673,11 @@ modNotLoadedError dflags m loc = throwGhcExceptionIO $ CmdLineError $ showSDoc d
    quotes (ppr (moduleName m)) <+>
    parens (text (expectJust "modNotLoadedError" (ml_hs_file loc)))
 
-renamePkgQualM :: GhcMonad m => Maybe FastString -> m PkgQual
-renamePkgQualM p = withSession $ \hsc_env -> pure (renamePkgQual (hsc_unit_env hsc_env) p)
+renamePkgQualM :: GhcMonad m => ModuleName -> Maybe FastString -> m PkgQual
+renamePkgQualM mn p = withSession $ \hsc_env -> pure (renamePkgQual (hsc_unit_env hsc_env) mn p)
 
-renameRawPkgQualM :: GhcMonad m => RawPkgQual -> m PkgQual
-renameRawPkgQualM p = withSession $ \hsc_env -> pure (renameRawPkgQual (hsc_unit_env hsc_env) p)
+renameRawPkgQualM :: GhcMonad m => ModuleName -> RawPkgQual -> m PkgQual
+renameRawPkgQualM mn p = withSession $ \hsc_env -> pure (renameRawPkgQual (hsc_unit_env hsc_env) mn p)
 
 -- | Like 'findModule', but differs slightly when the module refers to
 -- a source file, and the file has not been loaded via 'load'.  In
@@ -1688,7 +1688,7 @@ renameRawPkgQualM p = withSession $ \hsc_env -> pure (renameRawPkgQual (hsc_unit
 --
 lookupModule :: GhcMonad m => ModuleName -> Maybe FastString -> m Module
 lookupModule mod_name maybe_pkg = do
-  pkgqual <- renamePkgQualM maybe_pkg
+  pkgqual <- renamePkgQualM mod_name maybe_pkg
   lookupQualifiedModule pkgqual mod_name
 
 lookupQualifiedModule :: GhcMonad m => PkgQual -> ModuleName -> m Module
diff --git a/compiler/GHC/Driver/Backpack.hs b/compiler/GHC/Driver/Backpack.hs
index 1da05dbda8..8ca120e462 100644
--- a/compiler/GHC/Driver/Backpack.hs
+++ b/compiler/GHC/Driver/Backpack.hs
@@ -863,7 +863,7 @@ hsModuleToModSummary pn hsc_src modname
         implicit_imports = mkPrelImports modname loc
                                          implicit_prelude imps
 
-        rn_pkg_qual = renameRawPkgQual (hsc_unit_env hsc_env)
+        rn_pkg_qual = renameRawPkgQual (hsc_unit_env hsc_env) modname
         convImport (L _ i) = (rn_pkg_qual (ideclPkgQual i), reLoc $ ideclName i)
 
     extra_sig_imports <- liftIO $ findExtraSigImports hsc_env hsc_src modname
diff --git a/compiler/GHC/Driver/Make.hs b/compiler/GHC/Driver/Make.hs
index 209a6a9e76..3a37a06809 100644
--- a/compiler/GHC/Driver/Make.hs
+++ b/compiler/GHC/Driver/Make.hs
@@ -1955,7 +1955,7 @@ getPreprocessedImports hsc_env src_fn mb_phase maybe_buf = do
           mimps <- getImports popts imp_prelude pi_hspp_buf pi_hspp_fn src_fn
           return (first (mkMessages . fmap mkDriverPsHeaderMessage . getMessages) mimps)
   let rn_pkg_qual = renameRawPkgQual (hsc_unit_env hsc_env)
-  let rn_imps = fmap (first rn_pkg_qual)
+  let rn_imps = fmap (\(pk, lmn@(L _ mn)) -> (rn_pkg_qual mn pk, lmn))
   let pi_srcimps = rn_imps pi_srcimps'
   let pi_theimps = rn_imps pi_theimps'
   return PreprocessedImports {..}
diff --git a/compiler/GHC/Driver/Pipeline/Execute.hs b/compiler/GHC/Driver/Pipeline/Execute.hs
index 133f3005b2..95d2c35a0c 100644
--- a/compiler/GHC/Driver/Pipeline/Execute.hs
+++ b/compiler/GHC/Driver/Pipeline/Execute.hs
@@ -83,7 +83,6 @@ import GHC.Unit.Module.Env
 import GHC.Driver.Env.KnotVars
 import GHC.Driver.Config.Finder
 import GHC.Rename.Names
-import Data.Bifunctor (first)
 
 newtype HookedUse a = HookedUse { runHookedUse :: (Hooks, PhaseHook) -> IO a }
   deriving (Functor, Applicative, Monad, MonadIO, MonadThrow, MonadCatch) via (ReaderT (Hooks, PhaseHook) IO)
@@ -642,7 +641,7 @@ runHscPhase pipe_env hsc_env0 input_fn src_flavour = do
     let imp_prelude = xopt LangExt.ImplicitPrelude dflags
         popts = initParserOpts dflags
         rn_pkg_qual = renameRawPkgQual (hsc_unit_env hsc_env)
-        rn_imps = fmap (first rn_pkg_qual)
+        rn_imps = fmap (\(rpk, lmn@(L _ mn)) -> (rn_pkg_qual mn rpk, lmn))
     eimps <- getImports popts imp_prelude buf input_fn (basename <.> suff)
     case eimps of
         Left errs -> throwErrors (GhcPsMessage <$> errs)
diff --git a/compiler/GHC/Rename/Names.hs b/compiler/GHC/Rename/Names.hs
index 806a9e4dbf..34141ab9f4 100644
--- a/compiler/GHC/Rename/Names.hs
+++ b/compiler/GHC/Rename/Names.hs
@@ -328,7 +328,7 @@ rnImportDecl this_mod
         doc = ppr imp_mod_name <+> import_reason
 
     unit_env <- hsc_unit_env <$> getTopEnv
-    let pkg_qual = renameRawPkgQual unit_env raw_pkg_qual
+    let pkg_qual = renameRawPkgQual unit_env imp_mod_name raw_pkg_qual
 
     -- Check for self-import, which confuses the typechecker (#9032)
     -- ghc --make rejects self-import cycles already, but batch-mode may not
@@ -453,21 +453,21 @@ rnImportDecl this_mod
 
 
 -- | Rename raw package imports
-renameRawPkgQual :: UnitEnv -> RawPkgQual -> PkgQual
-renameRawPkgQual unit_env = \case
+renameRawPkgQual :: UnitEnv -> ModuleName -> RawPkgQual -> PkgQual
+renameRawPkgQual unit_env mn = \case
   NoRawPkgQual -> NoPkgQual
-  RawPkgQual p -> renamePkgQual unit_env (Just (sl_fs p))
+  RawPkgQual p -> renamePkgQual unit_env mn (Just (sl_fs p))
 
 -- | Rename raw package imports
-renamePkgQual :: UnitEnv -> Maybe FastString -> PkgQual
-renamePkgQual unit_env mb_pkg = case mb_pkg of
+renamePkgQual :: UnitEnv -> ModuleName -> Maybe FastString -> PkgQual
+renamePkgQual unit_env mn mb_pkg = case mb_pkg of
   Nothing -> NoPkgQual
   Just pkg_fs
     | Just uid <- homeUnitId <$> ue_home_unit unit_env
     , pkg_fs == fsLit "this" || pkg_fs == unitFS uid
     -> ThisPkg uid
 
-    | Just uid <- lookupPackageName (ue_units unit_env) (PackageName pkg_fs)
+    | Just uid <- resolvePackageImport (ue_units unit_env) mn (PackageName pkg_fs)
     -> OtherPkg uid
 
     | otherwise
diff --git a/compiler/GHC/Tc/Module.hs b/compiler/GHC/Tc/Module.hs
index dca730f6f0..6dfcf5d357 100644
--- a/compiler/GHC/Tc/Module.hs
+++ b/compiler/GHC/Tc/Module.hs
@@ -272,7 +272,7 @@ tcRnModuleTcRnM hsc_env mod_sum
 
         ; -- TODO This is a little skeevy; maybe handle a bit more directly
           let { simplifyImport (L _ idecl) =
-                  ( renameRawPkgQual (hsc_unit_env hsc_env) (ideclPkgQual idecl)
+                  ( renameRawPkgQual (hsc_unit_env hsc_env) (unLoc $ ideclName idecl) (ideclPkgQual idecl)
                   , reLoc $ ideclName idecl)
               }
         ; raw_sig_imports <- liftIO
@@ -2056,7 +2056,7 @@ runTcInteractive hsc_env thing_inside
             case i of                   -- force above: see #15111
                 IIModule n -> getOrphans n NoPkgQual
                 IIDecl i   -> getOrphans (unLoc (ideclName i))
-                                         (renameRawPkgQual (hsc_unit_env hsc_env) (ideclPkgQual i))
+                                         (renameRawPkgQual (hsc_unit_env hsc_env) (unLoc $ ideclName i) (ideclPkgQual i))
 
        ; let imports = emptyImportAvails {
                             imp_orphs = orphs
diff --git a/compiler/GHC/Unit/State.hs b/compiler/GHC/Unit/State.hs
index 55855da61f..e178bafea6 100644
--- a/compiler/GHC/Unit/State.hs
+++ b/compiler/GHC/Unit/State.hs
@@ -31,6 +31,7 @@ module GHC.Unit.State (
         unsafeLookupUnitId,
 
         lookupPackageName,
+        resolvePackageImport,
         improveUnit,
         searchPackageId,
         listVisibleModuleNames,
@@ -534,6 +535,8 @@ unsafeLookupUnitId state uid = case lookupUnitId state uid of
 
 -- | Find the unit we know about with the given package name (e.g. @foo@), if any
 -- (NB: there might be a locally defined unit name which overrides this)
+-- This function is unsafe to use in general because it doesn't respect package
+-- visibility.
 lookupPackageName :: UnitState -> PackageName -> Maybe UnitId
 lookupPackageName pkgstate n = lookupUFM (packageNameMap pkgstate) n
 
@@ -542,6 +545,21 @@ searchPackageId :: UnitState -> PackageId -> [UnitInfo]
 searchPackageId pkgstate pid = filter ((pid ==) . unitPackageId)
                                (listUnitInfo pkgstate)
 
+-- | Find the UnitId which an import qualified by a package import comes from.
+-- Compared to 'lookupPackageName', this function correctly accounts for visibility,
+-- renaming and thinning.
+resolvePackageImport :: UnitState -> ModuleName -> PackageName -> Maybe UnitId
+resolvePackageImport unit_st mn pn = do
+  -- 1. Find all modules providing the ModuleName (this accounts for visibility/thinning etc)
+  providers <- Map.lookup mn (moduleNameProvidersMap unit_st)
+  -- 2. Get the UnitIds of the candidates
+  let candidates_uid = map (toUnitId . moduleUnit) $ Map.keys providers
+  -- 3. Get the package names of the candidates
+  let candidates_units = map (\ui -> ((unitPackageName ui), unitId ui))
+                              $ mapMaybe (\uid -> Map.lookup uid (unitInfoMap unit_st)) candidates_uid
+  -- 4. Check to see if the PackageName helps us disambiguate any candidates.
+  lookup pn candidates_units
+
 -- | Create a Map UnitId UnitInfo
 --
 -- For each instantiated unit, we add two map keys: