Implement sharing of module graphs

Author: wz1000

Diff for: ghcide/src/Development/IDE/Core/Compile.hs

@@ -131,6 +131,7 @@ import           GHC                               (Anchor (anchor),
 import qualified GHC                               as G
 import           GHC.Hs                            (LEpaComment)
 import qualified GHC.Types.Error                   as Error
+import Development.IDE.Import.DependencyInformation
 #endif
 
 -- | Given a string buffer, return the string (after preprocessing) and the 'ParsedModule'.
@@ -1006,25 +1007,19 @@ loadModulesHome mod_infos e =
 -- Add the current ModSummary to the graph, along with the
 -- HomeModInfo's of all direct dependencies (by induction hypothesis all
 -- transitive dependencies will be contained in envs)
+mergeEnvs :: HscEnv -> ModuleGraph -> ModSummary -> [HomeModInfo] -> [HscEnv] -> IO HscEnv
+mergeEnvs env mg ms extraMods envs = do
 #if MIN_VERSION_ghc(9,3,0)
-mergeEnvs :: HscEnv -> (ModSummary, [NodeKey]) -> [HomeModInfo] -> [HscEnv] -> IO HscEnv
-mergeEnvs env (ms, deps) extraMods envs = do
     let im  = Compat.installedModule (toUnitId $ moduleUnit $ ms_mod ms) (moduleName (ms_mod ms))
         ifr = InstalledFound (ms_location ms) im
         curFinderCache = Compat.extendInstalledModuleEnv Compat.emptyInstalledModuleEnv im ifr
-        -- Very important to force this as otherwise the hsc_mod_graph field is not
-        -- forced and ends up retaining a reference to all the old hsc_envs we have merged to get
-        -- this new one, which in turn leads to the EPS referencing the HPT.
-        module_graph_nodes =
-          nubOrdOn mkNodeKey (ModuleNode deps ms : concatMap (mgModSummaries' . hsc_mod_graph) envs)
-
     newFinderCache <- concatFC curFinderCache (map hsc_FC envs)
-    liftRnf rwhnf module_graph_nodes `seq` (return $ loadModulesHome extraMods $
+    return $! loadModulesHome extraMods $
       let newHug = foldl' mergeHUG (hsc_HUG env) (map hsc_HUG envs) in
       (hscUpdateHUG (const newHug) env){
           hsc_FC = newFinderCache,
-          hsc_mod_graph = mkModuleGraph module_graph_nodes
-      })
+          hsc_mod_graph = mg
+      }
 
     where
         mergeHUG (UnitEnvGraph a) (UnitEnvGraph b) = UnitEnvGraph $ Map.unionWith mergeHUE a b
@@ -1043,30 +1038,16 @@ mergeEnvs env (ms, deps) extraMods envs = do
           pure $ FinderCache fcModules' fcFiles'
 
 #else
-mergeEnvs :: HscEnv -> ModSummary -> [HomeModInfo] -> [HscEnv] -> IO HscEnv
-mergeEnvs env ms extraMods envs = do
     prevFinderCache <- concatFC <$> mapM (readIORef . hsc_FC) envs
     let im  = Compat.installedModule (toUnitId $ moduleUnit $ ms_mod ms) (moduleName (ms_mod ms))
         ifr = InstalledFound (ms_location ms) im
-        -- Very important to force this as otherwise the hsc_mod_graph field is not
-        -- forced and ends up retaining a reference to all the old hsc_envs we have merged to get
-        -- this new one, which in turn leads to the EPS referencing the HPT.
-        module_graph_nodes =
-#if MIN_VERSION_ghc(9,2,0)
-        -- We don't do any instantiation for backpack at this point of time, so it is OK to use
-        -- 'extendModSummaryNoDeps'.
-        -- This may have to change in the future.
-          map extendModSummaryNoDeps $
-#endif
-          nubOrdOn ms_mod (ms : concatMap (mgModSummaries . hsc_mod_graph) envs)
-
     newFinderCache <- newIORef $! Compat.extendInstalledModuleEnv prevFinderCache im ifr
-    liftRnf rwhnf module_graph_nodes `seq` (return $ loadModulesHome extraMods $
+    return $! loadModulesHome extraMods $
       env{
           hsc_HPT = foldMapBy mergeUDFM emptyUDFM hsc_HPT envs,
           hsc_FC = newFinderCache,
-          hsc_mod_graph = mkModuleGraph module_graph_nodes
-      })
+          hsc_mod_graph = mg
+      }
 
     where
         mergeUDFM = plusUDFM_C combineModules
@@ -1460,8 +1441,8 @@ loadInterface session ms linkableNeeded RecompilationInfo{..} = do
            Just (old_hir, _)
              | isNothing linkableNeeded || isJust (hirCoreFp old_hir)
              -> do
-             -- Perform the fine grained recompilation check for TH
-             maybe_recomp <- checkLinkableDependencies get_linkable_hashes (hsc_mod_graph sessionWithMsDynFlags) (hirRuntimeModules old_hir)
+             -- Peform the fine grained recompilation check for TH
+             maybe_recomp <- checkLinkableDependencies session get_linkable_hashes (hirRuntimeModules old_hir)
              case maybe_recomp of
                Just msg -> do_regenerate msg
                Nothing  -> return ([], Just old_hir)
@@ -1472,8 +1453,8 @@ loadInterface session ms linkableNeeded RecompilationInfo{..} = do
              let runtime_deps
                    | not (mi_used_th iface) = emptyModuleEnv
                    | otherwise = parseRuntimeDeps (md_anns details)
-             -- Perform the fine grained recompilation check for TH
-             maybe_recomp <- checkLinkableDependencies get_linkable_hashes (hsc_mod_graph sessionWithMsDynFlags) runtime_deps
+             -- Peform the fine grained recompilation check for TH
+             maybe_recomp <- checkLinkableDependencies session get_linkable_hashes runtime_deps
              case maybe_recomp of
                Just msg -> do_regenerate msg
                Nothing
@@ -1510,13 +1491,21 @@ parseRuntimeDeps anns = mkModuleEnv $ mapMaybe go anns
 -- the runtime dependencies of the module, to check if any of them are out of date
 -- Hopefully 'runtime_deps' will be empty if the module didn't actually use TH
 -- See Note [Recompilation avoidance in the presence of TH]
-checkLinkableDependencies :: MonadIO m => ([NormalizedFilePath] -> m [BS.ByteString]) -> ModuleGraph -> ModuleEnv BS.ByteString -> m (Maybe RecompileRequired)
-checkLinkableDependencies get_linkable_hashes graph runtime_deps = do
-  let hs_files = mapM go (moduleEnvToList runtime_deps)
-      go (mod, hash) = do
-        ms <- mgLookupModule graph mod
-        let hs = fromJust $ ml_hs_file $ ms_location ms
-        pure (toNormalizedFilePath' hs, hash)
+checkLinkableDependencies :: MonadIO m => HscEnv -> ([NormalizedFilePath] -> m [BS.ByteString]) -> ModuleEnv BS.ByteString -> m (Maybe RecompileRequired)
+checkLinkableDependencies hsc_env get_linkable_hashes runtime_deps = do
+#if MIN_VERSION_ghc(9,3,0)
+  moduleLocs <- liftIO $ readIORef (fcModuleCache $ hsc_FC hsc_env)
+#else
+  moduleLocs <- liftIO $ readIORef (hsc_FC hsc_env)
+#endif
+  let go (mod, hash) = do
+        ifr <- lookupInstalledModuleEnv moduleLocs $ Compat.installedModule (toUnitId $ moduleUnit mod) (moduleName mod)
+        case ifr of
+          InstalledFound loc _ -> do
+            hs <- ml_hs_file loc
+            pure (toNormalizedFilePath' hs,hash)
+          _ -> Nothing
+      hs_files = mapM go (moduleEnvToList runtime_deps)
   case hs_files of
     Nothing -> error "invalid module graph"
     Just fs -> do

Diff for: ghcide/src/Development/IDE/Core/Rules.hs

@@ -14,6 +14,10 @@ module Development.IDE.Core.Rules(
     IdeState, GetParsedModule(..), TransitiveDependencies(..),
     Priority(..), GhcSessionIO(..), GetClientSettings(..),
     -- * Functions
+    --
+    --
+    --
+    -- 
     priorityTypeCheck,
     priorityGenerateCore,
     priorityFilesOfInterest,
@@ -23,7 +27,6 @@ module Development.IDE.Core.Rules(
     defineEarlyCutOffNoFile,
     mainRule,
     RulesConfig(..),
-    getDependencies,
     getParsedModule,
     getParsedModuleWithComments,
     getClientConfigAction,
@@ -155,6 +158,7 @@ import qualified Development.IDE.Types.Shake as Shake
 import           Development.IDE.GHC.CoreFile
 import           Data.Time.Clock.POSIX             (posixSecondsToUTCTime)
 import Control.Monad.IO.Unlift
+import qualified Data.IntMap as IM
 #if MIN_VERSION_ghc(9,3,0)
 import GHC.Unit.Module.Graph
 import GHC.Unit.Env
@@ -204,12 +208,6 @@ toIdeResult = either (, Nothing) (([],) . Just)
 ------------------------------------------------------------
 -- Exposed API
 ------------------------------------------------------------
--- | Get all transitive file dependencies of a given module.
--- Does not include the file itself.
-getDependencies :: NormalizedFilePath -> Action (Maybe [NormalizedFilePath])
-getDependencies file =
-    fmap transitiveModuleDeps . (`transitiveDeps` file) <$> useNoFile_ GetModuleGraph
-
 getSourceFileSource :: NormalizedFilePath -> Action BS.ByteString
 getSourceFileSource nfp = do
     (_, msource) <- getFileContents nfp
@@ -417,17 +415,17 @@ type RawDepM a = StateT (RawDependencyInformation, IntMap ArtifactsLocation) Act
 execRawDepM :: Monad m => StateT (RawDependencyInformation, IntMap a1) m a2 -> m (RawDependencyInformation, IntMap a1)
 execRawDepM act =
     execStateT act
-        ( RawDependencyInformation IntMap.empty emptyPathIdMap IntMap.empty IntMap.empty
+        ( RawDependencyInformation IntMap.empty emptyPathIdMap IntMap.empty
         , IntMap.empty
         )
 
 -- | Given a target file path, construct the raw dependency results by following
 -- imports recursively.
-rawDependencyInformation :: [NormalizedFilePath] -> Action RawDependencyInformation
+rawDependencyInformation :: [NormalizedFilePath] -> Action (RawDependencyInformation, BootIdMap)
 rawDependencyInformation fs = do
     (rdi, ss) <- execRawDepM (goPlural fs)
     let bm = IntMap.foldrWithKey (updateBootMap rdi) IntMap.empty ss
-    return (rdi { rawBootMap = bm })
+    return (rdi, bm)
   where
     goPlural ff = do
         mss <- lift $ (fmap.fmap) msrModSummary <$> uses GetModSummaryWithoutTimestamps ff
@@ -446,9 +444,9 @@ rawDependencyInformation fs = do
           fId <- getFreshFid al
           -- Record this module and its location
           whenJust msum $ \ms ->
-            modifyRawDepInfo (\rd -> rd { rawModuleNameMap = IntMap.insert (getFilePathId fId)
-                                                                           (ShowableModuleName (moduleName $ ms_mod ms))
-                                                                           (rawModuleNameMap rd)})
+            modifyRawDepInfo (\rd -> rd { rawModuleMap = IntMap.insert (getFilePathId fId)
+                                                                           (ShowableModule $ ms_mod ms)
+                                                                           (rawModuleMap rd)})
           -- Adding an edge to the bootmap so we can make sure to
           -- insert boot nodes before the real files.
           addBootMap al fId
@@ -670,8 +668,30 @@ knownFilesRule recorder = defineEarlyCutOffNoFile (cmapWithPrio LogShake recorde
 getModuleGraphRule :: Recorder (WithPriority Log) -> Rules ()
 getModuleGraphRule recorder = defineNoFile (cmapWithPrio LogShake recorder) $ \GetModuleGraph -> do
   fs <- toKnownFiles <$> useNoFile_ GetKnownTargets
-  rawDepInfo <- rawDependencyInformation (HashSet.toList fs)
-  pure $ processDependencyInformation rawDepInfo
+  (rawDepInfo, bm) <- rawDependencyInformation (HashSet.toList fs)
+  let (all_fs, _all_ids) = unzip $ HM.toList $ pathToIdMap $ rawPathIdMap rawDepInfo
+  mss <- map (fmap msrModSummary) <$> uses GetModSummaryWithoutTimestamps all_fs
+#if MIN_VERSION_ghc(9,3,0)
+  let deps = map (\i -> IM.lookup (getFilePathId i) (rawImports rawDepInfo)) _all_ids
+      nodeKeys = IM.fromList $ catMaybes $ zipWith (\fi mms -> (getFilePathId fi,) . NodeKey_Module . msKey <$> mms) _all_ids mss
+      mns = catMaybes $ zipWith go mss deps
+      go (Just ms) (Just (Right (ModuleImports xs))) = Just $ ModuleNode this_dep_keys ms
+        where this_dep_ids = mapMaybe snd xs
+              this_dep_keys = mapMaybe (\fi -> IM.lookup (getFilePathId fi) nodeKeys) this_dep_ids
+      go (Just ms) _ = Just $ ModuleNode [] ms
+      go _ _ = Nothing
+      mg = mkModuleGraph mns
+#else
+  let mg = mkModuleGraph $
+#if MIN_VERSION_ghc(9,2,0)
+        -- We don't do any instantiation for backpack at this point of time, so it is OK to use
+        -- 'extendModSummaryNoDeps'.
+        -- This may have to change in the future.
+          map extendModSummaryNoDeps $
+#endif
+          (catMaybes mss)
+#endif
+  pure $ processDependencyInformation rawDepInfo bm mg
 
 -- This is factored out so it can be directly called from the GetModIface
 -- rule. Directly calling this rule means that on the initial load we can
@@ -772,19 +792,8 @@ ghcSessionDepsDefinition fullModSummary GhcSessionDepsConfig{..} env file = do
             depSessions <- map hscEnv <$> uses_ (GhcSessionDeps_ fullModSummary) deps
             ifaces <- uses_ GetModIface deps
             let inLoadOrder = map (\HiFileResult{..} -> HomeModInfo hirModIface hirModDetails Nothing) ifaces
-#if MIN_VERSION_ghc(9,3,0)
-            -- On GHC 9.4+, the module graph contains not only ModSummary's but each `ModuleNode` in the graph
-            -- also points to all the direct descendants of the current module. To get the keys for the descendants
-            -- we must get their `ModSummary`s
-            !final_deps <- do
-              dep_mss <- map msrModSummary <$> uses_ GetModSummaryWithoutTimestamps deps
-             -- Don't want to retain references to the entire ModSummary when just the key will do
-              return $!! map (NodeKey_Module . msKey) dep_mss
-            let moduleNode = (ms, final_deps)
-#else
-            let moduleNode = ms
-#endif
-            session' <- liftIO $ mergeEnvs hsc moduleNode inLoadOrder depSessions
+            mg <- depModuleGraph <$> useNoFile_ GetModuleGraph
+            session' <- liftIO $ mergeEnvs hsc mg ms inLoadOrder depSessions
 
             Just <$> liftIO (newHscEnvEqWithImportPaths (envImportPaths env) session' [])
 

Diff for: ghcide/src/Development/IDE/GHC/Orphans.hs

@@ -40,6 +40,7 @@ import           Data.String                (IsString (fromString))
 import           Data.Text                  (unpack)
 #if MIN_VERSION_ghc(9,0,0)
 import           GHC.ByteCode.Types
+import GHC (ModuleGraph)
 #else
 import           ByteCodeTypes
 #endif
@@ -207,6 +208,9 @@ instance Hashable OccName where hashWithSalt s n = hashWithSalt s (getKey $ getU
 
 instance Show HomeModInfo where show = show . mi_module . hm_iface
 
+instance Show ModuleGraph where show _ = "ModuleGraph {..}"
+instance NFData ModuleGraph where rnf = rwhnf
+
 instance NFData HomeModInfo where
   rnf (HomeModInfo iface dets link) = rwhnf iface `seq` rnf dets `seq` rnf link
 

Diff for: ghcide/src/Development/IDE/Import/DependencyInformation.hs

@@ -10,8 +10,9 @@ module Development.IDE.Import.DependencyInformation
   , TransitiveDependencies(..)
   , FilePathId(..)
   , NamedModuleDep(..)
-  , ShowableModuleName(..)
-  , PathIdMap
+  , ShowableModule(..)
+  , ShowableModuleEnv(..)
+  , PathIdMap (..)
   , emptyPathIdMap
   , getPathId
   , lookupPathToId
@@ -23,7 +24,7 @@ module Development.IDE.Import.DependencyInformation
   , transitiveDeps
   , transitiveReverseDependencies
   , immediateReverseDependencies
-
+  , lookupModuleFile
   , BootIdMap
   , insertBootId
   ) where
@@ -53,6 +54,7 @@ import           Development.IDE.Types.Diagnostics
 import           Development.IDE.Types.Location
 
 import           GHC
+import Development.IDE.GHC.Compat
 
 -- | The imports for a given module.
 newtype ModuleImports = ModuleImports
@@ -128,15 +130,14 @@ data RawDependencyInformation = RawDependencyInformation
     -- corresponding hs file. It is used when topologically sorting as we
     -- need to add edges between .hs-boot and .hs so that the .hs files
     -- appear later in the sort.
-    , rawBootMap   :: !BootIdMap
-    , rawModuleNameMap :: !(FilePathIdMap ShowableModuleName)
+    , rawModuleMap :: !(FilePathIdMap ShowableModule)
     } deriving Show
 
 data DependencyInformation =
   DependencyInformation
     { depErrorNodes        :: !(FilePathIdMap (NonEmpty NodeError))
     -- ^ Nodes that cannot be processed correctly.
-    , depModuleNames       :: !(FilePathIdMap ShowableModuleName)
+    , depModules       :: !(FilePathIdMap ShowableModule)
     , depModuleDeps        :: !(FilePathIdMap FilePathIdSet)
     -- ^ For a non-error node, this contains the set of module immediate dependencies
     -- in the same package.
@@ -146,13 +147,24 @@ data DependencyInformation =
     -- ^ Map from FilePath to FilePathId
     , depBootMap           :: !BootIdMap
     -- ^ Map from hs-boot file to the corresponding hs file
+    , depModuleFiles       :: !(ShowableModuleEnv FilePathId)
+    -- ^ Map from Module to the corresponding non-boot hs file
+    , depModuleGraph       :: !ModuleGraph
     } deriving (Show, Generic)
 
-newtype ShowableModuleName =
-  ShowableModuleName {showableModuleName :: ModuleName}
+newtype ShowableModule =
+  ShowableModule {showableModule :: Module}
   deriving NFData
 
-instance Show ShowableModuleName where show = moduleNameString . showableModuleName
+newtype ShowableModuleEnv a =
+  ShowableModuleEnv {showableModuleEnv :: ModuleEnv a}
+
+instance Show a => Show (ShowableModuleEnv a) where
+  show (ShowableModuleEnv x) = show (moduleEnvToList x)
+instance NFData a => NFData (ShowableModuleEnv a) where
+  rnf = rwhnf
+
+instance Show ShowableModule where show = moduleNameString . moduleName . showableModule
 
 reachableModules :: DependencyInformation -> [NormalizedFilePath]
 reachableModules DependencyInformation{..} =
@@ -215,15 +227,17 @@ instance Semigroup NodeResult where
    SuccessNode _ <> ErrorNode errs   = ErrorNode errs
    SuccessNode a <> SuccessNode _    = SuccessNode a
 
-processDependencyInformation :: RawDependencyInformation -> DependencyInformation
-processDependencyInformation RawDependencyInformation{..} =
+processDependencyInformation :: RawDependencyInformation -> BootIdMap -> ModuleGraph -> DependencyInformation
+processDependencyInformation RawDependencyInformation{..} rawBootMap mg =
   DependencyInformation
     { depErrorNodes = IntMap.fromList errorNodes
     , depModuleDeps = moduleDeps
     , depReverseModuleDeps = reverseModuleDeps
-    , depModuleNames = rawModuleNameMap
+    , depModules = rawModuleMap
     , depPathIdMap = rawPathIdMap
     , depBootMap = rawBootMap
+    , depModuleFiles = ShowableModuleEnv reverseModuleMap
+    , depModuleGraph = mg
     }
   where resultGraph = buildResultGraph rawImports
         (errorNodes, successNodes) = partitionNodeResults $ IntMap.toList resultGraph
@@ -240,6 +254,7 @@ processDependencyInformation RawDependencyInformation{..} =
           foldr (\(p, cs) res ->
             let new = IntMap.fromList (map (, IntSet.singleton (coerce p)) (coerce cs))
             in IntMap.unionWith IntSet.union new res ) IntMap.empty successEdges
+        reverseModuleMap = mkModuleEnv $ map (\(i,sm) -> (showableModule sm, FilePathId i)) $ IntMap.toList rawModuleMap
 
 
 -- | Given a dependency graph, buildResultGraph detects and propagates errors in that graph as follows:
@@ -351,6 +366,10 @@ transitiveDeps DependencyInformation{..} file = do
 
     vs = topSort g
 
+lookupModuleFile :: Module -> DependencyInformation -> Maybe NormalizedFilePath
+lookupModuleFile mod DependencyInformation{..}
+  = idToPath depPathIdMap <$> lookupModuleEnv (showableModuleEnv depModuleFiles) mod
+
 newtype TransitiveDependencies = TransitiveDependencies
   { transitiveModuleDeps :: [NormalizedFilePath]
   -- ^ Transitive module dependencies in topological order.