[Dependency Scanning] Cache compiler instance/module depenency cache used for batched versioned-PCM scans.

Author: artemcm

include/swift/DependencyScan/DependencyScanningTool.h

@@ -14,10 +14,10 @@
 #define SWIFT_DEPENDENCY_SCANNING_TOOL_H
 
 #include "swift-c/DependencyScan/DependencyScan.h"
+#include "swift/Frontend/Frontend.h"
 #include "swift/AST/ModuleDependencies.h"
 #include "swift/DependencyScan/ScanDependencies.h"
 #include "swift/Frontend/PrintingDiagnosticConsumer.h"
-#include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/StringSaver.h"
 
@@ -63,10 +63,14 @@ class DependencyScanningTool {
   llvm::ErrorOr<std::unique_ptr<CompilerInstance>>
   initCompilerInstanceForScan(ArrayRef<const char *> Command);
 
-  /// Shared cache of module dependencies, re-used by individual queries
+  /// Shared cache of module dependencies, re-used by individual full-scan queries
   /// during the lifetime of this Tool.
   std::unique_ptr<ModuleDependenciesCache> SharedCache;
 
+  /// Shared cache of compiler instances created during batch scanning, corresponding to
+  /// command-line options specified in the batch scan input entry.
+  std::unique_ptr<CompilerArgInstanceCacheMap> VersionedPCMInstanceCacheCache;
+
   /// A shared consumer that, for now, just prints the encountered diagnostics.
   PrintingDiagnosticConsumer PDC;
   llvm::BumpPtrAllocator Alloc;

include/swift/DependencyScan/ScanDependencies.h

@@ -15,6 +15,7 @@
 
 #include "swift-c/DependencyScan/DependencyScan.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Support/Error.h"
 
 namespace llvm {
@@ -29,6 +30,10 @@ class ModuleDependenciesCache;
 
 namespace dependencies {
 
+using CompilerArgInstanceCacheMap =
+    llvm::StringMap<std::pair<std::unique_ptr<CompilerInstance>,
+                              std::unique_ptr<ModuleDependenciesCache>>>;
+
 struct BatchScanInput {
   llvm::StringRef moduleName;
   llvm::StringRef arguments;
@@ -64,13 +69,15 @@ performModulePrescan(CompilerInstance &instance);
 
 /// Batch scan the dependencies for modules specified in \c batchInputFile.
 std::vector<llvm::ErrorOr<swiftscan_dependency_graph_t>>
-performBatchModuleScan(CompilerInstance &instance,
-                       ModuleDependenciesCache &cache, llvm::StringSaver &saver,
+performBatchModuleScan(CompilerInstance &invocationInstance,
+                       ModuleDependenciesCache &invocationCache,
+                       CompilerArgInstanceCacheMap *versionedPCMInstanceCache,
+                       llvm::StringSaver &saver,
                        const std::vector<BatchScanInput> &BatchInput);
 
 /// Batch prescan the imports of modules specified in \c batchInputFile.
 std::vector<llvm::ErrorOr<swiftscan_import_set_t>>
-performBatchModulePrescan(CompilerInstance &instance,
+performBatchModulePrescan(CompilerInstance &invocationInstance,
                           ModuleDependenciesCache &cache,
                           llvm::StringSaver &saver,
                           const std::vector<BatchScanInput> &BatchInput);

lib/DependencyScan/DependencyScanningTool.cpp

@@ -15,7 +15,6 @@
 #include "swift/AST/DiagnosticsFrontend.h"
 #include "swift/Basic/LLVMInitialize.h"
 #include "swift/DependencyScan/DependencyScanImpl.h"
-#include "swift/Frontend/Frontend.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
 
@@ -25,8 +24,10 @@ namespace swift {
 namespace dependencies {
 
 DependencyScanningTool::DependencyScanningTool()
-    : SharedCache(std::make_unique<ModuleDependenciesCache>()), PDC(), Alloc(),
-      Saver(Alloc) {}
+    : SharedCache(std::make_unique<ModuleDependenciesCache>()),
+      VersionedPCMInstanceCacheCache(
+          std::make_unique<CompilerArgInstanceCacheMap>()),
+      PDC(), Alloc(), Saver(Alloc) {}
 
 llvm::ErrorOr<swiftscan_dependency_graph_t>
 DependencyScanningTool::getDependencies(
@@ -76,8 +77,9 @@ DependencyScanningTool::getDependencies(
         BatchInput.size(), std::make_error_code(std::errc::invalid_argument));
   auto Instance = std::move(*InstanceOrErr);
 
-  auto batchScanResults =
-      performBatchModuleScan(*Instance.get(), *SharedCache, Saver, BatchInput);
+  auto batchScanResults = performBatchModuleScan(
+      *Instance.get(), *SharedCache, VersionedPCMInstanceCacheCache.get(),
+      Saver, BatchInput);
 
   return batchScanResults;
 }

lib/DependencyScan/ScanDependencies.cpp

@@ -925,53 +925,59 @@ using CompilerArgInstanceCacheMap =
                               std::unique_ptr<ModuleDependenciesCache>>>;
 
 static bool
-forEachBatchEntry(CompilerInstance &instance, ModuleDependenciesCache &cache,
+forEachBatchEntry(CompilerInstance &invocationInstance,
+                  ModuleDependenciesCache &invocationCache,
+                  CompilerArgInstanceCacheMap *versionedPCMInstanceCache,
                   llvm::StringSaver &saver,
                   const std::vector<BatchScanInput> &batchInput,
                   llvm::function_ref<void(BatchScanInput, CompilerInstance &,
                                           ModuleDependenciesCache &)>
                       scanningAction) {
-  const CompilerInvocation &invok = instance.getInvocation();
-  // Keep track of all compiler instances and dependency caches we have
-  // created.
-  // TODO: Re-use a single cache across all invocations, once `alreadySeen`
-  // state is no longer shared.
-  CompilerArgInstanceCacheMap subInstanceMap;
-  auto &diags = instance.getDiags();
-  ForwardingDiagnosticConsumer FDC(instance.getDiags());
+  const CompilerInvocation &invok = invocationInstance.getInvocation();
+  bool localSubInstanceMap = false;
+  CompilerArgInstanceCacheMap *subInstanceMap;
+  if (versionedPCMInstanceCache)
+    subInstanceMap = versionedPCMInstanceCache;
+  else {
+    subInstanceMap = new CompilerArgInstanceCacheMap;
+    localSubInstanceMap = true;
+  }
+
+  auto &diags = invocationInstance.getDiags();
+  ForwardingDiagnosticConsumer FDC(invocationInstance.getDiags());
 
   for (auto &entry : batchInput) {
     CompilerInstance *pInstance = nullptr;
     ModuleDependenciesCache *pCache = nullptr;
     if (entry.arguments.empty()) {
       // Use the compiler's instance if no arguments are specified.
-      pInstance = &instance;
-      pCache = &cache;
-    } else if (subInstanceMap.count(entry.arguments)) {
+      pInstance = &invocationInstance;
+      pCache = &invocationCache;
+    } else if (subInstanceMap->count(entry.arguments)) {
       // Use the previously created instance if we've seen the arguments
       // before.
-      pInstance = subInstanceMap[entry.arguments].first.get();
-      pCache = subInstanceMap[entry.arguments].second.get();
+      pInstance = (*subInstanceMap)[entry.arguments].first.get();
+      pCache = (*subInstanceMap)[entry.arguments].second.get();
     } else {
       // Create a new instance by the arguments and save it in the map.
-      subInstanceMap.insert(
+      subInstanceMap->insert(
           {entry.arguments,
            std::make_pair(std::make_unique<CompilerInstance>(),
                           std::make_unique<ModuleDependenciesCache>())});
 
-      pInstance = subInstanceMap[entry.arguments].first.get();
-      pCache = subInstanceMap[entry.arguments].second.get();
+      pInstance = (*subInstanceMap)[entry.arguments].first.get();
+      pCache = (*subInstanceMap)[entry.arguments].second.get();
       SmallVector<const char *, 4> args;
       llvm::cl::TokenizeGNUCommandLine(entry.arguments, saver, args);
       CompilerInvocation subInvok = invok;
       pInstance->addDiagnosticConsumer(&FDC);
       if (subInvok.parseArgs(args, diags)) {
-        instance.getDiags().diagnose(
+        invocationInstance.getDiags().diagnose(
             SourceLoc(), diag::scanner_arguments_invalid, entry.arguments);
         return true;
       }
       if (pInstance->setup(subInvok)) {
-        instance.getDiags().diagnose(
+        invocationInstance.getDiags().diagnose(
             SourceLoc(), diag::scanner_arguments_invalid, entry.arguments);
         return true;
       }
@@ -980,6 +986,9 @@ forEachBatchEntry(CompilerInstance &instance, ModuleDependenciesCache &cache,
     assert(pCache);
     scanningAction(entry, *pInstance, *pCache);
   }
+
+  if (localSubInstanceMap)
+    delete subInstanceMap;
   return false;
 }
 
@@ -1132,8 +1141,9 @@ bool swift::dependencies::batchScanDependencies(
   if (!batchInput.hasValue())
     return true;
 
-  auto batchScanResults =
-      performBatchModuleScan(instance, cache, saver, *batchInput);
+  auto batchScanResults = performBatchModuleScan(
+      instance, cache, /*versionedPCMInstanceCache*/ nullptr, saver,
+      *batchInput);
 
   // Write the result JSON to the specified output path, for each entry
   auto ientries = batchInput->cbegin();
@@ -1268,14 +1278,17 @@ swift::dependencies::performModulePrescan(CompilerInstance &instance) {
 
 std::vector<llvm::ErrorOr<swiftscan_dependency_graph_t>>
 swift::dependencies::performBatchModuleScan(
-    CompilerInstance &instance, ModuleDependenciesCache &cache,
+    CompilerInstance &invocationInstance,
+    ModuleDependenciesCache &invocationCache,
+    CompilerArgInstanceCacheMap *versionedPCMInstanceCache,
     llvm::StringSaver &saver, const std::vector<BatchScanInput> &batchInput) {
   std::vector<llvm::ErrorOr<swiftscan_dependency_graph_t>> batchScanResult;
   batchScanResult.reserve(batchInput.size());
 
   // Perform a full dependency scan for each batch entry module
   forEachBatchEntry(
-      instance, cache, saver, batchInput,
+      invocationInstance, invocationCache, versionedPCMInstanceCache, saver,
+      batchInput,
       [&batchScanResult](BatchScanInput entry, CompilerInstance &instance,
                          ModuleDependenciesCache &cache) {
         StringRef moduleName = entry.moduleName;
@@ -1343,7 +1356,7 @@ swift::dependencies::performBatchModulePrescan(
 
   // Perform a full dependency scan for each batch entry module
   forEachBatchEntry(
-      instance, cache, saver, batchInput,
+      instance, cache, /*versionedPCMInstanceCache*/ nullptr, saver, batchInput,
       [&batchPrescanResult](BatchScanInput entry, CompilerInstance &instance,
                             ModuleDependenciesCache &cache) {
         StringRef moduleName = entry.moduleName;