From 873b9cba0285d03d7cd5dcf0006fe3262bbf33ad Mon Sep 17 00:00:00 2001
From: Yuchen Pei <hi@ypei.me>
Date: Tue, 7 Jun 2022 19:18:51 +1000
Subject: Fixing problem with indexing a project with multiple packages

- it will only index the first package with the same source dir as the
  -p option, or if no such package exists, index the first package
- it will concat map all the modules of all the units of the package
  to index
---
 src/HaskellCodeExplorer/PackageInfo.hs | 48 ++++++++++++++++++++++++++++------
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/src/HaskellCodeExplorer/PackageInfo.hs b/src/HaskellCodeExplorer/PackageInfo.hs
index 279a9a1..4fe22ea 100644
--- a/src/HaskellCodeExplorer/PackageInfo.hs
+++ b/src/HaskellCodeExplorer/PackageInfo.hs
@@ -32,6 +32,7 @@ import Control.Monad.Logger
   , MonadLoggerIO(..)
   , logDebugN
   , logErrorN
+  , logWarnN
   , logInfoN
   , runStdoutLoggingT
   )
@@ -40,7 +41,12 @@ import qualified Data.HashMap.Strict as HM
 import Data.IORef (readIORef)
 import qualified Data.IntMap.Strict as IM
 import qualified Data.List as L
-import Data.Maybe (fromMaybe, isJust, maybeToList)
+import Data.Maybe
+  ( fromMaybe
+  , isJust
+  , maybeToList
+  , mapMaybe
+  )
 import qualified Data.Set as S
 import qualified Data.Text as T
 import qualified Data.Text.Encoding as TE
@@ -54,6 +60,11 @@ import Distribution.Helper
   , DistDir(..)
   , SCabalProjType(..)
   , allUnits
+  , projectPackages
+  , pPackageName
+  , pSourceDir
+  , pUnits
+  , uComponentName
   , UnitInfo(..)
   , ChComponentInfo(..)
   , mkQueryEnv
@@ -178,8 +189,30 @@ createPackageInfo packageDirectoryPath mbDistDirRelativePath sourceCodePreproces
                          mkQueryEnv
                          (ProjLocV2Dir packageDirectoryAbsPath)
                          (DistDirCabal SCV2 distDir)
-  ((packageName, packageVersion), (_packageCompilerName, packageCompilerVersion), compInfo) <-
-    liftIO $ NE.head <$>
+  packages <- liftIO $ NE.toList <$> runQuery projectPackages cabalHelperQueryEnv
+  logDebugN $ "packages: " <>
+    (T.pack $ show $ zip3 (pPackageName <$> packages) (pSourceDir <$> packages) ((mapMaybe uComponentName . NE.toList . pUnits) <$> packages))
+  mbPackage <- liftIO $
+    findM
+    (\pkg -> do
+        dir <- (fmap (</> "") . fmap normalise . makeAbsolute . pSourceDir) pkg
+        return $ dir == packageDirectoryAbsPath </> "")
+    packages
+  package <-
+    case mbPackage of
+      Just package' -> return package'
+      Nothing -> do
+        logWarnN $
+          "Cannot find a package with sourceDir in the same directory ("
+          <> T.pack (packageDirectoryAbsPath </> "")
+          <> "), indexing the first package by default."
+          <> "Alternatively, try using absolute path for -p."
+        return $ head packages
+  
+  -- ((packageName, packageVersion), (_packageCompilerName, packageCompilerVersion), compInfo) <-
+  units <-
+    liftIO $
+    (filter (\((pkgName, _), _, _) -> pkgName == pPackageName package)) . NE.toList <$>
     runQuery
       (allUnits
        (\unit ->
@@ -189,13 +222,12 @@ createPackageInfo packageDirectoryPath mbDistDirRelativePath sourceCodePreproces
                           (ciSourceDirs comp, ciComponentName comp))) $
             (Map.elems . uiComponents) unit)))
       cabalHelperQueryEnv
-      -- ((,,) <$> packageId <*> compilerVersion <*>
-      --  (zip3 <$> components ((,) <$> ghcOptions) <*>
-      --   components ((,) <$> entrypoints) <*>
-      --   components ((,) <$> sourceDirs)))
+  -- TODO: we are assuming all pakcageVersion and packageCompilerVersion are the same
+  let ((packageName, packageVersion), (_, packageCompilerVersion), _) = head units
+      compInfo = concatMap (\(_, _, comp) -> comp) units
   -- logInfoN $ "unitinfo: " <> (T.pack $ show (packageName, packageVersion))
   -- logInfoN $ "compinfo: " <> (T.pack $ show compInfo)
-  let currentPackageId = HCE.PackageId (T.pack packageName) packageVersion
+      currentPackageId = HCE.PackageId (T.pack packageName) packageVersion
   unless
     (take 3 (versionBranch packageCompilerVersion) ==
      take 3 (versionBranch ghcVersion)) $ do
-- 
cgit v1.2.3