[SymbolGraph] Put extending declarations in rootmost module

When extending another module's type in your module, serialize declarations in
the extension into the other module's "extension" symbol graph file, including
relationships. This mechanic should continue up to the rootmost module. For
example:

A.AStruct <- B.BStruct < C.CStruct

Both BStruct and CStruct should go in `@A` symbol graph files because AStruct
owns BStruct and by extension owns CStruct. This is reflected in
documentation curation in some form already.

rdar://60796811

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/BasicExtension@Swift.symbols.json
+
+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/B@A.symbols.json --check-prefix=MODULEBATA
+
+// RUN: %FileCheck %s --input-file %t/NestedExtensions@A.symbols.json --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"