[SymbolGraphGen] Emit symbols from exported modules

When emitting a symbol graph file for a module that import modules via
`@_exported import`, emits those modules' symbols as well.

SR-15753

rdar://89547374

Author: franklinsch

include/swift/AST/Module.h

@@ -939,6 +939,9 @@ inline SourceLoc extractNearestSourceLoc(const ModuleDecl *mod) {
   return extractNearestSourceLoc(static_cast<const Decl *>(mod));
 }
 
+/// Collects modules that this module imports via `@_exported import`.
+void collectParsedExportedImports(const ModuleDecl *M, SmallPtrSetImpl<ModuleDecl *> &Imports);
+
 } // end namespace swift
 
 #endif

lib/AST/Module.cpp

@@ -788,7 +788,7 @@ bool ModuleDecl::shouldCollectDisplayDecls() const {
   return true;
 }
 
-static void collectParsedExportedImports(const ModuleDecl *M, SmallPtrSetImpl<ModuleDecl *> &Imports) {
+void swift::collectParsedExportedImports(const ModuleDecl *M, SmallPtrSetImpl<ModuleDecl *> &Imports) {
   for (const FileUnit *file : M->getFiles()) {
     if (const SourceFile *source = dyn_cast<SourceFile>(file)) {
       if (source->hasImports()) {

lib/SymbolGraphGen/SymbolGraph.cpp

@@ -674,12 +674,12 @@ bool SymbolGraph::isUnconditionallyUnavailableOnAllPlatforms(const Decl *D) cons
 
 /// Returns `true` if the symbol should be included as a node in the graph.
 bool SymbolGraph::canIncludeDeclAsNode(const Decl *D) const {
-  // If this decl isn't in this module, don't record it,
+  // If this decl isn't in this module or module that this module imported with `@_exported`, don't record it,
   // as it will appear elsewhere in its module's symbol graph.
-  if (D->getModuleContext()->getName() != M.getName()) {
+  if (D->getModuleContext()->getName() != M.getName() && !Walker.isFromExportedImportedModule(D)) {
     return false;
   }
-
+  
   if (!isa<ValueDecl>(D)) {
     return false;
   }

lib/SymbolGraphGen/SymbolGraphASTWalker.cpp

@@ -22,9 +22,11 @@ using namespace swift;
 using namespace symbolgraphgen;
 
 SymbolGraphASTWalker::SymbolGraphASTWalker(ModuleDecl &M,
+                                           const SmallPtrSet<ModuleDecl *, 4> ExportedImportedModules,
                                            const SymbolGraphOptions &Options)
   : Options(Options),
     M(M),
+    ExportedImportedModules(ExportedImportedModules),
     MainGraph(*this, M, None, Ctx) {}
 
 /// Get a "sub" symbol graph for the parent module of a type that
@@ -51,6 +53,11 @@ SymbolGraph *SymbolGraphASTWalker::getModuleSymbolGraph(const Decl *D) {
     // should put actual extensions of that module into the main graph
     return &MainGraph;
   }
+  
+  if (isFromExportedImportedModule(D)) {
+    return &MainGraph;
+  }
+  
   auto Found = ExtendedModuleGraphs.find(M->getNameStr());
   if (Found != ExtendedModuleGraphs.end()) {
     return Found->getValue();
@@ -208,3 +215,11 @@ bool SymbolGraphASTWalker::walkToDeclPre(Decl *D, CharSourceRange Range) {
 
   return true;
 }
+
+bool SymbolGraphASTWalker::isFromExportedImportedModule(const Decl* D) const {
+    auto *M = D->getModuleContext();
+    
+    return llvm::any_of(ExportedImportedModules, [&M](const auto *MD) {
+        return M->getNameStr().equals(MD->getModuleContext()->getNameStr());
+    });
+}

lib/SymbolGraphGen/SymbolGraphASTWalker.h

@@ -46,6 +46,8 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
 
   /// The module that this symbol graph will represent.
   const ModuleDecl &M;
+    
+  const SmallPtrSet<ModuleDecl *, 4> ExportedImportedModules;
 
   /// The symbol graph for the main module of interest.
   SymbolGraph MainGraph;
@@ -55,7 +57,9 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
 
   // MARK: - Initialization
 
-  SymbolGraphASTWalker(ModuleDecl &M, const SymbolGraphOptions &Options);
+  SymbolGraphASTWalker(ModuleDecl &M,
+                       const SmallPtrSet<ModuleDecl *, 4> ExportedImportedModules,
+                       const SymbolGraphOptions &Options);
   virtual ~SymbolGraphASTWalker() {}
 
   // MARK: - Utilities
@@ -87,6 +91,11 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
   // MARK: - SourceEntityWalker
 
   virtual bool walkToDeclPre(Decl *D, CharSourceRange Range) override;
+    
+  // MARK: - Utilities
+  
+  /// Returns whether the given declaration comes from an `@_exported import` module.
+  virtual bool isFromExportedImportedModule(const Decl *D) const;
 };
 
 } // end namespace symbolgraphgen

lib/SymbolGraphGen/SymbolGraphGen.cpp

@@ -55,13 +55,17 @@ int serializeSymbolGraph(SymbolGraph &SG,
 int
 symbolgraphgen::emitSymbolGraphForModule(ModuleDecl *M,
                                          const SymbolGraphOptions &Options) {
-  SymbolGraphASTWalker Walker(*M, Options);
   SmallVector<Decl *, 64> ModuleDecls;
   swift::getTopLevelDeclsForDisplay(M, ModuleDecls, /*recursive*/true);
+  
+  SmallPtrSet<ModuleDecl *, 4> ExportedImportedModules;
+  swift::collectParsedExportedImports(M, ExportedImportedModules);
 
   if (Options.PrintMessages)
     llvm::errs() << ModuleDecls.size()
         << " top-level declarations in this module.\n";
+    
+  SymbolGraphASTWalker Walker(*M, ExportedImportedModules, Options);
 
   for (auto *Decl : ModuleDecls) {
     Walker.walk(Decl);
@@ -98,7 +102,7 @@ printSymbolGraphForDecl(const ValueDecl *D, Type BaseTy,
 
   llvm::json::OStream JOS(OS, Options.PrettyPrint ? 2 : 0);
   ModuleDecl *MD = D->getModuleContext();
-  SymbolGraphASTWalker Walker(*MD, Options);
+  SymbolGraphASTWalker Walker(*MD, {}, Options);
   markup::MarkupContext MarkupCtx;
   SymbolGraph Graph(Walker, *MD, None, MarkupCtx, None,
                     /*IsForSingleNode=*/true);

test/SymbolGraph/Module/BasicExtension.swift

@@ -1,7 +1,11 @@
 // RUN: %empty-directory(%t)
 // RUN: %target-build-swift %s -module-name BasicExtension -emit-module -emit-module-path %t/
 // RUN: %target-swift-symbolgraph-extract -module-name BasicExtension -I %t -pretty-print -output-dir %t
-// RUN: %FileCheck %s --input-file %t/BasicExtension@Swift.symbols.json
+// RUN: %FileCheck %s --input-file %t/BasicExtension@Swift.symbols.json --check-prefix EXTRACT
+
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %s -module-name BasicExtension -emit-module -emit-module-path %t/ -emit-symbol-graph -emit-symbol-graph-dir %t
+// RUN: %FileCheck %s --input-file %t/BasicExtension@Swift.symbols.json --check-prefix BUILD
 
 extension String {
   /// Return something.
@@ -10,17 +14,28 @@ extension String {
   }
 }
 
-// CHECK: module
-// CHECK-NEXT: "name": "BasicExtension"
+// EXTRACT: module
+// EXTRACT-NEXT: "name": "BasicExtension"
+
+// BUILD: module
+// BUILD: "name":"BasicExtension"
 
-// CHECK: "precise": "s:SS14BasicExtensionE9somethingSSvp"
+// EXTRACT: "precise": "s:SS14BasicExtensionE9somethingSSvp"
 
-// CHECK: "kind": "memberOf"
-// CHECK-NEXT: "source": "s:SS14BasicExtensionE9somethingSSvp"
-// CHECK-NEXT: "target": "s:SS"
+// BUILD: "precise":"s:SS14BasicExtensionE9somethingSSvp"
+
+// EXTRACT: "kind": "memberOf"
+// EXTRACT-NEXT: "source": "s:SS14BasicExtensionE9somethingSSvp"
+// EXTRACT-NEXT: "target": "s:SS"
+
+// BUILD: "kind":"memberOf"
+// BUILD: "source":"s:SS14BasicExtensionE9somethingSSvp"
+// BUILD: "target":"s:SS"
 
 // Extending `String` creates a memberOf relationship above.
 // However, it should not be included as a node because `String`
 // is owned by the Swift module.
 // rdar://58876107
-// CHECK-NOT: "precise": "s:SS"
+// EXTRACT-NOT: "precise": "s:SS"
+
+// BUILD-NOT: "precise":"s:SS"

test/SymbolGraph/Module/ExportedImport.swift

@@ -0,0 +1,15 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend %S/Inputs/ExportedImport/A.swift -module-name A -emit-module -emit-module-path %t/A.swiftmodule
+// RUN: %target-swift-frontend %s -module-name ExportedImport -emit-module -emit-module-path /dev/null -I %t -emit-symbol-graph -emit-symbol-graph-dir %t/
+// RUN: %FileCheck %s --input-file %t/ExportedImport.symbols.json
+// RUN: ls %t | %FileCheck %s --check-prefix FILES
+
+@_exported import A
+
+// CHECK: "precise":"s:1A11SymbolFromAV"
+// CHECK-NOT: InternalSymbolFromA
+
+// FIXME: Symbols from `@_exported import` do not get emitted when using swift-symbolgraph-extract
+// This is tracked by https://bugs.swift.org/browse/SR-15921.
+
+// FILES-NOT: ExportedImport@A.symbols.json

test/SymbolGraph/Module/Inputs/ExportedImport/A.swift

@@ -0,0 +1,3 @@
+public struct SymbolFromA {}
+
+struct InternalSymbolFromA {}