[SymbolGraphGen] Refactor export-import logic (#61049)

rdar://98808363

lib/SymbolGraphGen/SymbolGraph.cpp

@@ -322,7 +322,9 @@ void SymbolGraph::recordConformanceSynthesizedMemberRelationships(Symbol S) {
 
       // We are only interested in synthesized members that come from an
       // extension that we defined in our module.
-      if (Info.EnablingExt && Info.EnablingExt->getModuleContext() != &M) {
+      if (Info.EnablingExt) {
+        const auto *ExtM = Info.EnablingExt->getModuleContext();
+        if (!Walker.isOurModule(ExtM))
           continue;
       }
 
@@ -399,7 +401,7 @@ void SymbolGraph::recordDefaultImplementationRelationships(Symbol S) {
           // If P is from a different module, and it's being added to a type
           // from the current module, add a `memberOf` relation to the extended
           // protocol.
-          if (MemberVD->getModuleContext()->getNameStr() != M.getNameStr() && VD->getDeclContext()) {
+          if (!Walker.isOurModule(MemberVD->getModuleContext()) && VD->getDeclContext()) {
             if (auto *ExP = VD->getDeclContext()->getSelfNominalTypeDecl()) {
               recordEdge(Symbol(this, VD, nullptr),
                          Symbol(this, ExP, nullptr),

lib/SymbolGraphGen/SymbolGraphASTWalker.cpp

@@ -25,12 +25,10 @@ namespace {
 
 /// Compare the two \c ModuleDecl instances to see whether they are the same.
 ///
-/// Pass \c true to the \c ignoreUnderlying argument to consider two modules the same even if
-/// one is a Swift module and the other a non-Swift module. This allows a Swift module and its
-/// underlying Clang module to compare as equal.
-bool areModulesEqual(const ModuleDecl *lhs, const ModuleDecl *rhs, bool ignoreUnderlying = false) {
-  return lhs->getNameStr() == rhs->getNameStr()
-    && (ignoreUnderlying || lhs->isNonSwiftModule() == rhs->isNonSwiftModule());
+/// This does a by-name comparison to consider a module's underlying Clang module to be equivalent
+/// to the wrapping module of the same name.
+bool areModulesEqual(const ModuleDecl *lhs, const ModuleDecl *rhs) {
+  return lhs->getNameStr() == rhs->getNameStr();
 }
 
 } // anonymous namespace
@@ -50,11 +48,13 @@ SymbolGraphASTWalker::SymbolGraphASTWalker(ModuleDecl &M,
 SymbolGraph *SymbolGraphASTWalker::getModuleSymbolGraph(const Decl *D) {
   auto *M = D->getModuleContext();
   const auto *DC = D->getDeclContext();
+  SmallVector<const NominalTypeDecl *, 2> ParentTypes = {};
   const Decl *ExtendedNominal = nullptr;
   while (DC) {
     M = DC->getParentModule();
     if (const auto *NTD = dyn_cast_or_null<NominalTypeDecl>(DC->getAsDecl())) {
       DC = NTD->getDeclContext();
+      ParentTypes.push_back(NTD);
     } else if (const auto *Ext = dyn_cast_or_null<ExtensionDecl>(DC->getAsDecl())) {
       DC = Ext->getExtendedNominal()->getDeclContext();
       if (!ExtendedNominal)
@@ -64,10 +64,10 @@ SymbolGraph *SymbolGraphASTWalker::getModuleSymbolGraph(const Decl *D) {
     }
   }
 
-  if (areModulesEqual(&this->M, M, true)) {
+  if (areModulesEqual(&this->M, M)) {
     return &MainGraph;
   } else if (MainGraph.DeclaringModule.hasValue() &&
-             areModulesEqual(MainGraph.DeclaringModule.getValue(), M, true)) {
+             areModulesEqual(MainGraph.DeclaringModule.getValue(), M)) {
     // Cross-import overlay modules already appear as "extensions" of their declaring module; we
     // should put actual extensions of that module into the main graph
     return &MainGraph;
@@ -79,9 +79,8 @@ SymbolGraph *SymbolGraphASTWalker::getModuleSymbolGraph(const Decl *D) {
     return &MainGraph;
   }
 
-  if (ExtendedNominal && isFromExportedImportedModule(ExtendedNominal)) {
-    return &MainGraph;
-  } else if (!ExtendedNominal && isConsideredExportedImported(D)) {
+  // If this type is the child of a type which was re-exported in a qualified export, use the main graph.
+  if (llvm::any_of(ParentTypes, [&](const NominalTypeDecl *NTD){ return isQualifiedExportedImport(NTD); })) {
     return &MainGraph;
   }
   
@@ -230,7 +229,7 @@ bool SymbolGraphASTWalker::walkToDeclPre(Decl *D, CharSourceRange Range) {
     if (const auto *ExtendedNominal = Extension->getExtendedNominal()) {
       auto ExtendedModule = ExtendedNominal->getModuleContext();
       auto ExtendedSG = getModuleSymbolGraph(ExtendedNominal);
-      if (ExtendedModule != &M) {
+      if (!isOurModule(ExtendedModule)) {
         ExtendedSG->recordNode(Symbol(ExtendedSG, VD, nullptr));
         return true;
       }
@@ -257,22 +256,10 @@ bool SymbolGraphASTWalker::walkToDeclPre(Decl *D, CharSourceRange Range) {
 }
 
 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.
+  // Check to see if this decl is an extension of something else that was re-exported.
+  // Do this first in case there's a chain of extensions that leads somewhere that's not a re-export.
   // FIXME: this considers synthesized members of extensions to be valid
+  const auto *DC = D->getDeclContext();
   const Decl *ExtendedNominal = nullptr;
   while (DC && !ExtendedNominal) {
     if (const auto *ED = dyn_cast_or_null<ExtensionDecl>(DC->getAsDecl())) {
@@ -282,10 +269,23 @@ bool SymbolGraphASTWalker::isConsideredExportedImported(const Decl *D) const {
     }
   }
 
-  if (ExtendedNominal && isFromExportedImportedModule(ExtendedNominal)) {
+  if (ExtendedNominal && isConsideredExportedImported(ExtendedNominal)) {
     return true;
   }
 
+  // Check to see if the decl is a child symbol of another decl that was re-exported.
+  DC = D->getDeclContext();
+  if (DC) {
+    if (const auto *VD = dyn_cast_or_null<ValueDecl>(DC->getAsDecl())) {
+      if (isConsideredExportedImported(VD))
+        return true;
+    }
+  }
+
+  // Check the decl itself to see if it was directly re-exported.
+  if (isFromExportedImportedModule(D) || isQualifiedExportedImport(D))
+    return true;
+
   // If none of the other checks passed, this wasn't from a re-export.
   return false;
 }
@@ -306,3 +306,7 @@ bool SymbolGraphASTWalker::isExportedImportedModule(const ModuleDecl *M) const {
     return areModulesEqual(M, MD->getModuleContext());
   });
 }
+
+bool SymbolGraphASTWalker::isOurModule(const ModuleDecl *M) const {
+  return areModulesEqual(M, &this->M) || isExportedImportedModule(M);
+}

lib/SymbolGraphGen/SymbolGraphASTWalker.h

@@ -47,6 +47,7 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
   /// The module that this symbol graph will represent.
   const ModuleDecl &M;
 
+  // FIXME: these should be tracked per-graph, rather than at the top level
   const SmallPtrSet<ModuleDecl *, 4> ExportedImportedModules;
 
   const llvm::SmallDenseMap<ModuleDecl *, SmallPtrSet<Decl *, 4>, 4> QualifiedExportedImports;
@@ -109,6 +110,9 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
 
   /// Returns whether the given module is an `@_exported import` module.
   virtual bool isExportedImportedModule(const ModuleDecl *M) const;
+
+  /// Returns whether the given module is the main module, or is an `@_exported import` module.
+  virtual bool isOurModule(const ModuleDecl *M) const;
 };
 
 } // end namespace symbolgraphgen

test/SymbolGraph/ClangImporter/ExportedImport.swift

@@ -0,0 +1,12 @@
+// RUN: %empty-directory(%t)
+// RUN: cp -r %S/Inputs/ExportedImport/ObjcProperty.framework %t
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -enable-objc-interop -emit-module -o %t/ObjcProperty.framework/Modules/ObjcProperty.swiftmodule/%target-swiftmodule-name -import-underlying-module -F %t -module-name ObjcProperty %S/Inputs/ExportedImport/A.swift
+// RUN: %target-swift-frontend(mock-sdk: %clang-importer-sdk) -enable-objc-interop -emit-module -o %t/ExportedImport.swiftmodule -F %t -module-name ExportedImport %s -emit-symbol-graph -emit-symbol-graph-dir %t
+// RUN: %FileCheck %s --input-file %t/ExportedImport.symbols.json
+
+// REQUIRES: objc_interop
+
+// CHECK-DAG: "precise":"s:So11ClangStructa12ObjcPropertyE05InnerB0V"
+// CHECK-DAG: "precise":"s:12ObjcProperty12SomeProtocolPAAE8someFuncyyF::SYNTHESIZED::s:So11ClangStructa12ObjcPropertyE05InnerB0V06NestedB0V",
+
+@_exported import ObjcProperty

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

@@ -0,0 +1,15 @@
+extension SwiftStruct {
+    public struct InnerStruct {}
+}
+
+extension SwiftStruct.InnerStruct {
+    public struct NestedStruct {}
+}
+
+public protocol SomeProtocol {}
+
+extension SomeProtocol {
+    public func someFunc() {}
+}
+
+extension SwiftStruct.InnerStruct.NestedStruct: SomeProtocol {}

test/SymbolGraph/ClangImporter/Inputs/ExportedImport/ObjcProperty.framework/Headers/ObjcProperty.apinotes

@@ -0,0 +1,5 @@
+---
+Name: ObjcProperty
+Typedefs:
+- Name: ClangStruct
+  SwiftName: SwiftStruct

test/SymbolGraph/ClangImporter/Inputs/ExportedImport/ObjcProperty.framework/Headers/ObjcProperty.h

@@ -0,0 +1,3 @@
+typedef struct {
+    unsigned filler;
+} ClangStruct;

test/SymbolGraph/ClangImporter/Inputs/ExportedImport/ObjcProperty.framework/module.map

@@ -0,0 +1,4 @@
+framework module ObjcProperty {
+  header "ObjcProperty.h"
+  export *
+}

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

@@ -3,3 +3,7 @@ import A
 public extension SomeStruct {
     struct InnerStruct: Equatable {}
 }
+
+public extension SomeStruct.InnerStruct {
+    struct NestedStruct: Equatable {}
+}

test/SymbolGraph/Module/ThirdOrder.swift

@@ -12,4 +12,7 @@
 @_exported import B
 
 // BASE-NOT: "s:SQsE2neoiySbx_xtFZ::SYNTHESIZED::s:1A10SomeStructV1BE05InnerB0V"
-// EXT: "s:SQsE2neoiySbx_xtFZ::SYNTHESIZED::s:1A10SomeStructV1BE05InnerB0V"
+// EXT-DAG: "s:SQsE2neoiySbx_xtFZ::SYNTHESIZED::s:1A10SomeStructV1BE05InnerB0V"
+
+// BASE-NOT: "s:1A10SomeStructV1BE05InnerB0V06NestedB0V"
+// EXT-DAG: "s:1A10SomeStructV1BE05InnerB0V06NestedB0V"