From 55d2e5f754bec2e095706c165dacdcc63c63e608 Mon Sep 17 00:00:00 2001
From: Zubin Duggal <zubin.duggal@gmail.com>
Date: Fri, 5 Jan 2024 15:55:03 +0530
Subject: [PATCH] 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
---
 ghcide/src/Development/IDE/GHC/Compat/Env.hs  | 14 +++-
 .../src/Development/IDE/Import/FindImports.hs | 69 +++++++++++++------
 .../data/multi-unit-reexport/a-1.0.0-inplace  | 18 +++++
 ghcide/test/data/multi-unit-reexport/a/A.hs   |  3 +
 .../data/multi-unit-reexport/b-1.0.0-inplace  | 21 ++++++
 ghcide/test/data/multi-unit-reexport/b/B.hs   |  3 +
 .../data/multi-unit-reexport/c-1.0.0-inplace  | 19 +++++
 ghcide/test/data/multi-unit-reexport/c/C.hs   |  4 ++
 .../data/multi-unit-reexport/cabal.project    |  2 +
 ghcide/test/data/multi-unit-reexport/hie.yaml |  6 ++
 ghcide/test/exe/CradleTests.hs                | 13 ++++
 11 files changed, 150 insertions(+), 22 deletions(-)
 create mode 100644 ghcide/test/data/multi-unit-reexport/a-1.0.0-inplace
 create mode 100644 ghcide/test/data/multi-unit-reexport/a/A.hs
 create mode 100644 ghcide/test/data/multi-unit-reexport/b-1.0.0-inplace
 create mode 100644 ghcide/test/data/multi-unit-reexport/b/B.hs
 create mode 100644 ghcide/test/data/multi-unit-reexport/c-1.0.0-inplace
 create mode 100644 ghcide/test/data/multi-unit-reexport/c/C.hs
 create mode 100644 ghcide/test/data/multi-unit-reexport/cabal.project
 create mode 100644 ghcide/test/data/multi-unit-reexport/hie.yaml

diff --git a/ghcide/src/Development/IDE/GHC/Compat/Env.hs b/ghcide/src/Development/IDE/GHC/Compat/Env.hs
index 5466d5fc22e..7b4125bea90 100644
--- a/ghcide/src/Development/IDE/GHC/Compat/Env.hs
+++ b/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)
diff --git a/ghcide/src/Development/IDE/Import/FindImports.hs b/ghcide/src/Development/IDE/Import/FindImports.hs
index 358666a0e9c..e7db3e0358d 100644
--- a/ghcide/src/Development/IDE/Import/FindImports.hs
+++ b/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 (\(uid, _, 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
@@ -235,3 +254,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
diff --git a/ghcide/test/data/multi-unit-reexport/a-1.0.0-inplace b/ghcide/test/data/multi-unit-reexport/a-1.0.0-inplace
new file mode 100644
index 00000000000..a54ea9bc4bb
--- /dev/null
+++ b/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
diff --git a/ghcide/test/data/multi-unit-reexport/a/A.hs b/ghcide/test/data/multi-unit-reexport/a/A.hs
new file mode 100644
index 00000000000..9a7d7e33c91
--- /dev/null
+++ b/ghcide/test/data/multi-unit-reexport/a/A.hs
@@ -0,0 +1,3 @@
+module A(foo) where
+import Data.Text
+foo = ()
diff --git a/ghcide/test/data/multi-unit-reexport/b-1.0.0-inplace b/ghcide/test/data/multi-unit-reexport/b-1.0.0-inplace
new file mode 100644
index 00000000000..d656a2539b6
--- /dev/null
+++ b/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
diff --git a/ghcide/test/data/multi-unit-reexport/b/B.hs b/ghcide/test/data/multi-unit-reexport/b/B.hs
new file mode 100644
index 00000000000..2c6d4b28a22
--- /dev/null
+++ b/ghcide/test/data/multi-unit-reexport/b/B.hs
@@ -0,0 +1,3 @@
+module B(module B) where
+import A
+qux = foo
diff --git a/ghcide/test/data/multi-unit-reexport/c-1.0.0-inplace b/ghcide/test/data/multi-unit-reexport/c-1.0.0-inplace
new file mode 100644
index 00000000000..e60a95eda0c
--- /dev/null
+++ b/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
diff --git a/ghcide/test/data/multi-unit-reexport/c/C.hs b/ghcide/test/data/multi-unit-reexport/c/C.hs
new file mode 100644
index 00000000000..1b2d305296f
--- /dev/null
+++ b/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
diff --git a/ghcide/test/data/multi-unit-reexport/cabal.project b/ghcide/test/data/multi-unit-reexport/cabal.project
new file mode 100644
index 00000000000..96f52330c92
--- /dev/null
+++ b/ghcide/test/data/multi-unit-reexport/cabal.project
@@ -0,0 +1,2 @@
+packages: a b c
+multi-repl: True
diff --git a/ghcide/test/data/multi-unit-reexport/hie.yaml b/ghcide/test/data/multi-unit-reexport/hie.yaml
new file mode 100644
index 00000000000..34858b5f641
--- /dev/null
+++ b/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"
+               ]
diff --git a/ghcide/test/exe/CradleTests.hs b/ghcide/test/exe/CradleTests.hs
index 9274e807c99..94d271b85bf 100644
--- a/ghcide/test/exe/CradleTests.hs
+++ b/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'