[Dependency Scanning] Keep track of which direct dependencies are '@testable' for a Source dependency

Author: artemcm

include/swift/AST/Decl.h

@@ -1422,6 +1422,10 @@ class ImportDecl final : public Decl,
     return getAttrs().hasAttribute<ExportedAttr>();
   }
 
+  bool isTestable() const {
+    return getAttrs().hasAttribute<TestableAttr>();
+  }
+
   ModuleDecl *getModule() const { return Mod; }
   void setModule(ModuleDecl *M) { Mod = M; }
 

include/swift/AST/ModuleDependencies.h

@@ -211,10 +211,14 @@ class SwiftSourceModuleDependenciesStorage :
   /// Details common to Swift textual (interface or source) modules
   CommonSwiftTextualModuleDependencyDetails textualModuleDetails;
 
+  /// Collection of module imports that were detected to be `@Testable`
+  llvm::StringSet<> testableImports;
+
   SwiftSourceModuleDependenciesStorage(
     ArrayRef<StringRef> extraPCMArgs
   ) : ModuleDependencyInfoStorageBase(ModuleDependencyKind::SwiftSource),
-      textualModuleDetails(extraPCMArgs) {}
+      textualModuleDetails(extraPCMArgs),
+      testableImports(llvm::StringSet<>()) {}
 
   ModuleDependencyInfoStorageBase *clone() const override {
     return new SwiftSourceModuleDependenciesStorage(*this);
@@ -223,6 +227,10 @@ class SwiftSourceModuleDependenciesStorage :
   static bool classof(const ModuleDependencyInfoStorageBase *base) {
     return base->dependencyKind == ModuleDependencyKind::SwiftSource;
   }
+
+  void addTestableImport(ImportPath::Module module) {
+    testableImports.insert(module.front().Item.str());
+  }
 };
 
 /// Describes the dependencies of a pre-built Swift module (with no .swiftinterface).
@@ -460,6 +468,13 @@ class ModuleDependencyInfo {
     storage->resolved = isResolved;
   }
 
+  /// For a Source dependency, register a `Testable` import
+  void addTestableImport(ImportPath::Module module);
+
+  /// Whether or not a queried module name is a `@Testable` import dependency
+  /// of this module. Can only return `true` for Swift source modules.
+  bool isTestableImport(StringRef moduleName) const;
+
   /// Whether the dependencies are for a Swift module: either Textual, Source, Binary, or Placeholder.
   bool isSwiftModule() const;
 

lib/AST/ModuleDependencies.cpp

@@ -85,6 +85,18 @@ ModuleDependencyInfo::getAsPlaceholderDependencyModule() const {
   return dyn_cast<SwiftPlaceholderModuleDependencyStorage>(storage.get());
 }
 
+void ModuleDependencyInfo::addTestableImport(ImportPath::Module module) {
+  assert(getAsSwiftSourceModule() && "Expected source module for addTestableImport.");
+  dyn_cast<SwiftSourceModuleDependenciesStorage>(storage.get())->addTestableImport(module);
+}
+
+bool ModuleDependencyInfo::isTestableImport(StringRef moduleName) const {
+  if (auto swiftSourceDepStorage = getAsSwiftSourceModule())
+    return swiftSourceDepStorage->testableImports.contains(moduleName);
+  else
+    return false;
+}
+
 void ModuleDependencyInfo::addModuleDependency(ModuleDependencyID dependencyID) {
   storage->resolvedModuleDependencies.push_back(dependencyID);
 }
@@ -114,6 +126,10 @@ void ModuleDependencyInfo::addModuleImport(
     ImportPath::Builder scratch;
     auto realPath = importDecl->getRealModulePath(scratch);
     addModuleImport(realPath, &alreadyAddedModules);
+
+    if (getKind() == swift::ModuleDependencyKind::SwiftSource &&
+        importDecl->isTestable())
+      addTestableImport(realPath);
   }
 
   auto fileName = sf.getFilename();