[SymbolGraphGen] only include the given symbol for qualified imports (#59852)

* only include the given symbol for qualified imports

rdar://96309088

* re-exporting one type should not allow unrelated types to sneak in

* ensure that children of re-exported types are also re-exported

Author: QuietMisdreavus

include/swift/AST/Module.h

@@ -949,7 +949,9 @@ inline SourceLoc extractNearestSourceLoc(const ModuleDecl *mod) {
 }
 
 /// Collects modules that this module imports via `@_exported import`.
-void collectParsedExportedImports(const ModuleDecl *M, SmallPtrSetImpl<ModuleDecl *> &Imports);
+void collectParsedExportedImports(const ModuleDecl *M,
+                                  SmallPtrSetImpl<ModuleDecl *> &Imports,
+                                  llvm::SmallDenseMap<ModuleDecl *, SmallPtrSet<Decl *, 4>, 4> &QualifiedImports);
 
 } // end namespace swift
 

lib/AST/Module.cpp

@@ -792,15 +792,33 @@ bool ModuleDecl::shouldCollectDisplayDecls() const {
   return true;
 }
 
-void swift::collectParsedExportedImports(const ModuleDecl *M, SmallPtrSetImpl<ModuleDecl *> &Imports) {
+void swift::collectParsedExportedImports(const ModuleDecl *M,
+                                         SmallPtrSetImpl<ModuleDecl *> &Imports,
+                                         llvm::SmallDenseMap<ModuleDecl *, SmallPtrSet<Decl *, 4>, 4> &QualifiedImports) {
   for (const FileUnit *file : M->getFiles()) {
     if (const SourceFile *source = dyn_cast<SourceFile>(file)) {
       if (source->hasImports()) {
         for (auto import : source->getImports()) {
           if (import.options.contains(ImportFlags::Exported) &&
-              !Imports.contains(import.module.importedModule) &&
               import.module.importedModule->shouldCollectDisplayDecls()) {
-            Imports.insert(import.module.importedModule);
+            auto *TheModule = import.module.importedModule;
+
+            if (import.module.getAccessPath().size() > 0) {
+              if (QualifiedImports.find(TheModule) == QualifiedImports.end()) {
+                QualifiedImports.try_emplace(TheModule);
+              }
+              auto collectDecls = [&](ValueDecl *VD,
+                                      DeclVisibilityKind reason) {
+                if (reason == DeclVisibilityKind::VisibleAtTopLevel)
+                  QualifiedImports[TheModule].insert(VD);
+              };
+              auto consumer = makeDeclConsumer(std::move(collectDecls));
+              TheModule->lookupVisibleDecls(
+                  import.module.getAccessPath(), consumer,
+                  NLKind::UnqualifiedLookup);
+            } else if (!Imports.contains(TheModule)) {
+              Imports.insert(TheModule);
+            }
           }
         }
       }
@@ -956,7 +974,12 @@ SourceFile::getExternalRawLocsForDecl(const Decl *D) const {
 void ModuleDecl::getDisplayDecls(SmallVectorImpl<Decl*> &Results, bool Recursive) const {
   if (Recursive && isParsedModule(this)) {
     SmallPtrSet<ModuleDecl *, 4> Modules;
-    collectParsedExportedImports(this, Modules);
+    llvm::SmallDenseMap<ModuleDecl *, SmallPtrSet<Decl *, 4>, 4> QualifiedImports;
+    collectParsedExportedImports(this, Modules, QualifiedImports);
+    for (const auto &QI : QualifiedImports) {
+      auto &Decls = QI.getSecond();
+      Results.append(Decls.begin(), Decls.end());
+    }
     for (const ModuleDecl *import : Modules) {
       import->getDisplayDecls(Results, Recursive);
     }

lib/SymbolGraphGen/SymbolGraph.cpp

@@ -732,7 +732,7 @@ bool SymbolGraph::isUnconditionallyUnavailableOnAllPlatforms(const Decl *D) cons
 bool SymbolGraph::canIncludeDeclAsNode(const Decl *D) const {
   // 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() && !Walker.isFromExportedImportedModule(D)) {
+  if (D->getModuleContext()->getName() != M.getName() && !Walker.isConsideredExportedImported(D)) {
     return false;
   }
 

lib/SymbolGraphGen/SymbolGraphASTWalker.cpp

@@ -37,23 +37,28 @@ bool areModulesEqual(const ModuleDecl *lhs, const ModuleDecl *rhs, bool ignoreUn
 
 SymbolGraphASTWalker::SymbolGraphASTWalker(ModuleDecl &M,
                                            const SmallPtrSet<ModuleDecl *, 4> ExportedImportedModules,
+                                           const llvm::SmallDenseMap<ModuleDecl *, SmallPtrSet<Decl *, 4>, 4> QualifiedExportedImports,
                                            const SymbolGraphOptions &Options)
   : Options(Options),
     M(M),
     ExportedImportedModules(ExportedImportedModules),
+    QualifiedExportedImports(QualifiedExportedImports),
     MainGraph(*this, M, None, Ctx) {}
 
 /// Get a "sub" symbol graph for the parent module of a type that
 /// the main module `M` is extending.
 SymbolGraph *SymbolGraphASTWalker::getModuleSymbolGraph(const Decl *D) {
   auto *M = D->getModuleContext();
   const auto *DC = D->getDeclContext();
+  const Decl *ExtendedNominal = nullptr;
   while (DC) {
     M = DC->getParentModule();
     if (const auto *NTD = dyn_cast_or_null<NominalTypeDecl>(DC->getAsDecl())) {
       DC = NTD->getDeclContext();
     } else if (const auto *Ext = dyn_cast_or_null<ExtensionDecl>(DC->getAsDecl())) {
       DC = Ext->getExtendedNominal()->getDeclContext();
+      if (!ExtendedNominal)
+        ExtendedNominal = Ext->getExtendedNominal();
     } else {
       DC = nullptr;
     }
@@ -67,8 +72,16 @@ SymbolGraph *SymbolGraphASTWalker::getModuleSymbolGraph(const Decl *D) {
     // should put actual extensions of that module into the main graph
     return &MainGraph;
   }
-  
-  if (isExportedImportedModule(M)) {
+
+  // Check the module and decl separately since the extension could be from a different module
+  // than the decl itself.
+  if (isExportedImportedModule(M) || isQualifiedExportedImport(D)) {
+    return &MainGraph;
+  }
+
+  if (ExtendedNominal && isFromExportedImportedModule(ExtendedNominal)) {
+    return &MainGraph;
+  } else if (!ExtendedNominal && isConsideredExportedImported(D)) {
     return &MainGraph;
   }
 
@@ -230,9 +243,49 @@ bool SymbolGraphASTWalker::walkToDeclPre(Decl *D, CharSourceRange Range) {
   return true;
 }
 
+bool SymbolGraphASTWalker::isConsideredExportedImported(const Decl *D) const {
+  // First check the decl itself to see if it was directly re-exported.
+  if (isFromExportedImportedModule(D))
+    return true;
+
+  const auto *DC = D->getDeclContext();
+
+  // Next, see if the decl is a child symbol of another decl that was re-exported.
+  if (DC) {
+    if (const auto *VD = dyn_cast_or_null<ValueDecl>(DC->getAsDecl())) {
+      if (isFromExportedImportedModule(VD))
+        return true;
+    }
+  }
+
+  // Finally, check to see if this decl is an extension of something else that was re-exported.
+  // FIXME: this considers synthesized members of extensions to be valid
+  const Decl *ExtendedNominal = nullptr;
+  while (DC && !ExtendedNominal) {
+    if (const auto *ED = dyn_cast_or_null<ExtensionDecl>(DC->getAsDecl())) {
+      ExtendedNominal = ED->getExtendedNominal();
+    } else {
+      DC = DC->getParent();
+    }
+  }
+
+  if (ExtendedNominal && isFromExportedImportedModule(ExtendedNominal)) {
+    return true;
+  }
+
+  // If none of the other checks passed, this wasn't from a re-export.
+  return false;
+}
+
 bool SymbolGraphASTWalker::isFromExportedImportedModule(const Decl* D) const {
   auto *M = D->getModuleContext();
-  return isExportedImportedModule(M);
+  return isQualifiedExportedImport(D) || isExportedImportedModule(M);
+}
+
+bool SymbolGraphASTWalker::isQualifiedExportedImport(const Decl *D) const {
+  return llvm::any_of(QualifiedExportedImports, [&D](const auto &QI) {
+    return QI.getSecond().contains(D);
+  });
 }
 
 bool SymbolGraphASTWalker::isExportedImportedModule(const ModuleDecl *M) const {

lib/SymbolGraphGen/SymbolGraphASTWalker.h

@@ -49,6 +49,8 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
 
   const SmallPtrSet<ModuleDecl *, 4> ExportedImportedModules;
 
+  const llvm::SmallDenseMap<ModuleDecl *, SmallPtrSet<Decl *, 4>, 4> QualifiedExportedImports;
+
   /// The symbol graph for the main module of interest.
   SymbolGraph MainGraph;
 
@@ -59,6 +61,7 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
 
   SymbolGraphASTWalker(ModuleDecl &M,
                        const SmallPtrSet<ModuleDecl *, 4> ExportedImportedModules,
+                       const llvm::SmallDenseMap<ModuleDecl *, SmallPtrSet<Decl *, 4>, 4> QualifiedExportedImports,
                        const SymbolGraphOptions &Options);
   virtual ~SymbolGraphASTWalker() {}
 
@@ -93,10 +96,17 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
   virtual bool walkToDeclPre(Decl *D, CharSourceRange Range) override;
 
   // MARK: - Utilities
+
+  /// Returns whether the given declaration was itself imported via an `@_exported import`
+  /// statement, or if it is an extension or child symbol of something else that was.
+  virtual bool isConsideredExportedImported(const Decl *D) const;
 
   /// Returns whether the given declaration comes from an `@_exported import` module.
   virtual bool isFromExportedImportedModule(const Decl *D) const;
 
+  /// Returns whether the given declaration was imported via an `@_exported import <type>` declaration.
+  virtual bool isQualifiedExportedImport(const Decl *D) const;
+
   /// Returns whether the given module is an `@_exported import` module.
   virtual bool isExportedImportedModule(const ModuleDecl *M) const;
 };

lib/SymbolGraphGen/SymbolGraphGen.cpp

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

test/SymbolGraph/Module/ExportedImport.swift

@@ -1,13 +1,17 @@
 // 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/Inputs/ExportedImport/B.swift -module-name B -emit-module -emit-module-path %t/B.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
+@_exported import struct B.StructOne
 
-// CHECK: "precise":"s:1A11SymbolFromAV"
 // CHECK-NOT: InternalSymbolFromA
+// CHECK-NOT: StructTwo
+// CHECK-DAG: "precise":"s:1A11SymbolFromAV"
+// CHECK-DAG: "precise":"s:1B9StructOneV"
 
 // 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.

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

@@ -0,0 +1,3 @@
+public struct StructOne {}
+
+public struct StructTwo {}

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

@@ -0,0 +1,5 @@
+public struct StructOne {
+    public init() {}
+}
+
+public struct StructTwo {}

test/SymbolGraph/Module/QualifiedImportWithExtensions.swift

@@ -0,0 +1,31 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend %S/Inputs/QualifiedImportWithExtensions/A.swift -module-name A -emit-module -emit-module-path %t/A.swiftmodule
+// RUN: %target-swift-frontend %s -module-name QualifiedImportWithExtensions -emit-module -emit-module-path /dev/null -I %t -emit-symbol-graph -emit-symbol-graph-dir %t/
+// RUN: %FileCheck %s --input-file %t/QualifiedImportWithExtensions.symbols.json
+// RUN: %FileCheck %s --input-file %t/QualifiedImportWithExtensions@A.symbols.json --check-prefix EXT
+
+@_exported import struct A.StructOne
+import A
+
+// An extension to a type that's been re-exported with a qualified import should appear in the
+// main symbol graph.
+extension A.StructOne {
+    public struct ExtendedStruct {}
+}
+
+// An extension to a type that's _not_ been re-exported should still appear in the extension symbol
+// graph, even if a different type from that module has been re-exported.
+extension A.StructTwo {
+    public struct InnerStruct {}
+}
+
+// Make sure that the `ExtendedStruct` type is present, and also that `StructOne`'s initializer is
+// present as well. The `InnerStruct` type should not be present since it's extending a type
+// that's not being re-exported.
+
+// CHECK-NOT: InnerStruct
+// CHECK-DAG: ExtendedStruct
+// CHECK-DAG: "precise":"s:1A9StructOneVACycfc"
+
+// EXT: InnerStruct
+// EXT-NOT: ExtendedStruct