Fix references for local variables and parameters

Sources/SourceKitLSP/LanguageService.swift

@@ -336,6 +336,8 @@ package protocol LanguageService: AnyObject, Sendable {
 
   /// Crash the language server. Should be used for crash recovery testing only.
   func crash() async
+
+  func localReferences(at position: Position, in uri: DocumentURI, includeDeclaration: Bool) async throws -> [Location]
 }
 
 /// Default implementations for methods that satisfy the following criteria:
@@ -550,3 +552,13 @@ package extension LanguageService {
     logger.error("\(Self.self) cannot be crashed")
   }
 }
+
+extension LanguageService {
+  package func localReferences(
+    at position: Position,
+    in uri: DocumentURI,
+    includeDeclaration: Bool
+  ) async throws -> [Location] {
+    return []
+  }
+}

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -2248,7 +2248,8 @@ extension SourceKitLSPServer {
     guard let index = await workspaceForDocument(uri: req.textDocument.uri)?.index(checkedFor: .deletedFiles) else {
       return []
     }
-    let locations = try symbols.flatMap { (symbol) -> [Location] in
+
+    var locations = try symbols.flatMap { (symbol) -> [Location] in
       guard let usr = symbol.usr else { return [] }
       logger.info("Finding references for USR \(usr)")
       var roles: SymbolRole = [.reference]
@@ -2257,6 +2258,16 @@ extension SourceKitLSPServer {
       }
       return try index.occurrences(ofUSR: usr, roles: roles).compactMap { $0.location.lspLocation }
     }
+
+    if locations.isEmpty {
+      locations =
+        (try? await languageService.localReferences(
+          at: req.position,
+          in: req.textDocument.uri,
+          includeDeclaration: req.context.includeDeclaration
+        )) ?? []
+    }
+
     let copiedFileMap = await workspace.buildServerManager.cachedCopiedFileMap
     let remappedLocations = locations.adjusted(for: copiedFileMap)
     return remappedLocations.unique.sorted()

Sources/SwiftLanguageService/SwiftLanguageService.swift

@@ -1420,3 +1420,31 @@ extension SwiftLanguageService {
     return false
   }
 }
+
+extension SwiftLanguageService {
+  package func localReferences(
+    at position: Position,
+    in uri: DocumentURI,
+    includeDeclaration: Bool
+  ) async throws -> [Location] {
+    guard let snapshot = try? await latestSnapshot(for: uri) else {
+      return []
+    }
+
+    let response = try await self.relatedIdentifiers(
+      at: position,
+      in: snapshot,
+      includeNonEditableBaseNames: false
+    )
+
+    var identifiers = response.relatedIdentifiers
+
+    if !includeDeclaration {
+      identifiers = Array(identifiers.dropFirst())
+    }
+
+    return identifiers.map {
+      Location(uri: uri, range: $0.range)
+    }
+  }
+}

Tests/SourceKitLSPTests/ReferencesTests.swift

@@ -72,4 +72,184 @@ final class ReferencesTests: SourceKitLSPTestCase {
       ]
     )
   }
+
+  func testLocalVariableReferences() async throws {
+    let project = try await IndexedSingleSwiftFileTestProject(
+      """
+      func testReferences() {
+        let 1️⃣myLocalVariable = "Hello"
+        print(2️⃣myLocalVariable)
+
+        if true {
+          let 3️⃣myLocalVariable = "Shadowed"
+          print(4️⃣myLocalVariable)
+        }
+
+        let stringLength = 5️⃣myLocalVariable.count
+      }
+      """
+    )
+
+    let outerRefs = try await project.testClient.send(
+      ReferencesRequest(
+        textDocument: TextDocumentIdentifier(project.fileURI),
+        position: project.positions["1️⃣"],
+        context: ReferencesContext(includeDeclaration: true)
+      )
+    )
+
+    let outerRefPositions = outerRefs.map { $0.range.lowerBound }
+    XCTAssertEqual(outerRefPositions.count, 3)
+    XCTAssertTrue(outerRefPositions.contains(project.positions["1️⃣"]))
+    XCTAssertTrue(outerRefPositions.contains(project.positions["2️⃣"]))
+    XCTAssertTrue(outerRefPositions.contains(project.positions["5️⃣"]))
+    XCTAssertFalse(outerRefPositions.contains(project.positions["3️⃣"]))
+    XCTAssertFalse(outerRefPositions.contains(project.positions["4️⃣"]))
+
+    let innerRefs = try await project.testClient.send(
+      ReferencesRequest(
+        textDocument: TextDocumentIdentifier(project.fileURI),
+        position: project.positions["4️⃣"],
+        context: ReferencesContext(includeDeclaration: true)
+      )
+    )
+
+    let innerRefPositions = innerRefs.map { $0.range.lowerBound }
+    XCTAssertEqual(innerRefPositions.count, 2)
+    XCTAssertTrue(innerRefPositions.contains(project.positions["3️⃣"]))
+    XCTAssertTrue(innerRefPositions.contains(project.positions["4️⃣"]))
+
+    let outerRefsWithoutDecl = try await project.testClient.send(
+      ReferencesRequest(
+        textDocument: TextDocumentIdentifier(project.fileURI),
+        position: project.positions["1️⃣"],
+        context: ReferencesContext(includeDeclaration: false)
+      )
+    )
+
+    let outerRefsWithoutDeclPositions = outerRefsWithoutDecl.map { $0.range.lowerBound }
+    XCTAssertEqual(outerRefsWithoutDeclPositions.count, 2)
+    XCTAssertTrue(outerRefsWithoutDeclPositions.contains(project.positions["2️⃣"]))
+    XCTAssertTrue(outerRefsWithoutDeclPositions.contains(project.positions["5️⃣"]))
+    XCTAssertFalse(outerRefsWithoutDeclPositions.contains(project.positions["1️⃣"]))
+    XCTAssertFalse(outerRefsWithoutDeclPositions.contains(project.positions["3️⃣"]))
+    XCTAssertFalse(outerRefsWithoutDeclPositions.contains(project.positions["4️⃣"]))
+  }
+
+  func testSimpleLocalVariableReferences() async throws {
+    let project = try await IndexedSingleSwiftFileTestProject(
+      """
+      func foo() {
+        let 1️⃣x = 1
+        print(2️⃣x)
+        print(3️⃣x)
+      }
+      """
+    )
+
+    let refs = try await project.testClient.send(
+      ReferencesRequest(
+        textDocument: TextDocumentIdentifier(project.fileURI),
+        position: project.positions["2️⃣"],
+        context: ReferencesContext(includeDeclaration: true)
+      )
+    )
+
+    XCTAssertEqual(
+      Set(refs.map(\.range.lowerBound)),
+      Set([
+        project.positions["1️⃣"],
+        project.positions["2️⃣"],
+        project.positions["3️⃣"],
+      ])
+    )
+  }
+
+  func testLocalVariableReferencesWithoutDeclaration() async throws {
+    let project = try await IndexedSingleSwiftFileTestProject(
+      """
+      func foo() {
+        let 1️⃣x = 1
+        print(2️⃣x)
+        print(3️⃣x)
+      }
+      """
+    )
+
+    let responseFromUsage = try await project.testClient.send(
+      ReferencesRequest(
+        textDocument: TextDocumentIdentifier(project.fileURI),
+        position: project.positions["2️⃣"],
+        context: ReferencesContext(includeDeclaration: false)
+      )
+    )
+
+    XCTAssertEqual(
+      Set(responseFromUsage.map(\.range.lowerBound)),
+      Set([
+        project.positions["2️⃣"],
+        project.positions["3️⃣"],
+      ])
+    )
+
+    let responseFromDeclaration = try await project.testClient.send(
+      ReferencesRequest(
+        textDocument: TextDocumentIdentifier(project.fileURI),
+        position: project.positions["1️⃣"],
+        context: ReferencesContext(includeDeclaration: false)
+      )
+    )
+
+    XCTAssertEqual(
+      Set(responseFromDeclaration.map(\.range.lowerBound)),
+      Set([
+        project.positions["2️⃣"],
+        project.positions["3️⃣"],
+      ])
+    )
+  }
+
+  func testParameterReferences() async throws {
+    let project = try await IndexedSingleSwiftFileTestProject(
+      """
+      func foo(1️⃣x: Int) {
+        print(2️⃣x)
+        print(3️⃣x)
+      }
+      """
+    )
+
+    let responseWithoutDecl = try await project.testClient.send(
+      ReferencesRequest(
+        textDocument: TextDocumentIdentifier(project.fileURI),
+        position: project.positions["1️⃣"],
+        context: ReferencesContext(includeDeclaration: false)
+      )
+    )
+
+    XCTAssertEqual(
+      Set(responseWithoutDecl.map(\.range.lowerBound)),
+      Set([
+        project.positions["2️⃣"],
+        project.positions["3️⃣"],
+      ])
+    )
+
+    let responseWithDecl = try await project.testClient.send(
+      ReferencesRequest(
+        textDocument: TextDocumentIdentifier(project.fileURI),
+        position: project.positions["1️⃣"],
+        context: ReferencesContext(includeDeclaration: true)
+      )
+    )
+
+    XCTAssertEqual(
+      Set(responseWithDecl.map(\.range.lowerBound)),
+      Set([
+        project.positions["1️⃣"],
+        project.positions["2️⃣"],
+        project.positions["3️⃣"],
+      ])
+    )
+  }
 }