Add Run/Debug CodeLens Support

Adds a response to the textDocument/codeLens request that returns two
code lenses on the `@main` attribute of an application.

The LSP documentation breaks out the code lens requests into a
[`Code Lens Request`](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_codeLens)
and a
[`Code Lens Resolve Request`](https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#codeLens_resolve),
stating this is for performance reasons. However, there is no intensive
work we need to do in order to resolve the commands for a CodeLens;
we know them based on context at the time of discovery. For this reason
we return resolved lenses with Commands for code lens requests.

A missing piece is only returning code lenses if the file resides in an
executable product. To my knoledge Libraries and Plugins can't have an
`@main` entrypoint and so it doesn't make sense to provide these code
lenses in those contexts.

Some guidance is required on how to best determine if the textDocument
in the request is within an executable product.

`testCodeLensRequestWithInvalidProduct` asserts that no lenses are
returned with the `@main` attribute is on a file in a `.executable`, and
is currently failing until this is addressed.

Author: plemarquand

Sources/LanguageServerProtocol/SupportTypes/RegistrationOptions.swift

@@ -242,6 +242,48 @@ public struct DiagnosticRegistrationOptions: RegistrationOptions, TextDocumentRe
   }
 }
 
+/// Describe options to be used when registering for code lenses.
+public struct CodeLensRegistrationOptions: RegistrationOptions, TextDocumentRegistrationOptionsProtocol {
+  public var textDocumentRegistrationOptions: TextDocumentRegistrationOptions
+  public var codeLensOptions: CodeLensOptions
+
+  public init(
+    documentSelector: DocumentSelector? = nil,
+    codeLensOptions: CodeLensOptions
+  ) {
+    textDocumentRegistrationOptions = TextDocumentRegistrationOptions(documentSelector: documentSelector)
+    self.codeLensOptions = codeLensOptions
+  }
+
+  public init?(fromLSPDictionary dictionary: [String: LSPAny]) {
+    self.codeLensOptions = CodeLensOptions()
+
+    if case .bool(let resolveProvider) = dictionary["resolveProvider"] {
+      self.codeLensOptions.resolveProvider = resolveProvider
+    }
+
+    guard let textDocumentRegistrationOptions = TextDocumentRegistrationOptions(fromLSPDictionary: dictionary) else {
+      return nil
+    }
+
+    self.textDocumentRegistrationOptions = textDocumentRegistrationOptions
+  }
+
+  public func encodeToLSPAny() -> LSPAny {
+    var dict: [String: LSPAny] = [:]
+
+    if let resolveProvider = codeLensOptions.resolveProvider {
+      dict["resolveProvider"] = .bool(resolveProvider)
+    }
+
+    if case .dictionary(let dictionary) = textDocumentRegistrationOptions.encodeToLSPAny() {
+      dict.merge(dictionary) { (current, _) in current }
+    }
+
+    return .dictionary(dict)
+  }
+}
+
 /// Describe options to be used when registering for file system change events.
 public struct DidChangeWatchedFilesRegistrationOptions: RegistrationOptions {
   /// The watchers to register.

Sources/SourceKitLSP/CMakeLists.txt

@@ -56,6 +56,7 @@ target_sources(SourceKitLSP PRIVATE
   Swift/SemanticRefactoring.swift
   Swift/SemanticTokens.swift
   Swift/SourceKitD+ResponseError.swift
+  Swift/SwiftCodeLensScanner.swift
   Swift/SwiftCommand.swift
   Swift/SwiftLanguageService.swift
   Swift/SwiftTestingScanner.swift

Sources/SourceKitLSP/CapabilityRegistry.swift

@@ -38,6 +38,9 @@ public final actor CapabilityRegistry {
   /// Dynamically registered pull diagnostics options.
   private var pullDiagnostics: [CapabilityRegistration: DiagnosticRegistrationOptions] = [:]
 
+  /// Dynamically registered code lens options.
+  private var codeLens: [CapabilityRegistration: CodeLensRegistrationOptions] = [:]
+
   /// Dynamically registered file watchers.
   private var didChangeWatchedFiles: DidChangeWatchedFilesRegistrationOptions?
 
@@ -66,6 +69,10 @@ public final actor CapabilityRegistry {
     clientCapabilities.textDocument?.diagnostic?.dynamicRegistration == true
   }
 
+  public var clientHasDynamicDocumentCodeLensRegistration: Bool {
+    clientCapabilities.textDocument?.codeLens?.dynamicRegistration == true
+  }
+
   public var clientHasDynamicExecuteCommandRegistration: Bool {
     clientCapabilities.workspace?.executeCommand?.dynamicRegistration == true
   }
@@ -279,6 +286,7 @@ public final actor CapabilityRegistry {
     server: SourceKitLSPServer
   ) async {
     guard clientHasDynamicInlayHintRegistration else { return }
+
     await registerLanguageSpecificCapability(
       options: InlayHintRegistrationOptions(
         documentSelector: DocumentSelector(for: languages),
@@ -314,6 +322,29 @@ public final actor CapabilityRegistry {
     )
   }
 
+  /// Dynamically register code lens capabilities,
+  /// if the client supports it.
+  public func registerCodeLensIfNeeded(
+    options: CodeLensOptions,
+    for languages: [Language],
+    server: SourceKitLSPServer
+  ) async {
+    guard clientHasDynamicDocumentCodeLensRegistration else { return }
+
+    await registerLanguageSpecificCapability(
+      options: CodeLensRegistrationOptions(
+        // Code lenses should only apply to saved files
+        documentSelector: DocumentSelector(for: languages, scheme: "file"),
+        codeLensOptions: options
+      ),
+      forMethod: CodeLensRequest.method,
+      languages: languages,
+      in: server,
+      registrationDict: codeLens,
+      setRegistrationDict: { codeLens[$0] = $1 }
+    )
+  }
+
   /// Dynamically register executeCommand with the given IDs if the client supports
   /// it and we haven't yet registered the given command IDs yet.
   public func registerExecuteCommandIfNeeded(
@@ -345,7 +376,7 @@ public final actor CapabilityRegistry {
 }
 
 fileprivate extension DocumentSelector {
-  init(for languages: [Language]) {
-    self.init(languages.map { DocumentFilter(language: $0.rawValue) })
+  init(for languages: [Language], scheme: String? = nil) {
+    self.init(languages.map { DocumentFilter(language: $0.rawValue, scheme: scheme) })
   }
 }

Sources/SourceKitLSP/Clang/ClangLanguageService.swift

@@ -622,6 +622,10 @@ extension ClangLanguageService {
     return try await forwardRequestToClangd(req)
   }
 
+  func codeLens(_ req: CodeLensRequest) async throws -> [CodeLens] {
+    return try await forwardRequestToClangd(req) ?? []
+  }
+
   func foldingRange(_ req: FoldingRangeRequest) async throws -> [FoldingRange]? {
     guard self.capabilities?.foldingRangeProvider?.isSupported ?? false else {
       return nil

Sources/SourceKitLSP/LanguageService.swift

@@ -197,6 +197,7 @@ public protocol LanguageService: AnyObject, Sendable {
   func colorPresentation(_ req: ColorPresentationRequest) async throws -> [ColorPresentation]
   func codeAction(_ req: CodeActionRequest) async throws -> CodeActionRequestResponse?
   func inlayHint(_ req: InlayHintRequest) async throws -> [InlayHint]
+  func codeLens(_ req: CodeLensRequest) async throws -> [CodeLens]
   func documentDiagnostic(_ req: DocumentDiagnosticsRequest) async throws -> DocumentDiagnosticReport
   func documentFormatting(_ req: DocumentFormattingRequest) async throws -> [TextEdit]?
 

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -701,6 +701,8 @@ extension SourceKitLSPServer: MessageHandler {
       await self.handleRequest(for: request, requestHandler: self.prepareCallHierarchy)
     case let request as RequestAndReply<CodeActionRequest>:
       await self.handleRequest(for: request, requestHandler: self.codeAction)
+    case let request as RequestAndReply<CodeLensRequest>:
+      await self.handleRequest(for: request, requestHandler: self.codeLens)
     case let request as RequestAndReply<ColorPresentationRequest>:
       await self.handleRequest(for: request, requestHandler: self.colorPresentation)
     case let request as RequestAndReply<CompletionRequest>:
@@ -1101,6 +1103,7 @@ extension SourceKitLSPServer {
           supportsCodeActions: true
         )
       ),
+      codeLensProvider: CodeLensOptions(resolveProvider: false),
       documentFormattingProvider: .value(DocumentFormattingOptions(workDoneProgress: false)),
       renameProvider: .value(RenameOptions(prepareProvider: true)),
       colorProvider: .bool(true),
@@ -1161,6 +1164,9 @@ extension SourceKitLSPServer {
     if let diagnosticOptions = server.diagnosticProvider {
       await registry.registerDiagnosticIfNeeded(options: diagnosticOptions, for: languages, server: self)
     }
+    if let codeLensOptions = server.codeLensProvider {
+      await registry.registerCodeLensIfNeeded(options: codeLensOptions, for: languages, server: self)
+    }
     if let commandOptions = server.executeCommandProvider {
       await registry.registerExecuteCommandIfNeeded(commands: commandOptions.commands, server: self)
     }
@@ -1642,6 +1648,14 @@ extension SourceKitLSPServer {
     return req.injectMetadata(toResponse: response)
   }
 
+  func codeLens(
+    _ req: CodeLensRequest,
+    workspace: Workspace,
+    languageService: LanguageService
+  ) async throws -> [CodeLens] {
+    return try await languageService.codeLens(req)
+  }
+
   func inlayHint(
     _ req: InlayHintRequest,
     workspace: Workspace,

Sources/SourceKitLSP/Swift/SwiftCodeLensScanner.swift

@@ -0,0 +1,77 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import LanguageServerProtocol
+import SwiftSyntax
+
+/// Scans a source file for classes or structs annotated with `@main` and returns a code lens for them.
+final class SwiftCodeLensScanner: SyntaxVisitor {
+  /// The document snapshot of the syntax tree that is being walked.
+  private var snapshot: DocumentSnapshot
+
+  /// The collection of CodeLenses found in the document.
+  private var result: [CodeLens] = []
+
+  private init(snapshot: DocumentSnapshot) {
+    self.snapshot = snapshot
+    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 static func findCodeLenses(
+    in snapshot: DocumentSnapshot,
+    syntaxTreeManager: SyntaxTreeManager
+  ) async -> [CodeLens] {
+    guard snapshot.text.contains("@main") else {
+      // This is intended to filter out files that obviously do not contain an entry point.
+      return []
+    }
+
+    let syntaxTree = await syntaxTreeManager.syntaxTree(for: snapshot)
+    let visitor = SwiftCodeLensScanner(snapshot: snapshot)
+    visitor.walk(syntaxTree)
+    return visitor.result
+  }
+
+  override func visit(_ node: ClassDeclSyntax) -> SyntaxVisitorContinueKind {
+    node.attributes.forEach(self.captureLensFromAttribute)
+    return .skipChildren
+  }
+
+  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)
+
+      // 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", command: "swift.run", arguments: nil)
+        )
+      )
+
+      self.result.append(
+        CodeLens(
+          range: range,
+          command: Command(title: "Debug", command: "swift.debug", arguments: nil)
+        )
+      )
+    }
+  }
+}

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -921,6 +921,11 @@ extension SwiftLanguageService {
     return Array(hints)
   }
 
+  public func codeLens(_ req: CodeLensRequest) async throws -> [CodeLens] {
+    let snapshot = try documentManager.latestSnapshot(req.textDocument.uri)
+    return await SwiftCodeLensScanner.findCodeLenses(in: snapshot, syntaxTreeManager: self.syntaxTreeManager)
+  }
+
   public func documentDiagnostic(_ req: DocumentDiagnosticsRequest) async throws -> DocumentDiagnosticReport {
     do {
       await semanticIndexManager?.prepareFileForEditorFunctionality(req.textDocument.uri)

Tests/SourceKitLSPTests/CodeLensTests.swift

@@ -0,0 +1,70 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import LSPTestSupport
+import LanguageServerProtocol
+import SKTestSupport
+import XCTest
+
+final class CodeLensTests: XCTestCase {
+  func testNoLenses() async throws {
+    let project = try await SwiftPMTestProject(
+      files: [
+        "Test.swift": """
+        struct MyApp {
+          public static func main() {}
+        }
+        """
+      ]
+    )
+    let (uri, _) = try project.openDocument("Test.swift")
+
+    let response = try await project.testClient.send(
+      CodeLensRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+
+    XCTAssertEqual(response, [])
+  }
+
+  func testSuccessfulCodeLensRequest() async throws {
+    let project = try await SwiftPMTestProject(
+      files: [
+        "Test.swift": """
+        1️⃣@main2️⃣
+        struct MyApp {
+          public static func main() {}
+        }
+        """
+      ]
+    )
+
+    let (uri, positions) = try project.openDocument("Test.swift")
+
+    let response = try await project.testClient.send(
+      CodeLensRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+
+    XCTAssertEqual(
+      response,
+      [
+        CodeLens(
+          range: positions["1️⃣"]..<positions["2️⃣"],
+          command: Command(title: "Run", command: "swift.run", arguments: nil)
+        ),
+        CodeLens(
+          range: positions["1️⃣"]..<positions["2️⃣"],
+          command: Command(title: "Debug", command: "swift.debug", arguments: nil)
+        ),
+      ]
+    )
+  }
+}