Show protocol requirements satisfied in extensions in jump-to-definition

The key issue here was that we were looking at the symbol occurrence that had the `overrideOf` role and were using this symbol occurrence to also check that it’s a child of one of the receiver types. But if the protocol requirement is satisfied in an extension, we have an implicit symbol occurrence at the location where the protocol requirement is stated but the real declaration’s occurrence is inside the extension. The fix here is to just extract the USR from the `overrideOf` relation and then do another index lookup to find the symbol’s definition.

rdar://129412428

Author: ahoppen

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -1717,9 +1717,12 @@ extension SourceKitLSPServer {
         }
       }()
       occurrences += occurrences.flatMap {
-        let overrides = index.occurrences(relatedToUSR: $0.symbol.usr, roles: .overrideOf)
-        // Only contain overrides that are children of one of the receiver types or their subtypes.
-        return overrides.filter { override in
+        let overriddenUsrs = index.occurrences(relatedToUSR: $0.symbol.usr, roles: .overrideOf).map(\.symbol.usr)
+        let overriddenSymbolDefinitions = overriddenUsrs.compactMap {
+          index.primaryDefinitionOrDeclarationOccurrence(ofUSR: $0)
+        }
+        // Only contain overrides that are children of one of the receiver types or their subtypes or extensions.
+        return overriddenSymbolDefinitions.filter { override in
           override.relations.contains(where: {
             guard $0.roles.contains(.childOf) else {
               return false
@@ -2473,13 +2476,13 @@ fileprivate struct DocumentNotificationRequestQueue {
   }
 }
 
-/// Returns the USRs of the subtypes of `usrs` as well as their subtypes, transitively.
+/// Returns the USRs of the subtypes of `usrs` as well as their subtypes and extensions, transitively.
 fileprivate func transitiveSubtypeClosure(ofUsrs usrs: [String], index: CheckedIndex) -> [String] {
   var result: [String] = []
   for usr in usrs {
     result.append(usr)
-    let directSubtypes = index.occurrences(ofUSR: usr, roles: .baseOf).flatMap { occurrence in
-      occurrence.relations.filter { $0.roles.contains(.baseOf) }.map(\.symbol.usr)
+    let directSubtypes = index.occurrences(ofUSR: usr, roles: [.baseOf, .extendedBy]).flatMap { occurrence in
+      occurrence.relations.filter { $0.roles.contains(.baseOf) || $0.roles.contains(.extendedBy) }.map(\.symbol.usr)
     }
     let transitiveSubtypes = transitiveSubtypeClosure(ofUsrs: directSubtypes, index: index)
     result += transitiveSubtypes

Tests/SourceKitLSPTests/DefinitionTests.swift

@@ -601,4 +601,28 @@ class DefinitionTests: XCTestCase {
       ])
     )
   }
+
+  func testJumpToSatisfiedProtocolRequirementInExtension() async throws {
+    let project = try await IndexedSingleSwiftFileTestProject(
+      """
+      protocol TestProtocol {
+        func 1️⃣doThing()
+      }
+
+      struct TestImpl: TestProtocol {}
+      extension TestImpl {
+        func 2️⃣doThing() { }
+      }
+      """
+    )
+
+    let response = try await project.testClient.send(
+      DefinitionRequest(textDocument: TextDocumentIdentifier(project.fileURI), position: project.positions["1️⃣"])
+    )
+    guard case .locations(let locations) = response else {
+      XCTFail("Expected locations response")
+      return
+    }
+    XCTAssertEqual(locations, [Location(uri: project.fileURI, range: Range(project.positions["2️⃣"]))])
+  }
 }