Show overriden functions when performing jump-to-definition on a dynamic call

Fixes #809
rdar://114864256

Author: ahoppen

Sources/LanguageServerProtocol/Requests/SymbolInfoRequest.swift

@@ -81,17 +81,23 @@ public struct SymbolDetails: ResponseType, Hashable {
   /// The kind of the symbol
   public var kind: SymbolKind?
 
+  /// Whether the symbol is a dynamic call for which it isn't known which method will be invoked at runtime. This is
+  /// the case for protocol methods and class functions.
+  public var isDynamic: Bool
+
   public init(
     name: String?,
-    containerName: String? = nil,
+    containerName: String?,
     usr: String?,
-    bestLocalDeclaration: Location? = nil,
-    kind: SymbolKind? = nil
+    bestLocalDeclaration: Location?,
+    kind: SymbolKind?,
+    isDynamic: Bool
   ) {
     self.name = name
     self.containerName = containerName
     self.usr = usr
     self.bestLocalDeclaration = bestLocalDeclaration
     self.kind = kind
+    self.isDynamic = isDynamic
   }
 }

Sources/SourceKitD/sourcekitd_uids.swift

@@ -44,6 +44,7 @@ public struct sourcekitd_keys {
   public let groupname: sourcekitd_uid_t
   public let id: sourcekitd_uid_t
   public let is_system: sourcekitd_uid_t
+  public let isDynamic: sourcekitd_uid_t
   public let kind: sourcekitd_uid_t
   public let length: sourcekitd_uid_t
   public let line: sourcekitd_uid_t
@@ -122,6 +123,7 @@ public struct sourcekitd_keys {
     groupname = api.uid_get_from_cstr("key.groupname")!
     id = api.uid_get_from_cstr("key.id")!
     is_system = api.uid_get_from_cstr("key.is_system")!
+    isDynamic = api.uid_get_from_cstr("key.is_dynamic")!
     kind = api.uid_get_from_cstr("key.kind")!
     length = api.uid_get_from_cstr("key.length")!
     line = api.uid_get_from_cstr("key.line")!

Sources/SourceKitLSP/SourceKitServer.swift

@@ -1667,7 +1667,7 @@ extension SourceKitServer {
     symbols: [SymbolDetails],
     index: IndexStoreDB?,
     useLocalFallback: Bool = false,
-    extractOccurrences: (String, IndexStoreDB) -> [SymbolOccurrence]
+    extractOccurrences: (SymbolDetails, IndexStoreDB) -> [SymbolOccurrence]
   ) -> [(occurrence: SymbolOccurrence?, location: Location)] {
     guard let symbol = symbols.first else {
       return []
@@ -1680,11 +1680,11 @@ extension SourceKitServer {
       fallback = []
     }
 
-    guard let usr = symbol.usr, let index = index else {
+    guard let index = index else {
       return fallback
     }
 
-    let occurs = extractOccurrences(usr, index)
+    let occurs = extractOccurrences(symbol, index)
     let resolved = occurs.compactMap { occur in
       indexToLSPLocation(occur.location).map {
         (occurrence: occur, location: $0)
@@ -1725,12 +1725,18 @@ extension SourceKitServer {
     }
 
     let resolved = self.extractIndexedOccurrences(symbols: symbols, index: index, useLocalFallback: true) {
-      (usr, index) in
+      (symbolDetails, index) in
+      guard let usr = symbolDetails.usr else { return [] }
       logger.info("performing indexed jump-to-def with usr \(usr)")
       var occurs = index.occurrences(ofUSR: usr, roles: [.definition])
       if occurs.isEmpty {
         occurs = index.occurrences(ofUSR: usr, roles: [.declaration])
       }
+      if symbolDetails.isDynamic {
+        occurs += occurs.flatMap {
+          index.occurrences(relatedToUSR: $0.symbol.usr, roles: .overrideOf)
+        }
+      }
       return occurs
     }
 
@@ -1783,7 +1789,8 @@ extension SourceKitServer {
       )
     )
     let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index
-    let extractedResult = self.extractIndexedOccurrences(symbols: symbols, index: index) { (usr, index) in
+    let extractedResult = self.extractIndexedOccurrences(symbols: symbols, index: index) { (symbolDetails, index) in
+      guard let usr = symbolDetails.usr else { return [] }
       var occurs = index.occurrences(ofUSR: usr, roles: .baseOf)
       if occurs.isEmpty {
         occurs = index.occurrences(relatedToUSR: usr, roles: .overrideOf)
@@ -1806,7 +1813,8 @@ extension SourceKitServer {
       )
     )
     let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index
-    let extractedResult = self.extractIndexedOccurrences(symbols: symbols, index: index) { (usr, index) in
+    let extractedResult = self.extractIndexedOccurrences(symbols: symbols, index: index) { (symbolDetails, index) in
+      guard let usr = symbolDetails.usr else { return [] }
       logger.info("performing indexed jump-to-def with usr \(usr)")
       var roles: SymbolRole = [.reference]
       if req.context.includeDeclaration {
@@ -1852,8 +1860,9 @@ extension SourceKitServer {
     )
     let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index
     // For call hierarchy preparation we only locate the definition
-    let extractedResult = self.extractIndexedOccurrences(symbols: symbols, index: index) { (usr, index) in
-      index.occurrences(ofUSR: usr, roles: [.definition, .declaration])
+    let extractedResult = self.extractIndexedOccurrences(symbols: symbols, index: index) { (symbolDetails, index) in
+      guard let usr = symbolDetails.usr else { return [] }
+      return index.occurrences(ofUSR: usr, roles: [.definition, .declaration])
     }
     return extractedResult.compactMap { info -> CallHierarchyItem? in
       guard let occurrence = info.occurrence else {
@@ -2009,8 +2018,9 @@ extension SourceKitServer {
     guard let index = await self.workspaceForDocument(uri: req.textDocument.uri)?.index else {
       return []
     }
-    let extractedResult = self.extractIndexedOccurrences(symbols: symbols, index: index) { (usr, index) in
-      index.occurrences(ofUSR: usr, roles: [.definition, .declaration])
+    let extractedResult = self.extractIndexedOccurrences(symbols: symbols, index: index) { (symbolDetails, index) in
+      guard let usr = symbolDetails.usr else { return [] }
+      return index.occurrences(ofUSR: usr, roles: [.definition, .declaration])
     }
     return extractedResult.compactMap { info -> TypeHierarchyItem? in
       guard let occurrence = info.occurrence else {

Sources/SourceKitLSP/Swift/CursorInfo.swift

@@ -131,7 +131,8 @@ extension SwiftLanguageServer {
         containerName: nil,
         usr: dict[keys.usr],
         bestLocalDeclaration: location,
-        kind: kind.asSymbolKind(self.sourcekitd.values)
+        kind: kind.asSymbolKind(self.sourcekitd.values),
+        isDynamic: dict[keys.isDynamic] ?? false
       ),
       annotatedDeclaration: dict[keys.annotated_decl],
       documentationXML: dict[keys.doc_full_as_xml],

Tests/SourceKitLSPTests/DefinitionTests.swift

@@ -36,4 +36,37 @@ class DefinitionTests: XCTestCase {
     }
     XCTAssertEqual(locations, [Location(uri: uri, range: Range(positions["1️⃣"]))])
   }
+
+  func testJumpToDefinitionIncludesOverrides() async throws {
+    let ws = try await IndexedSingleSwiftFileWorkspace(
+      """
+      protocol TestProtocol {
+        func 1️⃣doThing()
+      }
+
+      struct TestImpl: TestProtocol { 
+        func 2️⃣doThing() { }
+      }
+
+      func anyTestProtocol(value: any TestProtocol) {
+        value.3️⃣doThing()
+      }
+      """
+    )
+
+    let response = try await ws.testClient.send(
+      DefinitionRequest(textDocument: TextDocumentIdentifier(ws.fileURI), position: ws.positions["3️⃣"])
+    )
+    guard case .locations(let locations) = response else {
+      XCTFail("Expected locations response")
+      return
+    }
+    XCTAssertEqual(
+      locations,
+      [
+        Location(uri: ws.fileURI, range: Range(ws.positions["1️⃣"])),
+        Location(uri: ws.fileURI, range: Range(ws.positions["2️⃣"])),
+      ]
+    )
+  }
 }