[clang][deps] Minor ModuleDepCollector refactorings NFC

* Factor module map and module file path functions out
* Use a secondary mapping to lookup module deps by ID instead of the
  preprocessor module map.
* Sink DirectPrebuiltModularDeps into MDC.

Differential Revision: https://reviews.llvm.org/D132617

Author: benlangmuir

clang/include/clang/Tooling/DependencyScanning/ModuleDepCollector.h

@@ -61,12 +61,6 @@ struct ModuleID {
   }
 };
 
-struct ModuleIDHasher {
-  std::size_t operator()(const ModuleID &MID) const {
-    return llvm::hash_combine(MID.ModuleName, MID.ContextHash);
-  }
-};
-
 /// An output from a module compilation, such as the path of the module file.
 enum class ModuleOutputKind {
   /// The module file (.pcm). Required.
@@ -153,8 +147,6 @@ class ModuleDepCollectorPP final : public PPCallbacks {
   ModuleDepCollector &MDC;
   /// Working set of direct modular dependencies.
   llvm::SetVector<const Module *> DirectModularDeps;
-  /// Working set of direct modular dependencies that have already been built.
-  llvm::SetVector<const Module *> DirectPrebuiltModularDeps;
 
   void handleImport(const Module *Imported);
 
@@ -211,6 +203,11 @@ class ModuleDepCollector final : public DependencyCollector {
   std::vector<std::string> FileDeps;
   /// Direct and transitive modular dependencies of the main source file.
   llvm::MapVector<const Module *, std::unique_ptr<ModuleDeps>> ModularDeps;
+  /// Secondary mapping for \c ModularDeps allowing lookup by ModuleID without
+  /// a preprocessor. Storage owned by \c ModularDeps.
+  llvm::DenseMap<ModuleID, ModuleDeps *> ModuleDepsByID;
+  /// Direct modular dependencies that have already been built.
+  llvm::MapVector<const Module *, PrebuiltModuleDep> DirectPrebuiltModularDeps;
   /// Options that control the dependency output generation.
   std::unique_ptr<DependencyOutputOptions> Opts;
   /// The original Clang invocation passed to dependency scanner.
@@ -235,12 +232,39 @@ class ModuleDepCollector final : public DependencyCollector {
       const ModuleDeps &Deps,
       llvm::function_ref<void(CompilerInvocation &)> Optimize) const;
 
+  /// Add module map files to the invocation, if needed.
+  void addModuleMapFiles(CompilerInvocation &CI,
+                         ArrayRef<ModuleID> ClangModuleDeps) const;
+  /// Add module files (pcm) to the invocation, if needed.
+  void addModuleFiles(CompilerInvocation &CI,
+                      ArrayRef<ModuleID> ClangModuleDeps) const;
+
   /// Add paths that require looking up outputs to the given dependencies.
   void addOutputPaths(CompilerInvocation &CI, ModuleDeps &Deps);
+
+  /// Compute the context hash for \p Deps, and create the mapping
+  /// \c ModuleDepsByID[Deps.ID] = &Deps.
+  void associateWithContextHash(const CompilerInvocation &CI, ModuleDeps &Deps);
 };
 
 } // end namespace dependencies
 } // end namespace tooling
 } // end namespace clang
 
+namespace llvm {
+template <> struct DenseMapInfo<clang::tooling::dependencies::ModuleID> {
+  using ModuleID = clang::tooling::dependencies::ModuleID;
+  static inline ModuleID getEmptyKey() { return ModuleID{"", ""}; }
+  static inline ModuleID getTombstoneKey() {
+    return ModuleID{"~", "~"}; // ~ is not a valid module name or context hash
+  }
+  static unsigned getHashValue(const ModuleID &ID) {
+    return hash_combine(ID.ModuleName, ID.ContextHash);
+  }
+  static bool isEqual(const ModuleID &LHS, const ModuleID &RHS) {
+    return LHS == RHS;
+  }
+};
+} // namespace llvm
+
 #endif // LLVM_CLANG_TOOLING_DEPENDENCYSCANNING_MODULEDEPCOLLECTOR_H

clang/lib/Tooling/DependencyScanning/ModuleDepCollector.cpp

@@ -57,15 +57,7 @@ void ModuleDepCollector::addOutputPaths(CompilerInvocation &CI,
   // These are technically *inputs* to the compilation, but we populate them
   // here in order to make \c getModuleContextHash() independent of
   // \c lookupModuleOutput().
-  for (ModuleID MID : Deps.ClangModuleDeps) {
-    auto PCMPath =
-        Consumer.lookupModuleOutput(MID, ModuleOutputKind::ModuleFile);
-    if (EagerLoadModules)
-      CI.getFrontendOpts().ModuleFiles.push_back(PCMPath);
-    else
-      CI.getHeaderSearchOpts().PrebuiltModuleFiles.insert(
-          {MID.ModuleName, PCMPath});
-  }
+  addModuleFiles(CI, Deps.ClangModuleDeps);
 
   CI.getFrontendOpts().OutputFile =
       Consumer.lookupModuleOutput(Deps.ID, ModuleOutputKind::ModuleFile);
@@ -127,24 +119,12 @@ ModuleDepCollector::makeInvocationForModuleBuildWithoutOutputs(
   CI.getFrontendOpts().Inputs.emplace_back(Deps.ClangModuleMapFile,
                                            ModuleMapInputKind);
   CI.getFrontendOpts().ModuleMapFiles = Deps.ModuleMapFileDeps;
+  addModuleMapFiles(CI, Deps.ClangModuleDeps);
 
   // Report the prebuilt modules this module uses.
   for (const auto &PrebuiltModule : Deps.PrebuiltModuleDeps)
     CI.getFrontendOpts().ModuleFiles.push_back(PrebuiltModule.PCMFile);
 
-  if (!EagerLoadModules) {
-    ModuleMap &ModMap =
-        ScanInstance.getPreprocessor().getHeaderSearchInfo().getModuleMap();
-    for (ModuleID MID : Deps.ClangModuleDeps) {
-      const Module *M = ModMap.findModule(MID.ModuleName);
-      assert(M && "Modular dependency not found");
-      auto MDeps = ModularDeps.find(M);
-      assert(MDeps != ModularDeps.end() && "Inconsistent dependency info");
-      CI.getFrontendOpts().ModuleMapFiles.push_back(
-          MDeps->second->ClangModuleMapFile);
-    }
-  }
-
   // Remove any macro definitions that are explicitly ignored.
   if (!CI.getHeaderSearchOpts().ModulesIgnoreMacros.empty()) {
     llvm::erase_if(
@@ -171,6 +151,31 @@ ModuleDepCollector::makeInvocationForModuleBuildWithoutOutputs(
   return CI;
 }
 
+void ModuleDepCollector::addModuleMapFiles(
+    CompilerInvocation &CI, ArrayRef<ModuleID> ClangModuleDeps) const {
+  if (EagerLoadModules)
+    return; // Only pcm is needed for eager load.
+
+  for (const ModuleID &MID : ClangModuleDeps) {
+    ModuleDeps *MD = ModuleDepsByID.lookup(MID);
+    assert(MD && "Inconsistent dependency info");
+    CI.getFrontendOpts().ModuleMapFiles.push_back(MD->ClangModuleMapFile);
+  }
+}
+
+void ModuleDepCollector::addModuleFiles(
+    CompilerInvocation &CI, ArrayRef<ModuleID> ClangModuleDeps) const {
+  for (const ModuleID &MID : ClangModuleDeps) {
+    std::string PCMPath =
+        Consumer.lookupModuleOutput(MID, ModuleOutputKind::ModuleFile);
+    if (EagerLoadModules)
+      CI.getFrontendOpts().ModuleFiles.push_back(std::move(PCMPath));
+    else
+      CI.getHeaderSearchOpts().PrebuiltModuleFiles.insert(
+          {MID.ModuleName, std::move(PCMPath)});
+  }
+}
+
 static std::string getModuleContextHash(const ModuleDeps &MD,
                                         const CompilerInvocation &CI,
                                         bool EagerLoadModules) {
@@ -212,6 +217,14 @@ static std::string getModuleContextHash(const ModuleDeps &MD,
   return toString(llvm::APInt(sizeof(Words) * 8, Words), 36, /*Signed=*/false);
 }
 
+void ModuleDepCollector::associateWithContextHash(const CompilerInvocation &CI,
+                                                  ModuleDeps &Deps) {
+  Deps.ID.ContextHash = getModuleContextHash(Deps, CI, EagerLoadModules);
+  bool Inserted = ModuleDepsByID.insert({Deps.ID, &Deps}).second;
+  (void)Inserted;
+  assert(Inserted && "duplicate module mapping");
+}
+
 void ModuleDepCollectorPP::FileChanged(SourceLocation Loc,
                                        FileChangeReason Reason,
                                        SrcMgr::CharacteristicKind FileType,
@@ -263,7 +276,8 @@ void ModuleDepCollectorPP::handleImport(const Module *Imported) {
   const Module *TopLevelModule = Imported->getTopLevelModule();
 
   if (MDC.isPrebuiltModule(TopLevelModule))
-    DirectPrebuiltModularDeps.insert(TopLevelModule);
+    MDC.DirectPrebuiltModularDeps.insert(
+        {TopLevelModule, PrebuiltModuleDep{TopLevelModule}});
   else
     DirectModularDeps.insert(TopLevelModule);
 }
@@ -300,8 +314,8 @@ void ModuleDepCollectorPP::EndOfMainFile() {
   for (auto &&I : MDC.FileDeps)
     MDC.Consumer.handleFileDependency(I);
 
-  for (auto &&I : DirectPrebuiltModularDeps)
-    MDC.Consumer.handlePrebuiltModuleDependency(PrebuiltModuleDep{I});
+  for (auto &&I : MDC.DirectPrebuiltModularDeps)
+    MDC.Consumer.handlePrebuiltModuleDependency(I.second);
 }
 
 ModuleID ModuleDepCollectorPP::handleTopLevelModule(const Module *M) {
@@ -403,8 +417,8 @@ ModuleID ModuleDepCollectorPP::handleTopLevelModule(const Module *M) {
                                    *MDC.ScanInstance.getASTReader(), *MF);
       });
 
-  // Compute the context hash from the inputs. Requires dependencies.
-  MD.ID.ContextHash = getModuleContextHash(MD, CI, MDC.EagerLoadModules);
+  MDC.associateWithContextHash(CI, MD);
+
   // Finish the compiler invocation. Requires dependencies and the context hash.
   MDC.addOutputPaths(CI, MD);