From 28ea629c9e42b2835a1a5ad8d41ef9380a779d09 Mon Sep 17 00:00:00 2001 From: Pavel Kroupa <63880977+Tabonx@users.noreply.github.com> Date: Thu, 12 Mar 2026 00:10:26 +0100 Subject: [PATCH 1/2] Add references code lens for symbol declarations --- .../SwiftCodeLensScanner.swift | 287 ++++++++++++----- .../SwiftLanguageService.swift | 3 +- Tests/SourceKitLSPTests/CodeLensTests.swift | 289 ++++++++++++++++++ 3 files changed, 506 insertions(+), 73 deletions(-) diff --git a/Sources/SwiftLanguageService/SwiftCodeLensScanner.swift b/Sources/SwiftLanguageService/SwiftCodeLensScanner.swift index 0d2c33e14..e818c2e60 100644 --- a/Sources/SwiftLanguageService/SwiftCodeLensScanner.swift +++ b/Sources/SwiftLanguageService/SwiftCodeLensScanner.swift @@ -12,12 +12,16 @@ internal import BuildServerIntegration import BuildServerProtocol +internal import IndexStoreDB @_spi(SourceKitLSP) import LanguageServerProtocol +@_spi(SourceKitLSP) import SKLogging +import SemanticIndex import SourceKitLSP import SwiftSyntax import ToolchainRegistry -/// Scans a source file for classes or structs annotated with `@main` and returns a code lens for them. +/// Scans a source file for code lenses including `@main` run/debug actions, +/// symbol reference counts, and playground entries. final class SwiftCodeLensScanner: SyntaxVisitor { /// The document snapshot of the syntax tree that is being walked. private let snapshot: DocumentSnapshot @@ -25,119 +29,258 @@ final class SwiftCodeLensScanner: SyntaxVisitor { /// The collection of CodeLenses found in the document. private var result: [CodeLens] = [] + /// The display name of the build target containing this document, if available. private let targetName: String? - /// The map of supported commands and their client side command names + /// The language service used to resolve cursor info for symbols. + private let languageService: SwiftLanguageService + + /// The map of supported commands and their client side command names. private let supportedCommands: [SupportedCodeLensCommand: String] + /// Symbols collected during the syntax walk, processed asynchronously afterward. + private var symbolsToProcess: [(nameToken: TokenSyntax, displayRange: Range)] = [] + + private let workspace: Workspace? + private init( snapshot: DocumentSnapshot, targetName: String?, - supportedCommands: [SupportedCodeLensCommand: String] + supportedCommands: [SupportedCodeLensCommand: String], + workspace: Workspace?, + languageService: SwiftLanguageService ) { self.snapshot = snapshot self.targetName = targetName self.supportedCommands = supportedCommands + self.workspace = workspace + self.languageService = languageService super.init(viewMode: .fixedUp) } - /// Public entry point. Scans the syntax tree of the given snapshot for an `@main` annotation - /// and returns CodeLens's with Commands to run/debug the application. + /// Public entry point. Scans the syntax tree of the given snapshot and returns + /// all applicable code lenses including `@main` run/debug actions, reference counts, + /// and playground entries. public static func findCodeLenses( in snapshot: DocumentSnapshot, workspace: Workspace?, syntaxTreeManager: SyntaxTreeManager, supportedCommands: [SupportedCodeLensCommand: String], - toolchain: Toolchain + toolchain: Toolchain, + languageService: SwiftLanguageService ) async -> [CodeLens] { guard !supportedCommands.isEmpty else { return [] } - var targetDisplayName: String? = nil - if let workspace, - let target = await workspace.buildServerManager.canonicalTarget(for: snapshot.uri), - let buildTarget = await workspace.buildServerManager.buildTarget(named: target) - { - targetDisplayName = buildTarget.displayName - } + let targetDisplayName = await resolveTargetDisplayName(for: snapshot, workspace: workspace) - var codeLenses: [CodeLens] = [] - if snapshot.text.contains("@main") { - let visitor = SwiftCodeLensScanner( - snapshot: snapshot, - targetName: targetDisplayName, - supportedCommands: supportedCommands - ) - let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot) - visitor.walk(syntaxTree) - codeLenses += visitor.result - } + // Process @main annotations and symbol references + let visitor = SwiftCodeLensScanner( + snapshot: snapshot, + targetName: targetDisplayName, + supportedCommands: supportedCommands, + workspace: workspace, + languageService: languageService + ) + let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot) + visitor.walk(syntaxTree) - // "swift.play" CodeLens should be ignored if "swift-play" is not in the toolchain as the client has no way of running - if toolchain.swiftPlay != nil, - let workspace, - let playCommand = supportedCommands[SupportedCodeLensCommand.play] - { - let playgrounds = await SwiftPlaygroundsScanner.findDocumentPlaygrounds( - for: snapshot, - workspace: workspace, - syntaxTreeManager: syntaxTreeManager - ) - codeLenses += playgrounds.map({ - CodeLens( - range: $0.range, - command: Command( - title: "Play \"\($0.label ?? $0.id)\"", - command: playCommand, - arguments: [$0.encodeToLSPAny()] - ) - ) - }) + // Process collected symbols asynchronously for reference counts + for (nameToken, displayRange) in visitor.symbolsToProcess { + await visitor.captureReferenceLens(for: nameToken, at: displayRange) } + var codeLenses = visitor.result + + // Append playground lenses if swift-play is available in the toolchain + codeLenses += await playgroundLenses( + for: snapshot, + workspace: workspace, + toolchain: toolchain, + syntaxTreeManager: syntaxTreeManager, + supportedCommands: supportedCommands + ) + return codeLenses } override func visit(_ node: ClassDeclSyntax) -> SyntaxVisitorContinueKind { - node.attributes.forEach(self.captureLensFromAttribute) - return .skipChildren + node.attributes.forEach(captureMainAttributeLens) + symbolsToProcess.append((nameToken: node.name, displayRange: node.trimmedRange)) + return .visitChildren + } + + override func visit(_ node: EnumDeclSyntax) -> SyntaxVisitorContinueKind { + node.attributes.forEach(captureMainAttributeLens) + symbolsToProcess.append((nameToken: node.name, displayRange: node.trimmedRange)) + return .visitChildren + } + + override func visit(_ node: ActorDeclSyntax) -> SyntaxVisitorContinueKind { + symbolsToProcess.append((nameToken: node.name, displayRange: node.trimmedRange)) + return .visitChildren } override func visit(_ node: StructDeclSyntax) -> SyntaxVisitorContinueKind { - node.attributes.forEach(self.captureLensFromAttribute) - return .skipChildren - } - - private func captureLensFromAttribute(attribute: AttributeListSyntax.Element) { - if attribute.trimmedDescription == "@main" { - let range = self.snapshot.absolutePositionRange(of: attribute.trimmedRange) - var targetNameToAppend: String = "" - var arguments: [LSPAny] = [] - if let targetName { - targetNameToAppend = " \(targetName)" - arguments.append(.string(targetName)) + node.attributes.forEach(captureMainAttributeLens) + symbolsToProcess.append((nameToken: node.name, displayRange: node.trimmedRange)) + return .visitChildren + } + + override func visit(_ node: FunctionDeclSyntax) -> SyntaxVisitorContinueKind { + symbolsToProcess.append((nameToken: node.name, displayRange: node.trimmedRange)) + return .visitChildren + } + + override func visit(_ node: VariableDeclSyntax) -> SyntaxVisitorContinueKind { + // Only show references for member-level variables, not locals. + guard node.parent?.is(MemberBlockItemSyntax.self) == true else { + return .visitChildren + } + // Variable declarations can have multiple bindings (e.g., let x = 1, y = 2) + for binding in node.bindings { + if let identifier = binding.pattern.as(IdentifierPatternSyntax.self) { + symbolsToProcess.append((nameToken: identifier.identifier, displayRange: binding.trimmedRange)) } + } + return .visitChildren + } - if let runCommand = supportedCommands[SupportedCodeLensCommand.run] { - // Return commands for running/debugging the executable. - // These command names must be recognized by the client and so should not be chosen arbitrarily. - self.result.append( - CodeLens( - range: range, - command: Command(title: "Run" + targetNameToAppend, command: runCommand, arguments: arguments) - ) + override func visit(_ node: ProtocolDeclSyntax) -> SyntaxVisitorContinueKind { + symbolsToProcess.append((nameToken: node.name, displayRange: node.trimmedRange)) + return .visitChildren + } + + override func visit(_ node: TypeAliasDeclSyntax) -> SyntaxVisitorContinueKind { + symbolsToProcess.append((nameToken: node.name, displayRange: node.trimmedRange)) + return .visitChildren + } + + override func visit(_ node: InitializerDeclSyntax) -> SyntaxVisitorContinueKind { + symbolsToProcess.append((nameToken: node.initKeyword, displayRange: node.trimmedRange)) + return .visitChildren + } + + override func visit(_ node: SubscriptDeclSyntax) -> SyntaxVisitorContinueKind { + symbolsToProcess.append((nameToken: node.subscriptKeyword, displayRange: node.trimmedRange)) + return .visitChildren + } + + /// Adds run and debug code lenses for `@main` attributes. + private func captureMainAttributeLens(attribute: AttributeListSyntax.Element) { + guard attribute.trimmedDescription == "@main" else { + return + } + + let range = snapshot.absolutePositionRange(of: attribute.trimmedRange) + let suffix = targetName.map { " \($0)" } ?? "" + let arguments: [LSPAny] = targetName.map { [.string($0)] } ?? [] + + // Return commands for running/debugging the executable. + // These command names must be recognized by the client and so should not be chosen arbitrarily. + if let runCommand = supportedCommands[.run] { + result.append( + CodeLens( + range: range, + command: Command(title: "Run" + suffix, command: runCommand, arguments: arguments) + ) + ) + } + if let debugCommand = supportedCommands[.debug] { + result.append( + CodeLens( + range: range, + command: Command(title: "Debug" + suffix, command: debugCommand, arguments: arguments) ) + ) + } + } + + /// Queries the index for the number of references to a symbol and appends a code lens with the count. + private func captureReferenceLens(for nameToken: TokenSyntax, at displayRange: Range) async { + guard let referencesCommand = supportedCommands[.references] else { return } + + let lensRange = snapshot.absolutePositionRange(of: displayRange) + let nameRange = snapshot.absolutePositionRange(of: nameToken.trimmedRange) + + do { + let cursorInfoResults = try await languageService.cursorInfo( + snapshot, + compileCommand: await languageService.compileCommand(for: snapshot.uri, fallbackAfterTimeout: false), + nameRange + ) + .cursorInfo + + guard let cursorInfo = cursorInfoResults.first, + let usr = cursorInfo.symbolInfo.usr, + let index = await workspace?.index(checkedFor: .deletedFiles) + else { return } + + var referenceCount = 0 + index.forEachSymbolOccurrence(byUSR: usr, roles: .reference) { _ in + referenceCount += 1 + return true } - if let debugCommand = supportedCommands[SupportedCodeLensCommand.debug] { - self.result.append( - CodeLens( - range: range, - command: Command(title: "Debug" + targetNameToAppend, command: debugCommand, arguments: arguments) + let title = "\(referenceCount) reference\(referenceCount == 1 ? "" : "s")" + result.append( + CodeLens( + range: lensRange, + command: Command( + title: title, + command: referencesCommand, + arguments: [.string(snapshot.uri.stringValue), nameRange.lowerBound.encodeToLSPAny()] ) ) - } + ) + } catch { + logger.info("Failed to get cursor info for reference count: \(error.forLogging, privacy: .public)") + } + } + + /// Resolves the display name of the build target containing the given document. + private static func resolveTargetDisplayName(for snapshot: DocumentSnapshot, workspace: Workspace?) async -> String? { + guard let workspace, + let target = await workspace.buildServerManager.canonicalTarget(for: snapshot.uri), + let buildTarget = await workspace.buildServerManager.buildTarget(named: target) + else { + return nil + } + return buildTarget.displayName + } + + /// Returns playground code lenses if swift-play is available in the toolchain. + private static func playgroundLenses( + for snapshot: DocumentSnapshot, + workspace: Workspace?, + toolchain: Toolchain, + syntaxTreeManager: SyntaxTreeManager, + supportedCommands: [SupportedCodeLensCommand: String] + ) async -> [CodeLens] { + // "swift.play" CodeLens should be ignored if "swift-play" is not in the toolchain + // as the client has no way of running it. + guard toolchain.swiftPlay != nil, + let workspace, + let playCommand = supportedCommands[.play] + else { + return [] + } + + let playgrounds = await SwiftPlaygroundsScanner.findDocumentPlaygrounds( + for: snapshot, + workspace: workspace, + syntaxTreeManager: syntaxTreeManager + ) + return playgrounds.map { + CodeLens( + range: $0.range, + command: Command( + title: "Play \"\($0.label ?? $0.id)\"", + command: playCommand, + arguments: [$0.encodeToLSPAny()] + ) + ) } } } diff --git a/Sources/SwiftLanguageService/SwiftLanguageService.swift b/Sources/SwiftLanguageService/SwiftLanguageService.swift index c54bffe47..7e53ea239 100644 --- a/Sources/SwiftLanguageService/SwiftLanguageService.swift +++ b/Sources/SwiftLanguageService/SwiftLanguageService.swift @@ -1080,7 +1080,8 @@ extension SwiftLanguageService { workspace: workspace, syntaxTreeManager: self.syntaxTreeManager, supportedCommands: self.capabilityRegistry.supportedCodeLensCommands, - toolchain: toolchain + toolchain: toolchain, + languageService: self ) } diff --git a/Tests/SourceKitLSPTests/CodeLensTests.swift b/Tests/SourceKitLSPTests/CodeLensTests.swift index 6c3b9bc26..ce760e683 100644 --- a/Tests/SourceKitLSPTests/CodeLensTests.swift +++ b/Tests/SourceKitLSPTests/CodeLensTests.swift @@ -459,4 +459,293 @@ final class CodeLensTests: SourceKitLSPTestCase { ] ) } + + // MARK: - References Code Lens Tests + + func testReferencesLensForFunction() async throws { + var codeLensCapabilities = TextDocumentClientCapabilities.CodeLens() + codeLensCapabilities.supportedCommands = [ + SupportedCodeLensCommand.references: "swift.references" + ] + let capabilities = ClientCapabilities(textDocument: TextDocumentClientCapabilities(codeLens: codeLensCapabilities)) + + let project = try await SwiftPMTestProject( + files: [ + "Sources/MyLibrary/Lib.swift": """ + 1️⃣public func 3️⃣greet4️⃣() { + print("hello") + }2️⃣ + """, + "Sources/MyLibrary/Usage.swift": """ + func test() { + greet() + } + """, + ], + capabilities: capabilities, + enableBackgroundIndexing: true + ) + + let (uri, positions) = try project.openDocument("Lib.swift") + + let response = try await project.testClient.send( + CodeLensRequest(textDocument: TextDocumentIdentifier(uri)) + ) + + XCTAssertEqual( + response, + [ + CodeLens( + range: positions["1️⃣"].. Date: Fri, 10 Apr 2026 11:27:04 +0200 Subject: [PATCH 2/2] Use batch soucekit request --- Sources/SourceKitD/sourcekitd_uids.swift | 3 + .../SwiftCodeLensScanner.swift | 124 +++++++++++++----- 2 files changed, 91 insertions(+), 36 deletions(-) diff --git a/Sources/SourceKitD/sourcekitd_uids.swift b/Sources/SourceKitD/sourcekitd_uids.swift index b9a0326d6..e808f4c07 100644 --- a/Sources/SourceKitD/sourcekitd_uids.swift +++ b/Sources/SourceKitD/sourcekitd_uids.swift @@ -886,6 +886,8 @@ package struct sourcekitd_api_requests { package let collectExpressionType: sourcekitd_api_uid_t /// `source.request.variable.type` package let collectVariableType: sourcekitd_api_uid_t + /// `source.request.declaration.usr` + package let collectDeclarationUSR: sourcekitd_api_uid_t /// `source.request.configuration.global` package let globalConfiguration: sourcekitd_api_uid_t /// `source.request.dependency_updated` @@ -955,6 +957,7 @@ package struct sourcekitd_api_requests { testNotification = api.uid_get_from_cstr("source.request.test_notification")! collectExpressionType = api.uid_get_from_cstr("source.request.expression.type")! collectVariableType = api.uid_get_from_cstr("source.request.variable.type")! + collectDeclarationUSR = api.uid_get_from_cstr("source.request.declaration.usr")! globalConfiguration = api.uid_get_from_cstr("source.request.configuration.global")! dependencyUpdated = api.uid_get_from_cstr("source.request.dependency_updated")! diagnostics = api.uid_get_from_cstr("source.request.diagnostics")! diff --git a/Sources/SwiftLanguageService/SwiftCodeLensScanner.swift b/Sources/SwiftLanguageService/SwiftCodeLensScanner.swift index e818c2e60..327f4343f 100644 --- a/Sources/SwiftLanguageService/SwiftCodeLensScanner.swift +++ b/Sources/SwiftLanguageService/SwiftCodeLensScanner.swift @@ -16,6 +16,7 @@ internal import IndexStoreDB @_spi(SourceKitLSP) import LanguageServerProtocol @_spi(SourceKitLSP) import SKLogging import SemanticIndex +import SourceKitD import SourceKitLSP import SwiftSyntax import ToolchainRegistry @@ -32,7 +33,7 @@ final class SwiftCodeLensScanner: SyntaxVisitor { /// The display name of the build target containing this document, if available. private let targetName: String? - /// The language service used to resolve cursor info for symbols. + /// The language service used to resolve symbol metadata for code lenses. private let languageService: SwiftLanguageService /// The map of supported commands and their client side command names. @@ -86,10 +87,7 @@ final class SwiftCodeLensScanner: SyntaxVisitor { let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot) visitor.walk(syntaxTree) - // Process collected symbols asynchronously for reference counts - for (nameToken, displayRange) in visitor.symbolsToProcess { - await visitor.captureReferenceLens(for: nameToken, at: displayRange) - } + await visitor.captureReferenceLenses() var codeLenses = visitor.result @@ -197,45 +195,51 @@ final class SwiftCodeLensScanner: SyntaxVisitor { } } - /// Queries the index for the number of references to a symbol and appends a code lens with the count. - private func captureReferenceLens(for nameToken: TokenSyntax, at displayRange: Range) async { - guard let referencesCommand = supportedCommands[.references] else { return } - - let lensRange = snapshot.absolutePositionRange(of: displayRange) - let nameRange = snapshot.absolutePositionRange(of: nameToken.trimmedRange) + /// Queries sourcekitd once for declaration USRs, then looks up reference counts in the index. + private func captureReferenceLenses() async { + guard let referencesCommand = supportedCommands[.references], + let index = await workspace?.index(checkedFor: .deletedFiles) + else { + return + } do { - let cursorInfoResults = try await languageService.cursorInfo( + let declarationUsrs = try await languageService.declarationUSRs( snapshot, - compileCommand: await languageService.compileCommand(for: snapshot.uri, fallbackAfterTimeout: false), - nameRange + compileCommand: await languageService.compileCommand(for: snapshot.uri, fallbackAfterTimeout: false) + ) + let usrsByOffset = Dictionary( + declarationUsrs.map { ($0.offset, $0.usr) }, + uniquingKeysWith: { first, _ in first } ) - .cursorInfo - - guard let cursorInfo = cursorInfoResults.first, - let usr = cursorInfo.symbolInfo.usr, - let index = await workspace?.index(checkedFor: .deletedFiles) - else { return } - - var referenceCount = 0 - index.forEachSymbolOccurrence(byUSR: usr, roles: .reference) { _ in - referenceCount += 1 - return true - } - let title = "\(referenceCount) reference\(referenceCount == 1 ? "" : "s")" - result.append( - CodeLens( - range: lensRange, - command: Command( - title: title, - command: referencesCommand, - arguments: [.string(snapshot.uri.stringValue), nameRange.lowerBound.encodeToLSPAny()] + for (nameToken, displayRange) in symbolsToProcess { + guard let usr = usrsByOffset[nameToken.trimmedRange.lowerBound.utf8Offset] else { + continue + } + + var referenceCount = 0 + try index.forEachSymbolOccurrence(byUSR: usr, roles: .reference) { _ in + referenceCount += 1 + return true + } + + let lensRange = snapshot.absolutePositionRange(of: displayRange) + let nameRange = snapshot.absolutePositionRange(of: nameToken.trimmedRange) + let title = "\(referenceCount) reference\(referenceCount == 1 ? "" : "s")" + result.append( + CodeLens( + range: lensRange, + command: Command( + title: title, + command: referencesCommand, + arguments: [.string(snapshot.uri.stringValue), nameRange.lowerBound.encodeToLSPAny()] + ) ) ) - ) + } } catch { - logger.info("Failed to get cursor info for reference count: \(error.forLogging, privacy: .public)") + logger.info("Failed to get declaration USRs for reference count: \(error.forLogging, privacy: .public)") } } @@ -284,3 +288,51 @@ final class SwiftCodeLensScanner: SyntaxVisitor { } } } + +private struct DeclarationUSRInfo { + let offset: Int + let usr: String +} + +extension SwiftLanguageService { + fileprivate func declarationUSRs( + _ snapshot: DocumentSnapshot, + compileCommand: SwiftCompileCommand?, + _ range: Range? = nil + ) async throws -> [DeclarationUSRInfo] { + let skreq = sourcekitd.dictionary([ + keys.cancelOnSubsequentRequest: 0, + keys.filePath: snapshot.uri.sourcekitdSourceFile, + keys.compilerArgs: compileCommand?.compilerArgs as [any SKDRequestValue]?, + ]) + + if let range { + let start = snapshot.utf8Offset(of: range.lowerBound) + let end = snapshot.utf8Offset(of: range.upperBound) + skreq.set(keys.offset, to: start) + skreq.set(keys.length, to: end - start) + } + + let dict = try await send(sourcekitdRequest: \.collectDeclarationUSR, skreq, snapshot: snapshot) + guard let declarations: SKDResponseArray = dict[keys.declarations] else { + return [] + } + + var result: [DeclarationUSRInfo] = [] + result.reserveCapacity(declarations.count) + + // swift-format-ignore: ReplaceForEachWithForLoop + declarations.forEach { (_, declaration) -> Bool in + guard let offset: Int = declaration[keys.offset], + let usr: String = declaration[keys.usr] + else { + assertionFailure("DeclarationUSRInfo failed to deserialize") + return true + } + result.append(DeclarationUSRInfo(offset: offset, usr: usr)) + return true + } + + return result + } +}