When performing jump-to-definition on a method implementing a protocol requirement, jump to the requirement

rdar://129412482

Author: ahoppen

Sources/SourceKitLSP/Clang/ClangLanguageService.swift

@@ -338,6 +338,10 @@ actor ClangLanguageService: LanguageService, MessageHandler {
     return try await clangd.send(request)
   }
 
+  public func canonicalDeclarationPosition(of position: Position, in uri: DocumentURI) async -> Position? {
+    return nil
+  }
+
   func _crash() {
     // Since `clangd` doesn't have a method to crash it, kill it.
     #if os(Windows)

Sources/SourceKitLSP/LanguageService.swift

@@ -132,7 +132,7 @@ public protocol LanguageService: AnyObject, Sendable {
 
   // MARK: - Build System Integration
 
-  /// Sent when the `BuildSystem` has resolved build settings, such as for the intial build settings
+  /// Sent when the `BuildSystem` has resolved build settings, such as for the initial build settings
   /// or when the settings have changed (e.g. modified build system files). This may be sent before
   /// the respective `DocumentURI` has been opened.
   func documentUpdatedBuildSettings(_ uri: DocumentURI) async
@@ -226,6 +226,16 @@ public protocol LanguageService: AnyObject, Sendable {
   /// A return value of `nil` indicates that this language service does not support syntactic test discovery.
   func syntacticDocumentTests(for uri: DocumentURI, in workspace: Workspace) async throws -> [AnnotatedTestItem]?
 
+  /// A position that is canonical for all positions within a declaration. For example, if we have the following
+  /// declaration, then all `|` markers should return the same canonical position.
+  /// ```
+  /// func |fo|o(|ba|r: Int)
+  /// ```
+  /// The actual position returned by the method does not matter. All that's relevant is the canonicalization.
+  ///
+  /// Returns `nil` if no canonical position could be determined.
+  func canonicalDeclarationPosition(of position: Position, in uri: DocumentURI) async -> Position?
+
   /// Crash the language server. Should be used for crash recovery testing only.
   func _crash() async
 }

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -1757,23 +1757,87 @@ extension SourceKitLSPServer {
         position: req.position
       )
     )
-    let locations = try await symbols.asyncMap { (symbol) -> [Location] in
-      if symbol.bestLocalDeclaration != nil && !(symbol.isDynamic ?? true)
-        && symbol.usr?.hasPrefix("s:") ?? false /* Swift symbols have USRs starting with s: */
+
+    let canonicalOriginatorLocation = await languageService.canonicalDeclarationPosition(
+      of: req.position,
+      in: req.textDocument.uri
+    )
+
+    // Returns `true` if `location` points to the same declaration that the definition request was initiated from.
+    @Sendable func isAtCanonicalOriginatorLocation(_ location: Location) async -> Bool {
+      guard location.uri == req.textDocument.uri, let canonicalOriginatorLocation else {
+        return false
+      }
+      return await languageService.canonicalDeclarationPosition(of: location.range.lowerBound, in: location.uri)
+        == canonicalOriginatorLocation
+    }
+
+    var locations = try await symbols.asyncMap { (symbol) -> [Location] in
+      var locations: [Location]
+      if let bestLocalDeclaration = symbol.bestLocalDeclaration,
+        !(symbol.isDynamic ?? true),
+        symbol.usr?.hasPrefix("s:") ?? false /* Swift symbols have USRs starting with s: */
       {
         // If we have a known non-dynamic symbol within Swift, we don't need to do an index lookup.
         // For non-Swift symbols, we need to perform an index lookup because the best local declaration will point to
-        // a header file but jump-to-definition should prefer the implementation (there's the declaration request to jump
-        // to the function's declaration).
-        return [symbol.bestLocalDeclaration].compactMap { $0 }
+        // a header file but jump-to-definition should prefer the implementation (there's the declaration request to
+        // jump to the function's declaration).
+        locations = [bestLocalDeclaration]
+      } else {
+        locations = try await self.definitionLocations(
+          for: symbol,
+          in: req.textDocument.uri,
+          languageService: languageService
+        )
       }
-      return try await self.definitionLocations(for: symbol, in: req.textDocument.uri, languageService: languageService)
+
+      // If the symbol's location is is where we initiated rename from, also show the declarations that the symbol
+      // overrides.
+      if let location = locations.only,
+        let usr = symbol.usr,
+        let index = workspace.index(checkedFor: .deletedFiles),
+        await isAtCanonicalOriginatorLocation(location)
+      {
+        let baseUSRs = index.occurrences(ofUSR: usr, roles: .overrideOf).flatMap {
+          $0.relations.filter { $0.roles.contains(.overrideOf) }.map(\.symbol.usr)
+        }
+        locations += baseUSRs.compactMap {
+          guard let baseDeclOccurrence = index.primaryDefinitionOrDeclarationOccurrence(ofUSR: $0) else {
+            return nil
+          }
+          return indexToLSPLocation(baseDeclOccurrence.location)
+        }
+      }
+
+      return locations
     }.flatMap { $0 }
+
     // Remove any duplicate locations. We might end up with duplicate locations when performing a definition request
     // on eg. `MyStruct()` when no explicit initializer is declared. In this case we get two symbol infos, one for the
     // declaration of the `MyStruct` type and one for the initializer, which is implicit and thus has the location of
     // the `MyStruct` declaration itself.
-    return locations.unique
+    locations = locations.unique
+
+    // Try removing any results that would point back to the location we are currently at. This ensures that eg. in the
+    // following case we only show line 2 when performing jump-to-definition on `TestImpl.doThing`.
+    //
+    // ```
+    // protocol TestProtocol {
+    //   func doThing()
+    // }
+    // struct TestImpl: TestProtocol {
+    //   func doThing() { }
+    // }
+    // ```
+    //
+    // If this would result in no locations, don't apply the filter. This way, performing jump-to-definition in the
+    // middle of a function's base name takes us to the base name start, indicating that jump-to-definition was able to
+    // resolve the location and didn't fail.
+    let nonOriginatorLocations = await locations.asyncFilter { await !isAtCanonicalOriginatorLocation($0) }
+    if !nonOriginatorLocations.isEmpty {
+      locations = nonOriginatorLocations
+    }
+    return locations
   }
 
   func definition(

Sources/SourceKitLSP/Swift/CodeActions/SyntaxRefactoringCodeActionProvider.swift

@@ -121,13 +121,13 @@ extension RemoveSeparatorsFromIntegerLiteral: SyntaxRefactoringCodeActionProvide
   }
 }
 
-extension Syntax {
+extension SyntaxProtocol {
   /// Finds the innermost parent of the given type while not walking outside of nodes that satisfy `stoppingIf`.
   func findParentOfSelf<ParentType: SyntaxProtocol>(
     ofType: ParentType.Type,
     stoppingIf: (Syntax) -> Bool
   ) -> ParentType? {
-    var node: Syntax? = self
+    var node: Syntax? = Syntax(self)
     while let unwrappedNode = node, !stoppingIf(unwrappedNode) {
       if let expectedType = unwrappedNode.as(ParentType.self) {
         return expectedType

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -318,6 +318,21 @@ extension SwiftLanguageService {
     await self.sourcekitd.removeNotificationHandler(self)
   }
 
+  public func canonicalDeclarationPosition(of position: Position, in uri: DocumentURI) async -> Position? {
+    guard let snapshot = try? documentManager.latestSnapshot(uri) else {
+      return nil
+    }
+    let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
+    let decl = syntaxTree.token(at: snapshot.absolutePosition(of: position))?.findParentOfSelf(
+      ofType: DeclSyntax.self,
+      stoppingIf: { $0.is(CodeBlockSyntax.self) || $0.is(MemberBlockSyntax.self) }
+    )
+    guard let decl else {
+      return nil
+    }
+    return snapshot.position(of: decl.positionAfterSkippingLeadingTrivia)
+  }
+
   /// Tell sourcekitd to crash itself. For testing purposes only.
   public func _crash() async {
     let req = sourcekitd.dictionary([

Tests/SourceKitLSPTests/DefinitionTests.swift

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+@_spi(Testing) import LSPLogging
 import LanguageServerProtocol
 import SKTestSupport
 import XCTest
@@ -541,4 +542,63 @@ class DefinitionTests: XCTestCase {
       ])
     )
   }
+
+  func testJumpToDefinitionOnProtocolImplementationJumpsToRequirement() async throws {
+    let project = try await IndexedSingleSwiftFileTestProject(
+      """
+      protocol TestProtocol {
+        func 1️⃣doThing()
+      }
+
+      struct TestImpl: TestProtocol {
+        func 2️⃣do3️⃣Thing() { }
+      }
+      """
+    )
+
+    let definitionFromBaseName = try await project.testClient.send(
+      DefinitionRequest(textDocument: TextDocumentIdentifier(project.fileURI), position: project.positions["2️⃣"])
+    )
+    XCTAssertEqual(
+      definitionFromBaseName,
+      .locations([Location(uri: project.fileURI, range: Range(project.positions["1️⃣"]))])
+    )
+
+    let definitionFromInsideBaseName = try await project.testClient.send(
+      DefinitionRequest(textDocument: TextDocumentIdentifier(project.fileURI), position: project.positions["3️⃣"])
+    )
+    XCTAssertEqual(
+      definitionFromInsideBaseName,
+      .locations([Location(uri: project.fileURI, range: Range(project.positions["1️⃣"]))])
+    )
+  }
+
+  func testJumpToDefinitionOnProtocolImplementationShowsAllFulfilledRequirements() async throws {
+    let project = try await IndexedSingleSwiftFileTestProject(
+      """
+      protocol TestProtocol {
+        func 1️⃣doThing()
+      }
+
+      protocol OtherProtocol {
+        func 2️⃣doThing()
+      }
+
+      struct TestImpl: TestProtocol, OtherProtocol {
+        func 3️⃣doThing() { }
+      }
+      """
+    )
+
+    let result = try await project.testClient.send(
+      DefinitionRequest(textDocument: TextDocumentIdentifier(project.fileURI), position: project.positions["3️⃣"])
+    )
+    XCTAssertEqual(
+      result,
+      .locations([
+        Location(uri: project.fileURI, range: Range(project.positions["1️⃣"])),
+        Location(uri: project.fileURI, range: Range(project.positions["2️⃣"])),
+      ])
+    )
+  }
 }