From 477d33c89c930d4bc82990db869e0db736060a0e Mon Sep 17 00:00:00 2001 From: somiljain2006 Date: Sat, 30 May 2026 21:55:01 +0530 Subject: [PATCH 1/2] Fix references for local variables and parameters --- Sources/SourceKitLSP/LanguageService.swift | 12 ++ Sources/SourceKitLSP/SourceKitLSPServer.swift | 12 +- .../SwiftLanguageService.swift | 28 +++ Tests/SourceKitLSPTests/ReferencesTests.swift | 180 ++++++++++++++++++ 4 files changed, 231 insertions(+), 1 deletion(-) diff --git a/Sources/SourceKitLSP/LanguageService.swift b/Sources/SourceKitLSP/LanguageService.swift index e5733139e..41047f812 100644 --- a/Sources/SourceKitLSP/LanguageService.swift +++ b/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 [] + } +} diff --git a/Sources/SourceKitLSP/SourceKitLSPServer.swift b/Sources/SourceKitLSP/SourceKitLSPServer.swift index 59bf0b7fb..bbec4a7e2 100644 --- a/Sources/SourceKitLSP/SourceKitLSPServer.swift +++ b/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,15 @@ 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() diff --git a/Sources/SwiftLanguageService/SwiftLanguageService.swift b/Sources/SwiftLanguageService/SwiftLanguageService.swift index 56c8c2427..1461aac2e 100644 --- a/Sources/SwiftLanguageService/SwiftLanguageService.swift +++ b/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) + } + } +} diff --git a/Tests/SourceKitLSPTests/ReferencesTests.swift b/Tests/SourceKitLSPTests/ReferencesTests.swift index c75c106c7..078c56289 100644 --- a/Tests/SourceKitLSPTests/ReferencesTests.swift +++ b/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️⃣"], + ]) + ) + } } From ee42692b088bd59c1df5da157d6a3fca8216c01a Mon Sep 17 00:00:00 2001 From: somiljain2006 Date: Sat, 30 May 2026 22:03:04 +0530 Subject: [PATCH 2/2] Fix formatting --- Sources/SourceKitLSP/SourceKitLSPServer.swift | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/Sources/SourceKitLSP/SourceKitLSPServer.swift b/Sources/SourceKitLSP/SourceKitLSPServer.swift index bbec4a7e2..5f8aec89c 100644 --- a/Sources/SourceKitLSP/SourceKitLSPServer.swift +++ b/Sources/SourceKitLSP/SourceKitLSPServer.swift @@ -2260,11 +2260,12 @@ extension SourceKitLSPServer { } if locations.isEmpty { - locations = (try? await languageService.localReferences( - at: req.position, - in: req.textDocument.uri, - includeDeclaration: req.context.includeDeclaration - )) ?? [] + locations = + (try? await languageService.localReferences( + at: req.position, + in: req.textDocument.uri, + includeDeclaration: req.context.includeDeclaration + )) ?? [] } let copiedFileMap = await workspace.buildServerManager.cachedCopiedFileMap