[Dependency Scanning] Batching Clang Import Scans by Name

include/swift/AST/ModuleDependencies.h

@@ -1255,7 +1255,7 @@ class ModuleDependenciesCache {
                         ModuleDependencyInfo dependencies);
 
   /// Record dependencies for the given module collection.
-  void recordDependencies(ModuleDependencyVector moduleDependencies);
+  void recordDependencies(const ModuleDependencyVector &moduleDependencies);
 
   /// Update stored dependencies for the given module.
   void updateDependency(ModuleDependencyID moduleID,

include/swift/ClangImporter/ClangImporter.h

@@ -495,6 +495,19 @@ class ClangImporter final : public ClangModuleLoader {
                         llvm::PrefixMapper *mapper,
                         bool isTestableImport = false) override;
 
+  // The moduleOutputPath specifies a directory where the precompiled modules
+  // reside in. It should not vary accross the list of modules.
+  llvm::SmallVector<std::pair<ModuleDependencyID, ModuleDependencyInfo>, 1>
+  getModuleDependencies(
+      ArrayRef<StringRef> moduleNames, StringRef moduleOutputPath,
+      const llvm::DenseSet<clang::tooling::dependencies::ModuleID>
+          &alreadySeenClangModules,
+      clang::tooling::dependencies::DependencyScanningTool &clangScanningTool,
+      InterfaceSubContextDelegate &delegate, llvm::PrefixMapper *mapper,
+      std::vector<StringRef> &successModules,
+      std::vector<StringRef> &notFoundModules,
+      std::vector<StringRef> &errorModules, bool isTestableImport = false);
+
   void recordBridgingHeaderOptions(
       ModuleDependencyInfo &MDI,
       const clang::tooling::dependencies::TranslationUnitDeps &deps);

include/swift/DependencyScan/ModuleDependencyScanner.h

@@ -43,6 +43,16 @@ class ModuleDependencyScanningWorker {
                                          const llvm::DenseSet<clang::tooling::dependencies::ModuleID> &alreadySeenModules,
                                          llvm::PrefixMapper *prefixMapper);
 
+  /// Retrieve the module dependencies for a list of Clang modules given a list
+  /// of names.
+  ModuleDependencyVector scanFilesystemForClangModuleDependency(
+      ArrayRef<StringRef> moduleNames, StringRef moduleOutputPath,
+      const llvm::DenseSet<clang::tooling::dependencies::ModuleID>
+          &alreadySeenModules,
+      llvm::PrefixMapper *prefixMapper, std::vector<StringRef> &successModules,
+      std::vector<StringRef> &notFoundModules,
+      std::vector<StringRef> &errorModules);
+
   /// Retrieve the module dependencies for the Swift module with the given name.
   ModuleDependencyVector
   scanFilesystemForSwiftModuleDependency(Identifier moduleName, StringRef moduleOutputPath,
@@ -91,13 +101,6 @@ class ModuleDependencyScanner {
   performDependencyScan(ModuleDependencyID rootModuleID,
                         ModuleDependenciesCache &cache);
 
-  /// Query the module dependency info for the Clang module with the given name.
-  /// Explicit by-name lookups are useful for batch mode scanning.
-  std::optional<const ModuleDependencyInfo *>
-  getNamedClangModuleDependencyInfo(StringRef moduleName,
-                                    ModuleDependenciesCache &cache,
-                                    ModuleDependencyIDSetVector &discoveredClangModules);
-
   /// Query the module dependency info for the Swift module with the given name.
   /// Explicit by-name lookups are useful for batch mode scanning.
   std::optional<const ModuleDependencyInfo *>

lib/AST/ModuleDependencies.cpp

@@ -875,7 +875,7 @@ void ModuleDependenciesCache::recordDependency(
 }
 
 void ModuleDependenciesCache::recordDependencies(
-    ModuleDependencyVector dependencies) {
+    const ModuleDependencyVector &dependencies) {
   for (const auto &dep : dependencies) {
     if (!hasDependency(dep.first))
       recordDependency(dep.first.ModuleName, dep.second);

lib/ClangImporter/ClangModuleDependencyScanner.cpp

@@ -422,16 +422,16 @@ ClangImporter::getModuleDependencies(Identifier moduleName,
 
   auto lookupModuleOutput =
       [moduleOutputPath](const ModuleID &MID,
-                        ModuleOutputKind MOK) -> std::string {
+                         ModuleOutputKind MOK) -> std::string {
     return moduleCacheRelativeLookupModuleOutput(MID, MOK, moduleOutputPath);
   };
 
-  auto clangModuleDependencies =
+  auto [Error, clangModuleDependencies] =
       clangScanningTool.getModuleDependencies(
           moduleName.str(), commandLineArgs, workingDir,
           alreadySeenClangModules, lookupModuleOutput);
-  if (!clangModuleDependencies) {
-    auto errorStr = toString(clangModuleDependencies.takeError());
+  if (Error) {
+    auto errorStr = toString(std::move(Error));
     // We ignore the "module 'foo' not found" error, the Swift dependency
     // scanner will report such an error only if all of the module loaders
     // fail as well.
@@ -443,7 +443,78 @@ ClangImporter::getModuleDependencies(Identifier moduleName,
   }
 
   return bridgeClangModuleDependencies(clangScanningTool,
-                                       *clangModuleDependencies,
+                                       clangModuleDependencies,
+                                       moduleOutputPath, [&](StringRef path) {
+                                         if (mapper)
+                                           return mapper->mapToString(path);
+                                         return path.str();
+                                       });
+}
+
+ModuleDependencyVector ClangImporter::getModuleDependencies(
+    ArrayRef<StringRef> moduleNames, StringRef moduleOutputPath,
+    const llvm::DenseSet<clang::tooling::dependencies::ModuleID>
+        &alreadySeenClangModules,
+    clang::tooling::dependencies::DependencyScanningTool &clangScanningTool,
+    InterfaceSubContextDelegate &delegate, llvm::PrefixMapper *mapper,
+    std::vector<StringRef> &successModules,
+    std::vector<StringRef> &notFoundModules,
+    std::vector<StringRef> &errorModules, bool isTestableImport) {
+  auto &ctx = Impl.SwiftContext;
+  // Determine the command-line arguments for dependency scanning.
+  std::vector<std::string> commandLineArgs =
+      getClangDepScanningInvocationArguments(ctx);
+  auto optionalWorkingDir = computeClangWorkingDirectory(commandLineArgs, ctx);
+  if (!optionalWorkingDir) {
+    ctx.Diags.diagnose(SourceLoc(), diag::clang_dependency_scan_error,
+                       "Missing '-working-directory' argument");
+    return {};
+  }
+  std::string workingDir = *optionalWorkingDir;
+
+  auto lookupModuleOutput =
+      [moduleOutputPath](const ModuleID &MID,
+                         ModuleOutputKind MOK) -> std::string {
+    return moduleCacheRelativeLookupModuleOutput(MID, MOK, moduleOutputPath);
+  };
+
+  auto [error, clangModuleDependencies] =
+      clangScanningTool.getModuleDependencies(
+          moduleNames, commandLineArgs, workingDir, alreadySeenClangModules,
+          lookupModuleOutput);
+  if (error) {
+    auto errorStr = toString(std::move(error));
+    bool showError = false;
+
+    // We partition the modules into three sets, modules successfully scanned,
+    // modules not found, and modules that incur scanning errors.
+    for (const auto &N : moduleNames) {
+      // We ignore the "module 'foo' not found" error, the Swift dependency
+      // scanner will report such an error only if all of the module loaders
+      // fail as well.
+      if (errorStr.find("error: module '" + N.str() + "' not found") ==
+          std::string::npos) {
+        if (errorStr.find("'" + N.str() + "'") != std::string::npos) {
+          showError = true;
+          errorModules.push_back(N);
+        } else
+          successModules.push_back(N);
+      } else
+        notFoundModules.push_back(N);
+    }
+
+    if (showError)
+      ctx.Diags.diagnose(SourceLoc(), diag::clang_dependency_scan_error,
+                         errorStr);
+  } else
+    successModules.insert(successModules.end(), moduleNames.begin(),
+                          moduleNames.end());
+
+  if (!clangModuleDependencies.size())
+    return {};
+
+  return bridgeClangModuleDependencies(clangScanningTool,
+                                       clangModuleDependencies,
                                        moduleOutputPath, [&](StringRef path) {
                                          if (mapper)
                                            return mapper->mapToString(path);

lib/DependencyScan/ModuleDependencyScanner.cpp

@@ -259,6 +259,20 @@ ModuleDependencyScanningWorker::scanFilesystemForClangModuleDependency(
       *scanningASTDelegate, prefixMapper, false);
 }
 
+ModuleDependencyVector
+ModuleDependencyScanningWorker::scanFilesystemForClangModuleDependency(
+    ArrayRef<StringRef> moduleNames, StringRef moduleOutputPath,
+    const llvm::DenseSet<clang::tooling::dependencies::ModuleID>
+        &alreadySeenModules,
+    llvm::PrefixMapper *prefixMapper, std::vector<StringRef> &successModules,
+    std::vector<StringRef> &notFoundModules,
+    std::vector<StringRef> &errorModules) {
+  return clangScannerModuleLoader->getModuleDependencies(
+      moduleNames, moduleOutputPath, alreadySeenModules, clangScanningTool,
+      *scanningASTDelegate, prefixMapper, successModules, notFoundModules,
+      errorModules, false);
+}
+
 template <typename Function, typename... Args>
 auto ModuleDependencyScanner::withDependencyScanningWorker(Function &&F,
                                                            Args &&...ArgList) {
@@ -486,53 +500,6 @@ ModuleDependencyScanner::getMainModuleDependencyInfo(ModuleDecl *mainModule) {
   return mainDependencies;
 }
 
-/// Retrieve the module dependencies for the Clang module with the given name.
-std::optional<const ModuleDependencyInfo *>
-ModuleDependencyScanner::getNamedClangModuleDependencyInfo(
-    StringRef moduleName, ModuleDependenciesCache &cache,
-    ModuleDependencyIDSetVector &discoveredClangModules) {
-  // Check whether we've cached this result.
-  auto moduleID = ModuleDependencyID{moduleName.str(),
-                                     ModuleDependencyKind::Clang};
-  if (auto found = cache.findDependency(moduleID)) {
-    discoveredClangModules.insert(moduleID);
-    auto directClangDeps = cache.getImportedClangDependencies(moduleID);
-    ModuleDependencyIDSetVector reachableClangModules;
-    reachableClangModules.insert(directClangDeps.begin(),
-                                 directClangDeps.end());
-    for (unsigned currentModuleIdx = 0;
-         currentModuleIdx < reachableClangModules.size();
-         ++currentModuleIdx) {
-      auto moduleID = reachableClangModules[currentModuleIdx];
-      auto dependencies =
-        cache.findKnownDependency(moduleID).getImportedClangDependencies();
-      reachableClangModules.insert(dependencies.begin(), dependencies.end());
-    }
-    discoveredClangModules.insert(reachableClangModules.begin(),
-                                  reachableClangModules.end());
-    return found;
-  }
-
-  // Otherwise perform filesystem scan
-  auto moduleIdentifier = getModuleImportIdentifier(moduleName);
-  auto moduleDependencies = withDependencyScanningWorker(
-      [&cache, moduleIdentifier](ModuleDependencyScanningWorker *ScanningWorker) {
-        return ScanningWorker->scanFilesystemForClangModuleDependency(
-          moduleIdentifier, cache.getModuleOutputPath(),
-          cache.getAlreadySeenClangModules(),
-          cache.getScanService().getPrefixMapper());
-      });
-  if (moduleDependencies.empty())
-    return std::nullopt;
-
-  discoveredClangModules.insert(moduleID);
-  for (const auto &dep : moduleDependencies)
-    discoveredClangModules.insert(dep.first);
-
-  cache.recordDependencies(moduleDependencies);
-  return cache.findDependency(moduleID);
-}
-
 /// Retrieve the module dependencies for the Swift module with the given name.
 std::optional<const ModuleDependencyInfo *>
 ModuleDependencyScanner::getNamedSwiftModuleDependencyInfo(
@@ -896,48 +863,102 @@ ModuleDependencyScanner::resolveAllClangModuleDependencies(
   const llvm::DenseSet<clang::tooling::dependencies::ModuleID> seenClangModules =
       cache.getAlreadySeenClangModules();
   std::mutex cacheAccessLock;
-  auto scanForClangModuleDependency =
-      [this, &cache, &moduleLookupResult,
-       &cacheAccessLock, &seenClangModules](Identifier moduleIdentifier) {
-        auto moduleName = moduleIdentifier.str();
+
+  auto scanDependencyForListOfClangModules =
+      [this, &cache, &moduleLookupResult, &cacheAccessLock,
+       &seenClangModules](ArrayRef<StringRef> moduleIDs) {
+        std::vector<StringRef> unresolvedModules;
         {
           std::lock_guard<std::mutex> guard(cacheAccessLock);
-          if (cache.hasDependency(moduleName, ModuleDependencyKind::Clang))
+          // Check to see if there are any modules not resolved yet.
+          for (StringRef ID : moduleIDs) {
+            if (!cache.hasDependency(ID, ModuleDependencyKind::Clang))
+              unresolvedModules.push_back(ID);
+          }
+
+          // No unresolved modules remaining.
+          if (!unresolvedModules.size())
             return;
         }
 
+        std::vector<StringRef> successModules;
+        std::vector<StringRef> notFoundModules;
+        std::vector<StringRef> errorModules;
         auto moduleDependencies = withDependencyScanningWorker(
-            [&cache, &seenClangModules,
-             moduleIdentifier](ModuleDependencyScanningWorker *ScanningWorker) {
-              return ScanningWorker->scanFilesystemForClangModuleDependency(
-                  moduleIdentifier, cache.getModuleOutputPath(),
-                  seenClangModules, cache.getScanService().getPrefixMapper());
+            [&cache, &seenClangModules, &unresolvedModules, &successModules,
+             &notFoundModules,
+             &errorModules](ModuleDependencyScanningWorker *scanningWorker) {
+              return scanningWorker->scanFilesystemForClangModuleDependency(
+                  unresolvedModules, cache.getModuleOutputPath(),
+                  seenClangModules, cache.getScanService().getPrefixMapper(),
+                  successModules, notFoundModules, errorModules);
             });
 
-        // Update the `moduleLookupResult` and cache all discovered dependencies
-        // so that subsequent queries do not have to call into the scanner
-        // if looking for a module that was discovered as a transitive dependency
-        // in this scan.
         {
           std::lock_guard<std::mutex> guard(cacheAccessLock);
-          moduleLookupResult.insert_or_assign(moduleName, moduleDependencies);
-          if (!moduleDependencies.empty())
-            cache.recordDependencies(moduleDependencies);
+          for (const auto &N : successModules) {
+            moduleLookupResult.insert_or_assign(N, moduleDependencies);
+            if (!moduleDependencies.empty())
+              cache.recordDependencies(moduleDependencies);
+          }
+
+          // FIXME: document the reason why we use an empty dependency
+          // vector as a placeholder for modules not found.
+          for (const auto &N : notFoundModules) {
+            moduleLookupResult.insert_or_assign(N, ModuleDependencyVector());
+          }
+
+          for (const auto &N : errorModules) {
+            moduleLookupResult.insert_or_assign(N, ModuleDependencyVector());
+          }
         }
       };
 
-  // Enque asynchronous lookup tasks
-  for (const auto &unresolvedIdentifier : unresolvedImportIdentifiers)
-    ScanningThreadPool.async(
-        scanForClangModuleDependency,
-        getModuleImportIdentifier(unresolvedIdentifier.getKey()));
-  for (const auto &unresolvedIdentifier : unresolvedOptionalImportIdentifiers)
-    ScanningThreadPool.async(
-        scanForClangModuleDependency,
-        getModuleImportIdentifier(unresolvedIdentifier.getKey()));
+  auto batchResolveModuleDependencies = [&](const auto &unresolvedModuleNames) {
+    auto chunkResults = std::div((int)unresolvedModuleNames.size(), NumThreads);
+    const size_t chunkSize = chunkResults.quot;
+    const size_t remainingSize = chunkResults.rem;
+
+    size_t numOfUnresolvedModules = unresolvedModuleNames.size();
+    ArrayRef<StringRef> namesRef(unresolvedModuleNames);
+    size_t current = 0;
+    for (size_t i = 0; i + remainingSize < numOfUnresolvedModules;
+         i += chunkSize) {
+      ArrayRef<StringRef> batch = namesRef.slice(current, chunkSize);
+      ScanningThreadPool.async(scanDependencyForListOfClangModules, batch);
+      current += chunkSize;
+    }
+
+    if (remainingSize) {
+      ArrayRef<StringRef> batch = namesRef.slice(current, remainingSize);
+      ScanningThreadPool.async(scanDependencyForListOfClangModules, batch);
+    }
+  };
+
+  auto getModuleNames = [&](const auto &unresolvedIdentifiers,
+                            auto &unresolvedNames) {
+    llvm::for_each(unresolvedIdentifiers.keys(), [&](auto key) {
+      unresolvedNames.emplace_back(getModuleImportIdentifier(key).str());
+    });
+  };
+
+  std::vector<StringRef> unresolvedImportNames;
+  std::vector<StringRef> unresolvedOptionalImportNames;
+  getModuleNames(unresolvedImportIdentifiers, unresolvedImportNames);
+  getModuleNames(unresolvedOptionalImportIdentifiers,
+                 unresolvedOptionalImportNames);
+
+  batchResolveModuleDependencies(unresolvedImportNames);
+  batchResolveModuleDependencies(unresolvedOptionalImportNames);
   ScanningThreadPool.wait();
 
   // Use the computed scan results to update the dependency info
+  // Four things
+  // 1. For each unresolved import, update the direct dependencies. Stored in
+  // importedClangDependencies.
+  // 2. Every single module we discovered goes into allDiscoveredClangModules.
+  // 3. Every single module we discovered goes into the `cache`. Happened above.
+  // 4. Record every single module that failed to resolve.
   for (const auto &moduleID : swiftModuleDependents) {
     std::vector<ScannerImportStatementInfo> failedToResolveImports;
     ModuleDependencyIDSetVector importedClangDependencies;
@@ -970,6 +991,7 @@ ModuleDependencyScanner::resolveAllClangModuleDependencies(
 
     for (const auto &unresolvedImport : unresolvedImportsMap[moduleID])
       recordResolvedClangModuleImport(unresolvedImport, false);
+
     for (const auto &unresolvedImport : unresolvedOptionalImportsMap[moduleID])
       recordResolvedClangModuleImport(unresolvedImport, true);