connect memberOf relationship of default implementation targets to extension block symbol if extension block symbol format is used

Author: theMomax

lib/SymbolGraphGen/SymbolGraph.cpp

@@ -413,11 +413,20 @@ 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.
+          // protocol or the respective extension block.
           if (!Walker.isOurModule(MemberVD->getModuleContext()) && VD->getDeclContext()) {
-            if (auto *ExP = VD->getDeclContext()->getSelfNominalTypeDecl()) {
+            if (const auto *Extension
+                = dyn_cast_or_null<ExtensionDecl>(VD->getDeclContext())) {
+              if (this->Walker.shouldBeRecordedAsExtension(Extension)) {
+                recordEdge(Symbol(this, VD, nullptr),
+                           Symbol(this, Extension, nullptr),
+                           RelationshipKind::MemberOf());
+                continue;
+              }
+            }
+            if (auto *ExtendedProtocol = VD->getDeclContext()->getSelfNominalTypeDecl()) {
               recordEdge(Symbol(this, VD, nullptr),
-                         Symbol(this, ExP, nullptr),
+                         Symbol(this, ExtendedProtocol, nullptr),
                          RelationshipKind::MemberOf());
             }
           }

test/SymbolGraph/Relationships/DefaultImplementationOf/External.swift

@@ -0,0 +1,45 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %S/Inputs/RemoteP.swift -module-name RemoteP -emit-module -emit-module-path %t/
+// RUN: %target-build-swift %s -module-name External -emit-module -emit-module-path %t/ -I %t
+// RUN: %target-swift-symbolgraph-extract -module-name External -I %t -pretty-print -output-dir %t
+// RUN: %FileCheck %s --input-file %t/[email protected]
+// RUN: %FileCheck %s --input-file %t/[email protected] --check-prefix MEMBER
+
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %S/Inputs/RemoteP.swift -module-name RemoteP -emit-module -emit-module-path %t/
+// RUN: %target-build-swift %s -module-name External -emit-module -emit-module-path %t/ -I %t
+// RUN: %target-swift-symbolgraph-extract -module-name External -I %t -pretty-print -output-dir %t -emit-extension-block-symbols
+// RUN: %FileCheck %s --input-file %t/[email protected]
+// RUN: %FileCheck %s --input-file %t/[email protected] --check-prefix MEMBEREBS
+// RUN: %FileCheck %s --input-file %t/[email protected] --check-prefix EXTENSIONTOEBS
+
+import RemoteP
+
+public extension RemoteP {
+    func someFunc() {}
+}
+
+// Default implementations that are for protocols in a different module should have a `memberOf`
+// relation linking them to a local symbol. If the default implementation is not defined on a local
+// protocol, the extension block symbol defining the default implementation should be the target
+// of the memberOf relationship.
+
+// CHECK:           "kind": "defaultImplementationOf"
+// CHECK-NEXT:      "source": "s:7RemotePAAP8ExternalE8someFuncyyF"
+// CHECK-NEXT:      "target": "s:7RemotePAAP8someFuncyyF"
+// CHECK-NEXT:      "targetFallback": "RemoteP.RemoteP.someFunc()"
+
+// MEMBER:           "kind": "memberOf"
+// MEMBER-NEXT:      "source": "s:7RemotePAAP8ExternalE8someFuncyyF"
+// MEMBER-NEXT:      "target": "s:7RemotePAAP"
+// MEMBER-NEXT:      "targetFallback": "RemoteP.RemoteP"
+
+// MEMBEREBS:        "kind": "memberOf"
+// MEMBEREBS-NEXT:   "source": "s:7RemotePAAP8ExternalE8someFuncyyF"
+// MEMBEREBS-NEXT:   "target": "s:e:s:7RemotePAAP8ExternalE8someFuncyyF"
+// MEMBEREBS-NEXT:   "targetFallback": "RemoteP.RemoteP"
+
+// EXTENSIONTOEBS:   "kind": "extensionTo"
+// EXTENSIONTO-NEXT: "source": "s:7RemotePAAP8ExternalE8someFuncyyF"
+// EXTENSIONTO-NEXT: "target": "s:7RemotePAAP"
+// EXTENSIONTO-NEXT: "targetFallback": "RemoteP.RemoteP"

test/SymbolGraph/Relationships/DefaultImplementationOf/Remote.swift

@@ -5,6 +5,13 @@
 // RUN: %FileCheck %s --input-file %t/Remote.symbols.json
 // RUN: %FileCheck %s --input-file %t/Remote.symbols.json --check-prefix MEMBER
 
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %S/Inputs/RemoteP.swift -module-name RemoteP -emit-module -emit-module-path %t/
+// RUN: %target-build-swift %s -module-name Remote -emit-module -emit-module-path %t/ -I %t
+// RUN: %target-swift-symbolgraph-extract -module-name Remote -I %t -pretty-print -output-dir %t -emit-extension-block-symbols
+// RUN: %FileCheck %s --input-file %t/Remote.symbols.json
+// RUN: %FileCheck %s --input-file %t/Remote.symbols.json --check-prefix MEMBER
+
 import RemoteP
 
 public protocol LocalP: RemoteP {}
@@ -13,8 +20,9 @@ public extension LocalP {
     func someFunc() {}
 }
 
-// default implementations that are for protocols in a different module should have a `memberOf`
-// relation linking them to a local symbol, if one exists
+// Default implementations that are for protocols in a different module should have a `memberOf`
+// relation linking them to a local symbol. If the default implementation is defined on a local
+// protocol, this local protocol is the target of the memberOf relationship.
 
 // CHECK:           "kind": "defaultImplementationOf"
 // CHECK-NEXT:      "source": "s:6Remote6LocalPPAAE8someFuncyyF"