[Dependency Scanning] Teach dependency scanner to resolve optional dependencies of a module

Author: artemcm

include/swift/AST/ASTContext.h

@@ -1032,6 +1032,7 @@ class ASTContext final {
       StringRef moduleName,
       ModuleDependenciesCache &cache,
       InterfaceSubContextDelegate &delegate,
+      bool optionalDependencyLookup = false,
       llvm::Optional<std::pair<std::string, swift::ModuleDependencyKind>> dependencyOf = None);
 
   /// Retrieve the module dependencies for the Clang module with the given name.

include/swift/AST/ModuleDependencies.h

@@ -109,6 +109,11 @@ class ModuleDependencyInfoStorageBase {
   /// The set of modules on which this module depends.
   std::vector<std::string> moduleImports;
 
+  /// The set of modules which constitute optional module
+  /// dependencies for this module, such as `@_implementationOnly`
+  /// or `internal` imports.
+  std::vector<std::string> optionalModuleImports;
+
   /// The set of modules on which this module depends, resolved
   /// to Module IDs, qualified by module kind: Swift, Clang, etc.
   std::vector<ModuleDependencyID> resolvedModuleDependencies;
@@ -425,6 +430,11 @@ class ModuleDependencyInfo {
     return storage->moduleImports;
   }
 
+  /// Retrieve the module-level optional imports.
+  ArrayRef<std::string> getOptionalModuleImports() const {
+    return storage->optionalModuleImports;
+  }
+
   /// Retreive the module-level dependencies.
   const ArrayRef<ModuleDependencyID> getModuleDependencies() const {
     assert(storage->resolved);
@@ -488,6 +498,11 @@ class ModuleDependencyInfo {
   const SwiftPlaceholderModuleDependencyStorage *
     getAsPlaceholderDependencyModule() const;
 
+  /// Add a dependency on the given module, if it was not already in the set.
+  void addOptionalModuleImport(StringRef module,
+                               llvm::StringSet<> *alreadyAddedModules = nullptr);
+
+
   /// Add a dependency on the given module, if it was not already in the set.
   void addModuleImport(StringRef module,
                        llvm::StringSet<> *alreadyAddedModules = nullptr);

lib/AST/ASTContext.cpp

@@ -2023,6 +2023,7 @@ static void diagnoseScannerFailure(StringRef moduleName,
 Optional<const ModuleDependencyInfo*> ASTContext::getModuleDependencies(
     StringRef moduleName, ModuleDependenciesCache &cache,
     InterfaceSubContextDelegate &delegate,
+    bool optionalDependencyLookup,
     llvm::Optional<ModuleDependencyID> dependencyOf) {
   // Retrieve the dependencies for this module.
   // Check whether we've cached this result.
@@ -2048,8 +2049,9 @@ Optional<const ModuleDependencyInfo*> ASTContext::getModuleDependencies(
       return dependencies;
   }
 
-  diagnoseScannerFailure(moduleName, Diags, cache,
-                         dependencyOf);
+  if (!optionalDependencyLookup)
+    diagnoseScannerFailure(moduleName, Diags, cache,
+                           dependencyOf);
   return None;
 }
 

lib/AST/ModuleDependencies.cpp

@@ -89,6 +89,12 @@ void ModuleDependencyInfo::addModuleDependency(ModuleDependencyID dependencyID)
   storage->resolvedModuleDependencies.push_back(dependencyID);
 }
 
+void ModuleDependencyInfo::addOptionalModuleImport(
+    StringRef module, llvm::StringSet<> *alreadyAddedModules) {
+  if (!alreadyAddedModules || alreadyAddedModules->insert(module).second)
+    storage->optionalModuleImports.push_back(module.str());
+}
+
 void ModuleDependencyInfo::addModuleImport(
     StringRef module, llvm::StringSet<> *alreadyAddedModules) {
   if (!alreadyAddedModules || alreadyAddedModules->insert(module).second)

lib/DependencyScan/ScanDependencies.cpp

@@ -307,12 +307,12 @@ resolveExplicitModuleInputs(ModuleDependencyID moduleID,
 
 /// Resolve the direct dependencies of the given module.
 static ArrayRef<ModuleDependencyID>
-resolveDirectDependencies(CompilerInstance &instance, ModuleDependencyID module,
+resolveDirectDependencies(CompilerInstance &instance, ModuleDependencyID moduleID,
                           ModuleDependenciesCache &cache,
                           InterfaceSubContextDelegate &ASTDelegate) {
-  PrettyStackTraceStringAction trace("Resolving direct dependencies of: ", module.first);
+  PrettyStackTraceStringAction trace("Resolving direct dependencies of: ", moduleID.first);
   auto &ctx = instance.getASTContext();
-  auto optionalKnownDependencies = cache.findDependency(module.first, module.second);
+  auto optionalKnownDependencies = cache.findDependency(moduleID.first, moduleID.second);
   assert(optionalKnownDependencies.has_value());
   auto knownDependencies = optionalKnownDependencies.value();
 
@@ -329,18 +329,31 @@ resolveDirectDependencies(CompilerInstance &instance, ModuleDependencyID module,
   ModuleDependencyIDSetVector result;
   for (auto dependsOn : knownDependencies->getModuleImports()) {
     // Figure out what kind of module we need.
-    bool onlyClangModule = !isSwift || module.first == dependsOn;
+    bool onlyClangModule = !isSwift || moduleID.first == dependsOn;
     if (onlyClangModule) {
       if (auto found =
               ctx.getClangModuleDependencies(dependsOn, cache, ASTDelegate))
         result.insert({dependsOn, ModuleDependencyKind::Clang});
     } else {
       if (auto found =
-              ctx.getModuleDependencies(dependsOn, cache, ASTDelegate, module))
+              ctx.getModuleDependencies(dependsOn, cache, ASTDelegate,
+                                        /* optionalDependencyLookup */ false,
+                                        moduleID))
         result.insert({dependsOn, found.value()->getKind()});
     }
   }
 
+  // We may have a set of optional dependencies for this module, such as `@_implementationOnly`
+  // imports of a `@Testable` import. Attempt to locate those, but do not fail if they
+  // cannot be found.
+  for (auto optionallyDependsOn : knownDependencies->getOptionalModuleImports()) {
+    if (auto found =
+            ctx.getModuleDependencies(optionallyDependsOn, cache, ASTDelegate,
+                                      /* optionalDependencyLookup */ true,
+                                      moduleID))
+      result.insert({optionallyDependsOn, found.value()->getKind()});
+  }
+
   if (isSwiftInterfaceOrSource) {
     // A record of all of the Clang modules referenced from this Swift module.
     std::vector<std::string> allClangModules;
@@ -350,11 +363,11 @@ resolveDirectDependencies(CompilerInstance &instance, ModuleDependencyID module,
     if (knownDependencies->getBridgingHeader()) {
       auto clangImporter =
           static_cast<ClangImporter *>(ctx.getClangModuleLoader());
-      if (!clangImporter->addBridgingHeaderDependencies(module.first,
-                                                        module.second, cache)) {
+      if (!clangImporter->addBridgingHeaderDependencies(moduleID.first,
+                                                        moduleID.second, cache)) {
         // Grab the updated module dependencies.
         // FIXME: This is such a hack.
-        knownDependencies = *cache.findDependency(module.first, module.second);
+        knownDependencies = *cache.findDependency(moduleID.first, moduleID.second);
 
         // Add the Clang modules referenced from the bridging header to the
         // set of Clang modules we know about.
@@ -389,15 +402,15 @@ resolveDirectDependencies(CompilerInstance &instance, ModuleDependencyID module,
     for (const auto &clangDep : allClangModules) {
       if (auto found =
               ctx.getSwiftModuleDependencies(clangDep, cache, ASTDelegate)) {
-        if (clangDep != module.first)
+        if (clangDep != moduleID.first)
           result.insert({clangDep, found.value()->getKind()});
       }
     }
   }
 
   // Resolve the dependnecy info
-  cache.resolveDependencyImports(module, result.takeVector());
-  return cache.findDependency(module.first, module.second).value()->getModuleDependencies();
+  cache.resolveDependencyImports(moduleID, result.takeVector());
+  return cache.findDependency(moduleID.first, moduleID.second).value()->getModuleDependencies();
 }
 
 static void discoverCrossImportOverlayDependencies(

test/ScanDependencies/bin_mod_import.swift

@@ -17,12 +17,8 @@ import EWrapper
 // CHECK: "modulePath": "{{.*}}EWrapper.swiftmodule"
 // CHECK-NEXT: "directDependencies": [
 // CHECK-NEXT:   {
-// CHECK-NEXT:     "swift": "E"
-// CHECK-NEXT:   },
-// CHECK-NEXT:   {
-// CHECK-NEXT:     "swift": "Swift"
-// CHECK-NEXT:   },
-// CHECK-NEXT:   {
-// CHECK-NEXT:     "swift": "SwiftOnoneSupport"
+// CHECK-DAG:     "swift": "E"
+// CHECK-DAG:     "swift": "Swift"
+// CHECK-DAG:     "swift": "SwiftOnoneSupport"
 // CHECK-NEXT:   }
 // CHECK-NEXT: ],