DependenciesScanner: refactor cross import overlay resolver to a standalone function. NFC

This change facilitates resolving cross-import overlays from transitive dependencies.

lib/FrontendTool/ScanDependencies.cpp

@@ -62,23 +62,12 @@ static void findAllImportedClangModules(ASTContext &ctx, StringRef moduleName,
 /// Resolve the direct dependencies of the given module.
 static std::vector<ModuleDependencyID> resolveDirectDependencies(
     CompilerInstance &instance, ModuleDependencyID module,
-    ModuleDependenciesCache &cache) {
+    ModuleDependenciesCache &cache,
+    InterfaceSubContextDelegate &ASTDelegate) {
   auto &ctx = instance.getASTContext();
   auto knownDependencies = *cache.findDependencies(module.first, module.second);
   auto isSwift = knownDependencies.isSwiftModule();
-  auto ModuleCachePath = getModuleCachePathFromClang(ctx
-    .getClangModuleLoader()->getClangInstance());
-  auto &FEOpts = instance.getInvocation().getFrontendOptions();
-  ModuleInterfaceLoaderOptions LoaderOpts(FEOpts);
-  InterfaceSubContextDelegateImpl ASTDelegate(ctx.SourceMgr, ctx.Diags,
-                                              ctx.SearchPathOpts, ctx.LangOpts,
-                                              LoaderOpts,
-                                              ctx.getClangModuleLoader(),
-                                              /*buildModuleCacheDirIfAbsent*/false,
-                                              ModuleCachePath,
-                                              FEOpts.PrebuiltModuleCachePath,
-                                              FEOpts.SerializeModuleInterfaceDependencyHashes,
-                                              FEOpts.shouldTrackSystemDependencies());
+
   // Find the dependencies of every module this module directly depends on.
   std::vector<ModuleDependencyID> result;
   for (auto dependsOn : knownDependencies.getModuleDependencies()) {
@@ -138,41 +127,77 @@ static std::vector<ModuleDependencyID> resolveDirectDependencies(
       }
     }
   }
-  // Only resolve cross-import overlays when this is the main module.
-  // For other modules, these overlays are explicitly written.
-  bool isMainModule =
-    instance.getMainModule()->getName().str() == module.first &&
-    module.second == ModuleDependenciesKind::Swift;
-  if (isMainModule) {
-    // Modules explicitly imported. Only these can be secondary module.
-    std::vector<ModuleDependencyID> explicitImports = result;
-    for (unsigned I = 0; I != result.size(); ++I) {
-      auto dep = result[I];
-      auto moduleName = dep.first;
-      auto dependencies = *cache.findDependencies(moduleName, dep.second);
-      // Collect a map from secondary module name to cross-import overlay names.
-      auto overlayMap = dependencies.collectCrossImportOverlayNames(
-        instance.getASTContext(), moduleName);
-      if (overlayMap.empty())
-        continue;
-      std::for_each(explicitImports.begin(), explicitImports.end(),
-                    [&](ModuleDependencyID Id) {
-        // check if any explicitly imported modules can serve as a secondary
-        // module, and add the overlay names to the dependencies list.
-        for (auto overlayName: overlayMap[Id.first]) {
-          if (auto found = ctx.getModuleDependencies(overlayName.str(),
-                                                     /*onlyClangModule=*/false,
-                                                     cache,
-                                                     ASTDelegate)) {
-            result.emplace_back(overlayName.str(), found->getKind());
-          }
-        }
-      });
-    }
-  }
   return result;
 }
 
+static void discoverCrosssImportOverlayDependencies(
+    CompilerInstance &instance, StringRef mainModuleName,
+    ArrayRef<ModuleDependencyID> allDependencies,
+    ModuleDependenciesCache &cache, InterfaceSubContextDelegate &ASTDelegate,
+    llvm::function_ref<void(ModuleDependencyID)> action) {
+  // Modules explicitly imported. Only these can be secondary module.
+  llvm::SetVector<Identifier> newOverlays;
+  for (auto dep: allDependencies) {
+    auto moduleName = dep.first;
+    auto dependencies = *cache.findDependencies(moduleName, dep.second);
+    // Collect a map from secondary module name to cross-import overlay names.
+    auto overlayMap = dependencies.collectCrossImportOverlayNames(
+      instance.getASTContext(), moduleName);
+    if (overlayMap.empty())
+      continue;
+    std::for_each(allDependencies.begin(), allDependencies.end(),
+                  [&](ModuleDependencyID Id) {
+      // check if any explicitly imported modules can serve as a secondary
+      // module, and add the overlay names to the dependencies list.
+      for (auto overlayName: overlayMap[Id.first]) {
+        if (std::find_if(allDependencies.begin(), allDependencies.end(),
+            [&](ModuleDependencyID Id) { return Id.first == overlayName.str(); })
+                == allDependencies.end()) {
+          newOverlays.insert(overlayName);
+        }
+      }
+    });
+  }
+  // No new cross-import overlays are found, return.
+  if (newOverlays.empty())
+    return;
+  // Construct a dummy main to resolve the newly discovered cross import overlays.
+  StringRef dummyMainName = "DummyMainModuleForResolvingCrossImportOverlays";
+  auto dummyMainDependencies = ModuleDependencies::forMainSwiftModule({});
+
+  // Update main module's dependencies to include these new overlays.
+  auto mainDep = *cache.findDependencies(mainModuleName, ModuleDependenciesKind::Swift);
+  std::for_each(newOverlays.begin(), newOverlays.end(), [&](Identifier modName) {
+    dummyMainDependencies.addModuleDependency(modName.str());
+    mainDep.addModuleDependency(modName.str());
+  });
+  cache.updateDependencies({mainModuleName, ModuleDependenciesKind::Swift}, mainDep);
+
+  // Record the dummy main module's direct dependencies. The dummy main module
+  // only directly depend on these newly discovered overlay modules.
+  cache.recordDependencies(dummyMainName, dummyMainDependencies,
+                           ModuleDependenciesKind::Swift);
+  llvm::SetVector<ModuleDependencyID, std::vector<ModuleDependencyID>,
+                  std::set<ModuleDependencyID>> allModules;
+
+  // Seed the all module list from the dummpy main module.
+  allModules.insert({dummyMainName.str(), dummyMainDependencies.getKind()});
+
+  // Explore the dependencies of every module.
+  for (unsigned currentModuleIdx = 0;
+       currentModuleIdx < allModules.size();
+       ++currentModuleIdx) {
+    auto module = allModules[currentModuleIdx];
+    auto discoveredModules = resolveDirectDependencies(instance, module,
+                                                       cache, ASTDelegate);
+    allModules.insert(discoveredModules.begin(), discoveredModules.end());
+  }
+  // Report any discovered modules to the clients, which include all overlays
+  // and their dependencies.
+  std::for_each(/* +1 to exclude dummy main*/allModules.begin() + 1,
+                allModules.end(), action);
+}
+
 /// Write a single JSON field.
 namespace {
   template<typename T>
@@ -259,6 +284,7 @@ namespace {
 static void writeJSON(llvm::raw_ostream &out,
                       CompilerInstance &instance,
                       ModuleDependenciesCache &cache,
+                      InterfaceSubContextDelegate &ASTDelegate,
                       ArrayRef<ModuleDependencyID> allModules) {
   // Write out a JSON description of all of the dependencies.
   out << "{\n";
@@ -277,7 +303,8 @@ static void writeJSON(llvm::raw_ostream &out,
   };
   for (const auto &module : allModules) {
     auto directDependencies = resolveDirectDependencies(
-        instance, ModuleDependencyID(module.first, module.second), cache);
+      instance, ModuleDependencyID(module.first, module.second), cache,
+      ASTDelegate);
 
     // Grab the completed module dependencies.
     auto moduleDeps = *cache.findDependencies(module.first, module.second);
@@ -469,18 +496,18 @@ bool swift::scanDependencies(CompilerInstance &instance) {
       break;
 
     case ImplicitStdlibKind::Stdlib:
-      mainDependencies.addModuleDependency("Swift", alreadyAddedModules);
+      mainDependencies.addModuleDependency("Swift", &alreadyAddedModules);
       break;
     }
 
     // Add any implicit module names.
     for (const auto &moduleName : importInfo.ModuleNames) {
-      mainDependencies.addModuleDependency(moduleName.str(), alreadyAddedModules);
+      mainDependencies.addModuleDependency(moduleName.str(), &alreadyAddedModules);
     }
 
     // Already-loaded, implicitly imported module names.
     for (const auto &module : importInfo.AdditionalModules) {
-      mainDependencies.addModuleDependency(module.first->getNameStr(), alreadyAddedModules);
+      mainDependencies.addModuleDependency(module.first->getNameStr(), &alreadyAddedModules);
     }
 
     // Add the bridging header.
@@ -492,7 +519,7 @@ bool swift::scanDependencies(CompilerInstance &instance) {
     // add a dependency with the same name to trigger the search.
     if (importInfo.ShouldImportUnderlyingModule) {
       mainDependencies.addModuleDependency(mainModule->getName().str(),
-                                           alreadyAddedModules);
+                                           &alreadyAddedModules);
     }
   }
 
@@ -508,18 +535,40 @@ bool swift::scanDependencies(CompilerInstance &instance) {
   cache.recordDependencies(mainModuleName, std::move(mainDependencies),
                            ModuleDependenciesKind::Swift);
 
+  auto &ctx = instance.getASTContext();
+  auto ModuleCachePath = getModuleCachePathFromClang(ctx
+    .getClangModuleLoader()->getClangInstance());
+  auto &FEOpts = instance.getInvocation().getFrontendOptions();
+  ModuleInterfaceLoaderOptions LoaderOpts(FEOpts);
+  InterfaceSubContextDelegateImpl ASTDelegate(ctx.SourceMgr, ctx.Diags,
+                                              ctx.SearchPathOpts, ctx.LangOpts,
+                                              LoaderOpts,
+                                              ctx.getClangModuleLoader(),
+                                              /*buildModuleCacheDirIfAbsent*/false,
+                                              ModuleCachePath,
+                                              FEOpts.PrebuiltModuleCachePath,
+                                              FEOpts.SerializeModuleInterfaceDependencyHashes,
+                                              FEOpts.shouldTrackSystemDependencies());
+
   // Explore the dependencies of every module.
   for (unsigned currentModuleIdx = 0;
        currentModuleIdx < allModules.size();
        ++currentModuleIdx) {
     auto module = allModules[currentModuleIdx];
     auto discoveredModules =
-        resolveDirectDependencies(instance, module, cache);
+        resolveDirectDependencies(instance, module, cache, ASTDelegate);
     allModules.insert(discoveredModules.begin(), discoveredModules.end());
   }
 
+  // We have all explicit imports now, resolve cross import overlays.
+  discoverCrosssImportOverlayDependencies(instance, mainModuleName,
+      /*All transitive dependencies*/allModules.getArrayRef().slice(1), cache,
+      ASTDelegate, [&](ModuleDependencyID id) {
+    allModules.insert(id);
+  });
+
   // Write out the JSON description.
-  writeJSON(out, instance, cache, allModules.getArrayRef());
+  writeJSON(out, instance, cache, ASTDelegate, allModules.getArrayRef());
 
   // Update the dependency tracker.
   if (auto depTracker = instance.getDependencyTracker()) {