Fix multi unit session when some packages have reexported modules.

If we are loading multiple home packages, we need to explicitly take reexports
into account when searching for target files.

If we can't find a module via the usual mean, but it is listed as a reexport of
a unit in scope, we need to look for the module from the perspective of that unit.

This is not necessary for non-home modules because GHC already handles this for
modules in the package DB.

Unfortunately we can't fix this in GHC 9.2 because it doesn't support multiple
home units and we have no way of knowing if a unit reexports modules

Author: wz1000

ghcide/session-loader/Development/IDE/Session.hs

@@ -121,8 +121,7 @@ import           Development.IDE.GHC.Compat.CmdLine
 import qualified Data.Set                             as OS
 
 import           GHC.Data.Bag
-import           GHC.Driver.Env                       (hscSetActiveUnitId,
-                                                       hsc_all_home_unit_ids)
+import           GHC.Driver.Env                       (hsc_all_home_unit_ids)
 import           GHC.Driver.Errors.Types
 import           GHC.Driver.Make                      (checkHomeUnitsClosed)
 import           GHC.Types.Error                      (errMsgDiagnostic)

ghcide/src/Development/IDE/GHC/Compat/Env.hs

@@ -57,6 +57,8 @@ module Development.IDE.GHC.Compat.Env (
     Development.IDE.GHC.Compat.Env.platformDefaultBackend,
     workingDirectory,
     setWorkingDirectory,
+    hscSetActiveUnitId,
+    reexportedModules,
     ) where
 
 import           GHC                 (setInteractiveDynFlags)
@@ -78,10 +80,20 @@ import           GHC.Utils.TmpFs
 
 #if !MIN_VERSION_ghc(9,3,0)
 import           GHC.Driver.Env      (HscEnv, hsc_EPS)
+import qualified Data.Set            as S
 #endif
 
 #if MIN_VERSION_ghc(9,3,0)
-import           GHC.Driver.Env      (HscEnv)
+import           GHC.Driver.Env      (HscEnv, hscSetActiveUnitId)
+#endif
+
+
+#if !MIN_VERSION_ghc(9,3,0)
+hscSetActiveUnitId :: UnitId -> HscEnv -> HscEnv
+hscSetActiveUnitId _ env = env
+
+reexportedModules :: HscEnv -> S.Set a
+reexportedModules _ = S.empty
 #endif
 
 #if MIN_VERSION_ghc(9,3,0)

ghcide/src/Development/IDE/Import/FindImports.hs

@@ -23,7 +23,8 @@ import           Development.IDE.Types.Location
 -- standard imports
 import           Control.Monad.Extra
 import           Control.Monad.IO.Class
-import           Data.List                         (isSuffixOf)
+import           Data.List                         (isSuffixOf, find)
+import qualified Data.Set                          as S
 import           Data.Maybe
 import           System.FilePath
 
@@ -70,19 +71,30 @@ modSummaryToArtifactsLocation nfp ms = ArtifactsLocation nfp (ms_location <$> ms
       Just modSum -> isSource (ms_hsc_src modSum)
     mbMod = ms_mod <$> ms
 
+data LocateResult
+  = LocateNotFound
+  | LocateFoundReexport UnitId
+  | LocateFoundFile UnitId NormalizedFilePath
+
 -- | locate a module in the file system. Where we go from *daml to Haskell
 locateModuleFile :: MonadIO m
-             => [(UnitId, [FilePath])]
+             => [(UnitId, [FilePath], S.Set ModuleName)]
              -> [String]
              -> (ModuleName -> NormalizedFilePath -> m (Maybe NormalizedFilePath))
              -> Bool
              -> ModuleName
-             -> m (Maybe (UnitId, NormalizedFilePath))
+             -> m LocateResult
 locateModuleFile import_dirss exts targetFor isSource modName = do
   let candidates import_dirs =
         [ toNormalizedFilePath' (prefix </> moduleNameSlashes modName <.> maybeBoot ext)
            | prefix <- import_dirs , ext <- exts]
-  firstJustM go (concat [map (uid,) (candidates dirs) | (uid, dirs) <- import_dirss])
+  mf <- firstJustM go (concat [map (uid,) (candidates dirs) | (uid, dirs, _) <- import_dirss])
+  case mf of
+    Nothing ->
+      case find (\(_ , _, reexports) -> S.member modName reexports) import_dirss of
+        Just (uid,_,_) -> pure $ LocateFoundReexport uid
+        Nothing -> pure $ LocateNotFound
+    Just (uid,file) -> pure $ LocateFoundFile uid file
   where
     go (uid, candidate) = fmap ((uid,) <$>) $ targetFor modName candidate
     maybeBoot ext
@@ -94,11 +106,11 @@ locateModuleFile import_dirss exts targetFor isSource modName = do
 -- current module. In particular, it will return Nothing for 'main' components
 -- as they can never be imported into another package.
 #if MIN_VERSION_ghc(9,3,0)
-mkImportDirs :: HscEnv -> (UnitId, DynFlags) -> Maybe (UnitId, [FilePath])
-mkImportDirs _env (i, flags) = Just (i, importPaths flags)
+mkImportDirs :: HscEnv -> (UnitId, DynFlags) -> Maybe (UnitId, ([FilePath], S.Set ModuleName))
+mkImportDirs _env (i, flags) = Just (i, (importPaths flags, reexportedModules flags))
 #else
-mkImportDirs :: HscEnv -> (UnitId, DynFlags) -> Maybe (PackageName, (UnitId, [FilePath]))
-mkImportDirs env (i, flags) = (, (i, importPaths flags)) <$> getUnitName env i
+mkImportDirs :: HscEnv -> (UnitId, DynFlags) -> Maybe (PackageName, (UnitId, [FilePath], S.Set ModuleName))
+mkImportDirs env (i, flags) = (, (i, importPaths flags, S.empty)) <$> getUnitName env i
 #endif
 
 -- | locate a module in either the file system or the package database. Where we go from *daml to
@@ -125,16 +137,16 @@ locateModule env comp_info exts targetFor modName mbPkgName isSource = do
 #else
     Just "this" -> do
 #endif
-      lookupLocal (homeUnitId_ dflags) (importPaths dflags)
+      lookupLocal (homeUnitId_ dflags) (importPaths dflags) S.empty
     -- if a package name is given we only go look for a package
 #if MIN_VERSION_ghc(9,3,0)
     OtherPkg uid
-      | Just dirs <- lookup uid import_paths
-          -> lookupLocal uid dirs
+      | Just (dirs, reexports) <- lookup uid import_paths
+          -> lookupLocal uid dirs reexports
 #else
     Just pkgName
-      | Just (uid, dirs) <- lookup (PackageName pkgName) import_paths
-          -> lookupLocal uid dirs
+      | Just (uid, dirs, reexports) <- lookup (PackageName pkgName) import_paths
+          -> lookupLocal uid dirs reexports
 #endif
       | otherwise -> lookupInPackageDB
 #if MIN_VERSION_ghc(9,3,0)
@@ -143,10 +155,15 @@ locateModule env comp_info exts targetFor modName mbPkgName isSource = do
     Nothing -> do
 #endif
 
-      mbFile <- locateModuleFile ((homeUnitId_ dflags, importPaths dflags) : other_imports) exts targetFor isSource $ unLoc modName
+      -- Reexports for current unit have to be empty because they only apply to other units depending on the
+      -- current unit. If we set the reexports to be the actual reexports then we risk looping forever trying
+      -- to find the module from the perspective of the current unit.
+      mbFile <- locateModuleFile ((homeUnitId_ dflags, importPaths dflags, S.empty) : other_imports) exts targetFor isSource $ unLoc modName
       case mbFile of
-        Nothing          -> lookupInPackageDB
-        Just (uid, file) -> toModLocation uid file
+        LocateNotFound -> lookupInPackageDB
+        -- Lookup again with the perspective of the unit reexporting the file
+        LocateFoundReexport uid -> locateModule (hscSetActiveUnitId uid env) comp_info exts targetFor modName noPkgQual isSource
+        LocateFoundFile uid file -> toModLocation uid file
   where
     dflags = hsc_dflags env
     import_paths = mapMaybe (mkImportDirs env) comp_info
@@ -160,7 +177,7 @@ locateModule env comp_info exts targetFor modName mbPkgName isSource = do
       -- about which module unit a imports.
       -- Without multi-component support it is hard to recontruct the dependency environment so
       -- unit a will have both unit b and unit c in scope.
-      map (\uid -> (uid, importPaths (homeUnitEnv_dflags (ue_findHomeUnitEnv uid ue)))) hpt_deps
+      map (\uid -> let this_df = homeUnitEnv_dflags (ue_findHomeUnitEnv uid ue) in (uid, importPaths this_df, reexportedModules this_df)) hpt_deps
     ue = hsc_unit_env env
     units = homeUnitEnv_units $ ue_findHomeUnitEnv (homeUnitId_ dflags) ue
     hpt_deps :: [UnitId]
@@ -186,11 +203,13 @@ locateModule env comp_info exts targetFor modName mbPkgName isSource = do
         let genMod = mkModule (RealUnit $ Definite uid) (unLoc modName)  -- TODO support backpack holes
         return $ Right $ FileImport $ ArtifactsLocation file (Just loc) (not isSource) (Just genMod)
 
-    lookupLocal uid dirs = do
-      mbFile <- locateModuleFile [(uid, dirs)] exts targetFor isSource $ unLoc modName
+    lookupLocal uid dirs reexports = do
+      mbFile <- locateModuleFile [(uid, dirs, reexports)] exts targetFor isSource $ unLoc modName
       case mbFile of
-        Nothing   -> return $ Left $ notFoundErr env modName $ LookupNotFound []
-        Just (uid', file) -> toModLocation uid' file
+        LocateNotFound -> return $ Left $ notFoundErr env modName $ LookupNotFound []
+        -- Lookup again with the perspective of the unit reexporting the file
+        LocateFoundReexport uid' -> locateModule (hscSetActiveUnitId uid' env) comp_info exts targetFor modName noPkgQual isSource
+        LocateFoundFile uid' file -> toModLocation uid' file
 
     lookupInPackageDB = do
       case Compat.lookupModuleWithSuggestions env (unLoc modName) mbPkgName of
@@ -239,3 +258,11 @@ notFound = NotFound
   , fr_unusables = []
   , fr_suggestions = []
   }
+
+#if MIN_VERSION_ghc(9,3,0)
+noPkgQual :: PkgQual
+noPkgQual = NoPkgQual
+#else
+noPkgQual :: Maybe a
+noPkgQual = Nothing
+#endif

ghcide/test/data/multi-unit-reexport/a-1.0.0-inplace

@@ -0,0 +1,18 @@
+-this-package-name
+a
+-working-dir
+a
+-fbuilding-cabal-package
+-O0
+-i.
+-this-unit-id
+a-1.0.0-inplace
+-hide-all-packages
+-Wmissing-home-modules
+-no-user-package-db
+-package
+base
+-package
+text
+-XHaskell98
+A

ghcide/test/data/multi-unit-reexport/a/A.hs

@@ -0,0 +1,3 @@
+module A(foo) where
+import Data.Text
+foo = ()

ghcide/test/data/multi-unit-reexport/b-1.0.0-inplace

@@ -0,0 +1,21 @@
+-this-package-name
+b
+-working-dir
+b
+-fbuilding-cabal-package
+-O0
+-i
+-i.
+-this-unit-id
+b-1.0.0-inplace
+-hide-all-packages
+-Wmissing-home-modules
+-no-user-package-db
+-package-id
+a-1.0.0-inplace
+-reexported-module
+A
+-package
+base
+-XHaskell98
+B

ghcide/test/data/multi-unit-reexport/b/B.hs

@@ -0,0 +1,3 @@
+module B(module B) where
+import A
+qux = foo

ghcide/test/data/multi-unit-reexport/c-1.0.0-inplace

@@ -0,0 +1,19 @@
+-this-package-name
+c
+-working-dir
+c
+-fbuilding-cabal-package
+-O0
+-i
+-i.
+-this-unit-id
+c-1.0.0-inplace
+-hide-all-packages
+-Wmissing-home-modules
+-no-user-package-db
+-package-id
+b-1.0.0-inplace
+-package
+base
+-XHaskell98
+C

ghcide/test/data/multi-unit-reexport/c/C.hs

@@ -0,0 +1,4 @@
+module C(module C) where
+import A
+import B
+cux = foo `seq` qux

ghcide/test/data/multi-unit-reexport/cabal.project

@@ -0,0 +1,2 @@
+packages: a b c
+multi-repl: True

ghcide/test/data/multi-unit-reexport/hie.yaml

@@ -0,0 +1,6 @@
+cradle:
+  direct:
+    arguments: ["-unit" ,"@a-1.0.0-inplace"
+               ,"-unit" ,"@b-1.0.0-inplace" 
+               ,"-unit" ,"@c-1.0.0-inplace"
+               ]

ghcide/test/exe/CradleTests.hs

@@ -44,6 +44,8 @@ tests = testGroup "cradle"
     ,ignoreFor (BrokenForGHC [GHC92]) "multiple units not supported on 9.2"
        $ testGroup "multi-unit" (multiTests "multi-unit")
     ,testGroup "sub-directory"   [simpleSubDirectoryTest]
+    ,ignoreFor (BrokenForGHC [GHC92]) "multiple units not supported on 9.2"
+      $ testGroup "multi-unit-rexport" [multiRexportTest]
     ]
 
 loadCradleOnlyonce :: TestTree
@@ -187,6 +189,17 @@ simpleMultiDefTest variant = testCase (multiTestName variant "def-test") $ runWi
     checkDefs locs (pure [fooL])
     expectNoMoreDiagnostics 0.5
 
+multiRexportTest :: TestTree
+multiRexportTest =
+  testCase "multi-unit-reexport-test"  $ runWithExtraFiles "multi-unit-reexport" $ \dir -> do
+    let cPath = dir </> "c/C.hs"
+    cdoc <- openDoc cPath "haskell"
+    WaitForIdeRuleResult {} <- waitForAction "TypeCheck" cdoc
+    locs <- getDefinitions cdoc (Position 3 7)
+    let aPath = dir </> "a/A.hs"
+    let fooL = mkL (filePathToUri aPath) 2 0 2 3
+    checkDefs locs (pure [fooL])
+    expectNoMoreDiagnostics 0.5
 
 sessionDepsArePickedUp :: TestTree
 sessionDepsArePickedUp = testSession'