[Caching] Support CrossImport modules for caching build

Add support for cross import modules by ingesting swiftoverlay files for
the cross import into CAS file system.

The long-term better fix will be just passing the cross import
information from scanner to swift-frontend so frontend doesn't need to
read overlay files again to figure out the cross import module.

rdar://123839248

Author: cachemeifyoucan

include/swift/AST/ModuleDependencies.h

@@ -172,6 +172,10 @@ class ModuleDependencyInfoStorageBase {
   /// The cache key for the produced module.
   std::string moduleCacheKey;
 
+  /// Auxiliary files that help to construct other dependencies (e.g.
+  /// command-line), no need to be saved to reconstruct from cache.
+  std::vector<std::string> auxiliaryFiles;
+
   /// The direct dependency of the module is resolved by scanner.
   bool resolved;
   /// ModuleDependencyInfo is finalized (with all transitive dependencies
@@ -635,6 +639,24 @@ class ModuleDependencyInfo {
     llvm_unreachable("Unexpected type");
   }
 
+  void addAuxiliaryFile(const std::string &file) {
+    storage->auxiliaryFiles.emplace_back(file);
+  }
+
+  void updateCASFileSystemRootID(const std::string &rootID) {
+    if (isSwiftInterfaceModule())
+      cast<SwiftInterfaceModuleDependenciesStorage>(storage.get())
+          ->textualModuleDetails.CASFileSystemRootID = rootID;
+    else if (isSwiftSourceModule())
+      cast<SwiftSourceModuleDependenciesStorage>(storage.get())
+          ->textualModuleDetails.CASFileSystemRootID = rootID;
+    else if (isClangModule())
+      cast<ClangModuleDependencyStorage>(storage.get())->CASFileSystemRootID =
+          rootID;
+    else
+      llvm_unreachable("Unexpected type");
+  }
+
   bool isResolved() const {
     return storage->resolved;
   }
@@ -764,7 +786,8 @@ class ModuleDependencyInfo {
   /// Collect a map from a secondary module name to a list of cross-import
   /// overlays, when this current module serves as the primary module.
   llvm::StringMap<llvm::SmallSetVector<Identifier, 4>>
-  collectCrossImportOverlayNames(ASTContext &ctx, StringRef moduleName) const;
+  collectCrossImportOverlayNames(ASTContext &ctx, StringRef moduleName,
+                                 std::vector<std::string> &overlayFiles) const;
 };
 
 using ModuleDependencyVector = llvm::SmallVector<std::pair<ModuleDependencyID, ModuleDependencyInfo>, 1>;

include/swift/DependencyScan/ModuleDependencyScanner.h

@@ -72,9 +72,8 @@ class ModuleDependencyScanner {
                           DiagnosticEngine &diags, bool ParallelScan);
 
   /// Identify the scanner invocation's main module's dependencies
-  llvm::ErrorOr<ModuleDependencyInfo> getMainModuleDependencyInfo(
-      ModuleDecl *mainModule,
-      std::optional<SwiftDependencyTracker> tracker = std::nullopt);
+  llvm::ErrorOr<ModuleDependencyInfo>
+  getMainModuleDependencyInfo(ModuleDecl *mainModule);
 
   /// Resolve module dependencies of the given module, computing a full
   /// transitive closure dependency graph.

include/swift/Serialization/ScanningLoaders.h

@@ -41,20 +41,17 @@ class SwiftModuleScanner : public SerializedModuleLoaderBase {
   /// Location where pre-built moduels are to be built into.
   std::string moduleOutputPath;
 
-  std::optional<SwiftDependencyTracker> dependencyTracker;
-
 public:
   std::optional<ModuleDependencyInfo> dependencies;
 
-  SwiftModuleScanner(
-      ASTContext &ctx, ModuleLoadingMode LoadMode, Identifier moduleName,
-      InterfaceSubContextDelegate &astDelegate, StringRef moduleOutputPath,
-      ScannerKind kind = MDS_plain,
-      std::optional<SwiftDependencyTracker> tracker = std::nullopt)
+  SwiftModuleScanner(ASTContext &ctx, ModuleLoadingMode LoadMode,
+                     Identifier moduleName,
+                     InterfaceSubContextDelegate &astDelegate,
+                     StringRef moduleOutputPath, ScannerKind kind = MDS_plain)
       : SerializedModuleLoaderBase(ctx, nullptr, LoadMode,
                                    /*IgnoreSwiftSourceInfoFile=*/true),
         kind(kind), moduleName(moduleName), astDelegate(astDelegate),
-        moduleOutputPath(moduleOutputPath), dependencyTracker(tracker) {}
+        moduleOutputPath(moduleOutputPath) {}
 
   std::error_code findModuleFilesInDirectory(
       ImportPath::Element ModuleID, const SerializedModuleBaseName &BaseName,
@@ -91,13 +88,13 @@ class PlaceholderSwiftModuleScanner : public SwiftModuleScanner {
   llvm::BumpPtrAllocator Allocator;
 
 public:
-  PlaceholderSwiftModuleScanner(
-      ASTContext &ctx, ModuleLoadingMode LoadMode, Identifier moduleName,
-      StringRef PlaceholderDependencyModuleMap,
-      InterfaceSubContextDelegate &astDelegate, StringRef moduleOutputPath,
-      std::optional<SwiftDependencyTracker> tracker = std::nullopt)
+  PlaceholderSwiftModuleScanner(ASTContext &ctx, ModuleLoadingMode LoadMode,
+                                Identifier moduleName,
+                                StringRef PlaceholderDependencyModuleMap,
+                                InterfaceSubContextDelegate &astDelegate,
+                                StringRef moduleOutputPath)
       : SwiftModuleScanner(ctx, LoadMode, moduleName, astDelegate,
-                           moduleOutputPath, MDS_placeholder, tracker) {
+                           moduleOutputPath, MDS_placeholder) {
 
     // FIXME: Find a better place for this map to live, to avoid
     // doing the parsing on every module.

lib/AST/ModuleLoader.cpp

@@ -186,8 +186,9 @@ void ModuleLoader::findOverlayFiles(SourceLoc diagLoc, ModuleDecl *module,
 }
 
 llvm::StringMap<llvm::SmallSetVector<Identifier, 4>>
-ModuleDependencyInfo::collectCrossImportOverlayNames(ASTContext &ctx,
-                                                     StringRef moduleName) const {
+ModuleDependencyInfo::collectCrossImportOverlayNames(
+    ASTContext &ctx, StringRef moduleName,
+    std::vector<std::string> &overlayFiles) const {
   using namespace llvm::sys;
   using namespace file_types;
   std::optional<std::string> modulePath;
@@ -239,6 +240,7 @@ ModuleDependencyInfo::collectCrossImportOverlayNames(ASTContext &ctx,
       ModuleDecl::collectCrossImportOverlay(ctx, file, moduleName,
                                             bystandingModule);
     result[bystandingModule] = std::move(overlayNames);
+    overlayFiles.push_back(file.str());
   });
   return result;
 }

lib/DependencyScan/ModuleDependencyScanner.cpp

@@ -326,8 +326,7 @@ static void findAllImportedClangModules(StringRef moduleName,
 }
 
 llvm::ErrorOr<ModuleDependencyInfo>
-ModuleDependencyScanner::getMainModuleDependencyInfo(
-    ModuleDecl *mainModule, std::optional<SwiftDependencyTracker> tracker) {
+ModuleDependencyScanner::getMainModuleDependencyInfo(ModuleDecl *mainModule) {
   // Main module file name.
   auto newExt = file_types::getExtension(file_types::TY_SwiftModuleFile);
   llvm::SmallString<32> mainModulePath = mainModule->getName().str();
@@ -345,37 +344,19 @@ ModuleDependencyScanner::getMainModuleDependencyInfo(
         ExtraPCMArgs.begin(),
         {"-Xcc", "-target", "-Xcc", ScanASTContext.LangOpts.Target.str()});
 
-  std::string rootID;
-  if (tracker) {
-    tracker->startTracking();
-    for (auto fileUnit : mainModule->getFiles()) {
-      auto sf = dyn_cast<SourceFile>(fileUnit);
-      if (!sf)
-        continue;
-      tracker->trackFile(sf->getFilename());
-    }
-    tracker->addCommonSearchPathDeps(
-        ScanCompilerInvocation.getSearchPathOptions());
-    // Fetch some dependency files from clang importer.
-    std::vector<std::string> clangDependencyFiles;
+  auto mainDependencies =
+      ModuleDependencyInfo::forSwiftSourceModule({}, {}, {}, ExtraPCMArgs);
+
+  if (ScanASTContext.CASOpts.EnableCaching) {
     auto clangImporter =
         static_cast<ClangImporter *>(ScanASTContext.getClangModuleLoader());
+    std::vector<std::string> clangDependencyFiles;
     clangImporter->addClangInvovcationDependencies(clangDependencyFiles);
-    llvm::for_each(clangDependencyFiles,
-                   [&](std::string &file) { tracker->trackFile(file); });
-
-    auto root = tracker->createTreeFromDependencies();
-    if (!root) {
-      Diagnostics.diagnose(SourceLoc(), diag::error_cas,
-                           toString(root.takeError()));
-      return std::make_error_code(std::errc::io_error);
-    }
-    rootID = root->getID().toString();
+    llvm::for_each(clangDependencyFiles, [&](std::string &file) {
+      mainDependencies.addAuxiliaryFile(file);
+    });
   }
 
-  auto mainDependencies =
-      ModuleDependencyInfo::forSwiftSourceModule(rootID, {}, {}, ExtraPCMArgs);
-
   llvm::StringSet<> alreadyAddedModules;
   // Compute Implicit dependencies of the main module
   {
@@ -801,6 +782,7 @@ void ModuleDependencyScanner::discoverCrossImportOverlayDependencies(
     llvm::function_ref<void(ModuleDependencyID)> action) {
   // Modules explicitly imported. Only these can be secondary module.
   llvm::SetVector<Identifier> newOverlays;
+  std::vector<std::string> overlayFiles;
   for (auto dep : allDependencies) {
     auto moduleName = dep.ModuleName;
     // Do not look for overlays of main module under scan
@@ -810,7 +792,7 @@ void ModuleDependencyScanner::discoverCrossImportOverlayDependencies(
     auto dependencies = cache.findDependency(moduleName, dep.Kind).value();
     // Collect a map from secondary module name to cross-import overlay names.
     auto overlayMap = dependencies->collectCrossImportOverlayNames(
-        ScanASTContext, moduleName);
+        ScanASTContext, moduleName, overlayFiles);
     if (overlayMap.empty())
       continue;
     for (const auto &dependencyId : allDependencies) {
@@ -872,11 +854,9 @@ void ModuleDependencyScanner::discoverCrossImportOverlayDependencies(
 
   // Update main module's dependencies to include these new overlays.
   auto resolvedDummyDep =
-      *(cache.findDependency(dummyMainName, ModuleDependencyKind::SwiftSource)
-            .value());
+      **cache.findDependency(dummyMainName, ModuleDependencyKind::SwiftSource);
   auto mainDep =
-      *(cache.findDependency(mainModuleName, ModuleDependencyKind::SwiftSource)
-            .value());
+      **cache.findDependency(mainModuleName, ModuleDependencyKind::SwiftSource);
   auto newOverlayDeps = resolvedDummyDep.getDirectModuleDependencies();
   auto existingMainDeps = mainDep.getDirectModuleDependencies();
   ModuleDependencyIDSet existingMainDepsSet(existingMainDeps.begin(),
@@ -888,6 +868,11 @@ void ModuleDependencyScanner::discoverCrossImportOverlayDependencies(
                   if (!existingMainDepsSet.count(crossImportOverlayModID))
                     mainDep.addModuleDependency(crossImportOverlayModID);
                 });
+
+  llvm::for_each(overlayFiles, [&mainDep](const std::string &file) {
+    mainDep.addAuxiliaryFile(file);
+  });
+
   cache.updateDependency(
       {mainModuleName.str(), ModuleDependencyKind::SwiftSource}, mainDep);
 

lib/DependencyScan/ScanDependencies.cpp

@@ -207,13 +207,11 @@ static llvm::Error resolveExplicitModuleInputs(
   auto &service = cache.getScanService();
   auto remapPath = [&](StringRef path) { return service.remapPath(path); };
   std::vector<std::string> rootIDs;
-  if (auto ID = resolvingDepInfo.getCASFSRootID())
-    rootIDs.push_back(*ID);
-
   std::vector<std::string> includeTrees;
   if (auto ID = resolvingDepInfo.getClangIncludeTree())
     includeTrees.push_back(*ID);
 
+  auto tracker = cache.getScanService().createSwiftDependencyTracker();
   auto addBridgingHeaderDeps =
       [&](const ModuleDependencyInfo &depInfo) -> llvm::Error {
     auto sourceDepDetails = depInfo.getAsSwiftSourceModule();
@@ -223,8 +221,7 @@ static llvm::Error resolveExplicitModuleInputs(
     if (sourceDepDetails->textualModuleDetails
             .CASBridgingHeaderIncludeTreeRootID.empty()) {
       if (!sourceDepDetails->textualModuleDetails.bridgingSourceFiles.empty()) {
-        if (auto tracker =
-                cache.getScanService().createSwiftDependencyTracker()) {
+        if (tracker) {
           tracker->startTracking();
           for (auto &file :
                sourceDepDetails->textualModuleDetails.bridgingSourceFiles)
@@ -342,10 +339,36 @@ static llvm::Error resolveExplicitModuleInputs(
     // Merge CASFS from clang dependency.
     auto &CASFS = cache.getScanService().getSharedCachingFS();
     auto &CAS = CASFS.getCAS();
+    assert(tracker && "no caching tracker is available");
 
     // Update build command line.
     if (resolvingDepInfo.isSwiftInterfaceModule() ||
         resolvingDepInfo.isSwiftSourceModule()) {
+      // Compute the CASFS root ID for the resolving dependency.
+      tracker->startTracking();
+      tracker->addCommonSearchPathDeps(
+          instance.getInvocation().getSearchPathOptions());
+      if (auto *sourceDep = resolvingDepInfo.getAsSwiftSourceModule()) {
+        llvm::for_each(
+            sourceDep->sourceFiles,
+            [&tracker](const std::string &file) { tracker->trackFile(file); });
+        llvm::for_each(
+            sourceDep->auxiliaryFiles,
+            [&tracker](const std::string &file) { tracker->trackFile(file); });
+      } else if (auto *textualDep =
+                     resolvingDepInfo.getAsSwiftInterfaceModule()) {
+        tracker->trackFile(textualDep->swiftInterfaceFile);
+        llvm::for_each(
+            textualDep->auxiliaryFiles,
+            [&tracker](const std::string &file) { tracker->trackFile(file); });
+      }
+      auto root = tracker->createTreeFromDependencies();
+      if (!root)
+        return root.takeError();
+      auto rootID = root->getID().toString();
+      dependencyInfoCopy.updateCASFileSystemRootID(rootID);
+      rootIDs.push_back(rootID);
+
       // Update with casfs option.
       std::vector<std::string> newCommandLine =
           dependencyInfoCopy.getCommandline();
@@ -1898,9 +1921,8 @@ swift::dependencies::performModuleScan(CompilerInstance &instance,
 
   // Identify imports of the main module and add an entry for it
   // to the dependency graph.
-  auto mainModuleDepInfo = scanner.getMainModuleDependencyInfo(
-      instance.getMainModule(),
-      cache.getScanService().createSwiftDependencyTracker());
+  auto mainModuleDepInfo =
+      scanner.getMainModuleDependencyInfo(instance.getMainModule());
   auto mainModuleName = instance.getMainModule()->getNameStr();
   auto mainModuleID = ModuleDependencyID{mainModuleName.str(),
                                          ModuleDependencyKind::SwiftSource};

lib/Serialization/ScanningLoaders.cpp

@@ -192,29 +192,21 @@ SwiftModuleScanner::scanInterfaceFile(Twine moduleInterfacePath,
             *moduleDecl, SourceFileKind::Interface, bufferID, parsingOpts);
         moduleDecl->addAuxiliaryFile(*sourceFile);
 
-        std::string RootID;
-        if (dependencyTracker) {
-          dependencyTracker->startTracking();
-          dependencyTracker->addCommonSearchPathDeps(Ctx.SearchPathOpts);
+        std::vector<StringRef> ArgsRefs(Args.begin(), Args.end());
+        Result = ModuleDependencyInfo::forSwiftInterfaceModule(
+            outputPathBase.str().str(), InPath, compiledCandidates, ArgsRefs,
+            PCMArgs, Hash, isFramework, {}, /*module-cache-key*/ "");
+
+        if (Ctx.CASOpts.EnableCaching) {
           std::vector<std::string> clangDependencyFiles;
           auto clangImporter =
               static_cast<ClangImporter *>(Ctx.getClangModuleLoader());
           clangImporter->addClangInvovcationDependencies(clangDependencyFiles);
           llvm::for_each(clangDependencyFiles, [&](std::string &file) {
-            dependencyTracker->trackFile(file);
+            Result->addAuxiliaryFile(file);
           });
-          dependencyTracker->trackFile(moduleInterfacePath);
-          auto RootOrError = dependencyTracker->createTreeFromDependencies();
-          if (!RootOrError)
-            return llvm::errorToErrorCode(RootOrError.takeError());
-          RootID = RootOrError->getID().toString();
         }
 
-        std::vector<StringRef> ArgsRefs(Args.begin(), Args.end());
-        Result = ModuleDependencyInfo::forSwiftInterfaceModule(
-            outputPathBase.str().str(), InPath, compiledCandidates, ArgsRefs,
-            PCMArgs, Hash, isFramework, RootID, /*module-cache-key*/ "");
-
         // Walk the source file to find the import declarations.
         llvm::StringSet<> alreadyAddedModules;
         Result->addModuleImport(*sourceFile, alreadyAddedModules);
@@ -247,9 +239,6 @@ ModuleDependencyVector SerializedModuleLoaderBase::getModuleDependencies(
   ImportPath::Module::Builder builder(moduleName);
   auto modulePath = builder.get();
   auto moduleId = modulePath.front().Item;
-  std::optional<SwiftDependencyTracker> tracker = std::nullopt;
-  if (CacheFS)
-    tracker = SwiftDependencyTracker(*CacheFS, mapper);
 
   // Do not load interface module if it is testable import.
   ModuleLoadingMode MLM =
@@ -264,10 +253,10 @@ ModuleDependencyVector SerializedModuleLoaderBase::getModuleDependencies(
   // FIXME: submodules?
   scanners.push_back(std::make_unique<PlaceholderSwiftModuleScanner>(
       Ctx, MLM, moduleId, Ctx.SearchPathOpts.PlaceholderDependencyModuleMap,
-      delegate, moduleOutputPath, tracker));
+      delegate, moduleOutputPath));
   scanners.push_back(std::make_unique<SwiftModuleScanner>(
       Ctx, MLM, moduleId, delegate, moduleOutputPath,
-      SwiftModuleScanner::MDS_plain, tracker));
+      SwiftModuleScanner::MDS_plain));
 
   // Check whether there is a module with this name that we can import.
   assert(isa<PlaceholderSwiftModuleScanner>(scanners[0].get()) &&