Merge pull request #30857 from bitjammer/acgarland/rdar-60796811-rootmost-extension-sgf

[SymbolGraph] Put extending declarations in rootmost module

Author: bitjammer

lib/SymbolGraphGen/SymbolGraphASTWalker.cpp

@@ -29,7 +29,20 @@ SymbolGraphASTWalker::SymbolGraphASTWalker(ModuleDecl &M,
 
 /// Get a "sub" symbol graph for the parent module of a type that
 /// the main module `M` is extending.
-SymbolGraph *SymbolGraphASTWalker::getModuleSymbolGraph(ModuleDecl *M) {
+SymbolGraph *SymbolGraphASTWalker::getModuleSymbolGraph(const Decl *D) {
+  auto *M = D->getModuleContext();
+  const auto *DC = D->getDeclContext();
+  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();
+    } else {
+      DC = nullptr;
+    }
+  }
+
   if (this->M.getNameStr().equals(M->getNameStr())) {
     return &MainGraph;
   }
@@ -88,7 +101,7 @@ bool SymbolGraphASTWalker::walkToDeclPre(Decl *D, CharSourceRange Range) {
       return true;
   }
 
-  auto SG = getModuleSymbolGraph(D->getModuleContext());
+  auto SG = getModuleSymbolGraph(D);
 
   // If this is an extension, let's check that it implies some new conformances,
   // potentially with generic requirements.
@@ -107,7 +120,7 @@ bool SymbolGraphASTWalker::walkToDeclPre(Decl *D, CharSourceRange Range) {
     // grab them for some new conformsTo relationships.
     if (!Extension->getInherited().empty()) {
       auto ExtendedSG =
-          getModuleSymbolGraph(ExtendedNominal->getModuleContext());
+          getModuleSymbolGraph(ExtendedNominal);
 
       // The symbol graph to use to record these relationships.
       SmallVector<const ProtocolDecl *, 4> Protocols;
@@ -175,7 +188,7 @@ bool SymbolGraphASTWalker::walkToDeclPre(Decl *D, CharSourceRange Range) {
       = dyn_cast_or_null<ExtensionDecl>(VD->getDeclContext())) {
     if (const auto *ExtendedNominal = Extension->getExtendedNominal()) {
       auto ExtendedModule = ExtendedNominal->getModuleContext();
-      auto ExtendedSG = getModuleSymbolGraph(ExtendedModule);
+      auto ExtendedSG = getModuleSymbolGraph(ExtendedNominal);
       if (ExtendedModule != &M) {
         ExtendedSG->recordNode(Symbol(ExtendedSG, VD, nullptr));
         return true;

lib/SymbolGraphGen/SymbolGraphASTWalker.h

@@ -60,8 +60,29 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
 
   // MARK: - Utilities
 
-  /// Get a "sub" symbol graph for the parent module of a type that the main module `M` is extending.
-  SymbolGraph *getModuleSymbolGraph(ModuleDecl *M);
+  /// Get a "sub" symbol graph for the appropriate module concerning a declaration.
+  ///
+  /// This will get the "rootmost" module responsible for a declaration's
+  /// documentation. For example:
+  ///
+  /// Module A:
+  ///
+  /// ```swift
+  /// public struct AStruct {}
+  /// ```
+  ///
+  /// Module B:
+  ///
+  /// ```swift
+  /// import A
+  /// extension AStruct {
+  ///   public struct BStruct {}
+  /// }
+  ///
+  /// `BStruct` will go in module A's extension symbol graph, because `BStruct`
+  /// is a member of `AStruct`, and module A owns `AStruct`, and so on for
+  /// further nestings.
+  SymbolGraph *getModuleSymbolGraph(const Decl *D);
 
   // MARK: - SourceEntityWalker
 

test/SymbolGraph/Module/BasicExtension.swift

@@ -0,0 +1,26 @@
+// 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/[email protected]
+
+extension String {
+  /// Return something.
+  public var something: String {
+    return "something"
+  }
+}
+
+// CHECK: module
+// CHECK-NEXT: "name": "BasicExtension"
+
+// CHECK: "precise": "s:SS14BasicExtensionE9somethingSSvp"
+
+// CHECK: "kind": "memberOf"
+// CHECK-NEXT: "source": "s:SS14BasicExtensionE9somethingSSvp"
+// CHECK-NEXT: "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"

test/SymbolGraph/Module/Extension.swift

@@ -0,0 +1,14 @@
+public protocol P {
+  associatedtype Thing
+  func foo() -> Thing
+}
+
+public struct AStruct<Thing>: P {
+  public var thing: Thing
+  public init(thing: Thing) {
+    self.thing = thing
+  }
+  public func foo() -> Thing {
+    return thing
+  }
+}

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

@@ -0,0 +1,13 @@
+import A
+
+extension AStruct {
+  public struct BStruct {
+    public func foo() -> Double {
+      return 1.0
+    }
+  }
+}
+
+extension AStruct.BStruct: P where Thing == Double {
+  public func bar() {}
+}

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

@@ -0,0 +1,18 @@
+import A
+import B
+
+extension AStruct.BStruct {
+  public struct CStruct: P {
+    public func foo() -> UInt8 {
+      return 0
+    }
+  }
+}
+
+extension AStruct.BStruct {
+  public func baz() {}
+}
+
+extension AStruct.BStruct.CStruct where Thing: Equatable {
+  public func baz() {}
+}

test/SymbolGraph/Module/Inputs/NestedExtensions/C.swift

@@ -0,0 +1,44 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %S/Inputs/NestedExtensions/A.swift -I %t -module-name A -emit-module -emit-module-path %t/
+// RUN: %target-build-swift %S/Inputs/NestedExtensions/B.swift -I %t -module-name B -emit-module -emit-module-path %t/
+// RUN: %target-build-swift %s -module-name NestedExtensions -emit-module -I %t -emit-module-path %t/
+
+// RUN: %target-swift-symbolgraph-extract -module-name A -I %t -pretty-print -output-dir %t
+// RUN: %target-swift-symbolgraph-extract -module-name B -I %t -pretty-print -output-dir %t
+// RUN: %target-swift-symbolgraph-extract -module-name NestedExtensions -I %t -pretty-print -output-dir %t
+
+// RUN: %FileCheck %s --input-file %t/B.symbols.json --check-prefix=MODULEB
+// RUN: %FileCheck %s --input-file %t/[email protected] --check-prefix=MODULEBATA
+
+// RUN: %FileCheck %s --input-file %t/[email protected] --check-prefix=NESTEDATA
+
+// RUN: %FileCheck %s --input-file %t/NestedExtensions.symbols.json --check-prefix=NESTED
+
+import A
+import B
+
+extension AStruct.BStruct {
+  public struct CStruct: P {
+    public func foo() -> UInt8 {
+      return 0
+    }
+  }
+}
+
+extension AStruct.BStruct {
+  public func baz() {}
+}
+
+extension AStruct.BStruct.CStruct where Thing: Equatable {
+  public func baz() {}
+}
+
+// BStruct belongs to AStruct and so should only ever appear in B@A extension symbol graph files.
+// MODULEB-NOT: BStruct
+// MODULEBATA: "precise": "s:1A7AStructV1BE7BStructV"
+
+// CStruct belongs to BStruct, and BStruct belongs to AStruct, so should only appear in NestedExtension@A.
+// NESTED-NOT: BStruct
+// NESTED-NOT: CStruct
+// NESTEDATB-NOT: BStruct
+// NESTEDATA: "precise": "s:1A7AStructV1BE7BStructV16NestedExtensionsE7CStructV"