Filter overrides in DefinitionRequest by receiver types

Author: ahoppen

Sources/LanguageServerProtocol/Requests/SymbolInfoRequest.swift

@@ -87,19 +87,46 @@ public struct SymbolDetails: ResponseType, Hashable {
   /// Optional because `clangd` does not return whether a symbol is dynamic.
   public var isDynamic: Bool?
 
+  /// If the symbol is dynamic, the USRs of the types that might be called.
+  ///
+  /// This is relevant in the following cases
+  /// ```swift
+  /// class A {
+  ///   func doThing() {}
+  /// }
+  /// class B: A {}
+  /// class C: B {
+  ///   override func doThing() {}
+  /// }
+  /// class D: A {
+  ///   override func doThing() {}
+  /// }
+  /// func test(value: B) {
+  ///   value.doThing()
+  /// }
+  /// ```
+  ///
+  /// The USR of the called function in `value.doThing` is `A.doThing` (or its
+  /// mangled form) but it can never call `D.doThing`. In this case, the
+  /// receiver USR would be `B`, indicating that only overrides of subtypes in
+  /// `B` may be called dynamically.
+  public var receiverUsrs: [String]?
+
   public init(
     name: String?,
     containerName: String?,
     usr: String?,
     bestLocalDeclaration: Location?,
     kind: SymbolKind?,
-    isDynamic: Bool
+    isDynamic: Bool?,
+    receiverUsrs: [String]?
   ) {
     self.name = name
     self.containerName = containerName
     self.usr = usr
     self.bestLocalDeclaration = bestLocalDeclaration
     self.kind = kind
     self.isDynamic = isDynamic
+    self.receiverUsrs = receiverUsrs
   }
 }

Sources/SourceKitD/SKDResponseArray.swift

@@ -55,6 +55,27 @@ public final class SKDResponseArray {
     return true
   }
 
+  public func map<T>(_ transform: (SKDResponseDictionary) -> T) -> [T] {
+    var result: [T] = []
+    result.reserveCapacity(self.count)
+    self.forEach { _, element in
+      result.append(transform(element))
+      return true
+    }
+    return result
+  }
+
+  public func compactMap<T>(_ transform: (SKDResponseDictionary) -> T?) -> [T] {
+    var result: [T] = []
+    self.forEach { _, element in
+      if let transformed = transform(element) {
+        result.append(transformed)
+      }
+      return true
+    }
+    return result
+  }
+
   /// Attempt to access the item at `index` as a string.
   public subscript(index: Int) -> String? {
     if let cstr = sourcekitd.api.variant_array_get_string(array, index) {

Sources/SourceKitD/sourcekitd_uids.swift

@@ -57,6 +57,7 @@ public struct sourcekitd_keys {
   public let num_bytes_to_erase: sourcekitd_uid_t
   public let offset: sourcekitd_uid_t
   public let ranges: sourcekitd_uid_t
+  public let receivers: sourcekitd_uid_t
   public let refactor_actions: sourcekitd_uid_t
   public let request: sourcekitd_uid_t
   public let results: sourcekitd_uid_t
@@ -136,6 +137,7 @@ public struct sourcekitd_keys {
     num_bytes_to_erase = api.uid_get_from_cstr("key.num_bytes_to_erase")!
     offset = api.uid_get_from_cstr("key.offset")!
     ranges = api.uid_get_from_cstr("key.ranges")!
+    receivers = api.uid_get_from_cstr("key.receivers")!
     refactor_actions = api.uid_get_from_cstr("key.refactor_actions")!
     request = api.uid_get_from_cstr("key.request")!
     results = api.uid_get_from_cstr("key.results")!

Sources/SourceKitLSP/SourceKitServer.swift

@@ -1709,8 +1709,18 @@ extension SourceKitServer {
       occurances = index.occurrences(ofUSR: usr, roles: [.declaration])
     }
     if symbol.isDynamic ?? false {
+      lazy var transitiveReceiverUsrs: [String] = transitiveSubtypeClosure(
+        ofUsrs: symbol.receiverUsrs ?? [],
+        index: index
+      )
       occurances += occurances.flatMap {
-        index.occurrences(relatedToUSR: $0.symbol.usr, roles: .overrideOf)
+        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
+          override.relations.contains(where: {
+            $0.roles.contains(.childOf) && transitiveReceiverUsrs.contains($0.symbol.usr)
+          })
+        }
       }
     }
 
@@ -2206,3 +2216,17 @@ fileprivate struct DocumentNotificationRequestQueue {
     queue = []
   }
 }
+
+/// Returns the USRs of the subtypes of `usrs` as well as their subtypes, transitively.
+fileprivate func transitiveSubtypeClosure(ofUsrs usrs: [String], index: IndexStoreDB) -> [String] {
+  var result: [String] = []
+  for usr in usrs {
+    result.append(usr)
+    let directSubtypes = index.occurrences(ofUSR: usr, roles: .baseOf).flatMap { occurance in
+      occurance.relations.filter { $0.roles.contains(.baseOf) }.map(\.symbol.usr)
+    }
+    let transitiveSubtypes = transitiveSubtypeClosure(ofUsrs: directSubtypes, index: index)
+    result += transitiveSubtypes
+  }
+  return result
+}

Sources/SourceKitLSP/Swift/CursorInfo.swift

@@ -130,14 +130,18 @@ extension SwiftLanguageServer {
 
     let refactorActionsArray: SKDResponseArray? = dict[keys.refactor_actions]
 
+    let receiversArray: SKDResponseArray? = dict[keys.receivers]
+    let receiverUsrs = receiversArray?.compactMap { $0[keys.usr] as String? } ?? []
+
     return CursorInfo(
       SymbolDetails(
         name: dict[keys.name],
         containerName: nil,
         usr: dict[keys.usr],
         bestLocalDeclaration: location,
         kind: kind.asSymbolKind(self.sourcekitd.values),
-        isDynamic: dict[keys.isDynamic] ?? false
+        isDynamic: dict[keys.isDynamic] ?? false,
+        receiverUsrs: receiverUsrs
       ),
       annotatedDeclaration: dict[keys.annotated_decl],
       documentationXML: dict[keys.doc_full_as_xml],

Tests/SourceKitLSPTests/DefinitionTests.swift

@@ -69,4 +69,40 @@ class DefinitionTests: XCTestCase {
       ]
     )
   }
+
+  func testJumpToDefinitionFiltersByReceiver() async throws {
+    let ws = try await IndexedSingleSwiftFileWorkspace(
+      """
+      class A {
+        func 1️⃣doThing() {}
+      }
+      class B: A {}
+      class C: B {
+        override func 2️⃣doThing() {}
+      }
+      class D: A {
+        override func doThing() {}
+      }
+
+      func test(value: B) {
+        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️⃣"])),
+      ]
+    )
+  }
 }