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?
 
@@ -314,6 +317,26 @@ public final actor CapabilityRegistry {
     )
   }
 
+  public func registerCodeLensIfNeeded(
+    options: CodeLensOptions,
+    for languages: [Language],
+    server: SourceKitLSPServer
+  ) async {
+    guard clientHasDynamicDocumentDiagnosticsRegistration else { return }
+
+    await registerLanguageSpecificCapability(
+      options: CodeLensRegistrationOptions(
+        documentSelector: DocumentSelector(for: languages),
+        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(

Sources/SourceKitLSP/Clang/ClangLanguageService.swift

@@ -616,6 +616,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

@@ -191,6 +191,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

@@ -743,6 +743,8 @@ extension SourceKitLSPServer: MessageHandler {
       initialized = true
     case let request as RequestAndReply<InlayHintRequest>:
       await self.handleRequest(for: request, requestHandler: self.inlayHint)
+    case let request as RequestAndReply<CodeLensRequest>:
+      await self.handleRequest(for: request, requestHandler: self.codeLens)
     case let request as RequestAndReply<PollIndexRequest>:
       await request.reply { try await pollIndex(request.params) }
     case let request as RequestAndReply<PrepareRenameRequest>:
@@ -1161,6 +1163,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)
     }
@@ -1622,6 +1627,14 @@ extension SourceKitLSPServer {
     return try await languageService.inlayHint(req)
   }
 
+  func codeLens(
+    _ req: CodeLensRequest,
+    workspace: Workspace,
+    languageService: LanguageService
+  ) async throws -> [CodeLens] {
+    return try await languageService.codeLens(req)
+  }
+
   func documentDiagnostic(
     _ req: DocumentDiagnosticsRequest,
     workspace: Workspace,

Sources/SourceKitLSP/Swift/SwiftCodeLensScanner.swift

@@ -0,0 +1,78 @@
+//===----------------------------------------------------------------------===//
+//
+// 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(attr: AttributeListSyntax.Element) {
+    if attr.trimmedDescription == "@main" {
+      let range = self.snapshot.absolutePositionRange(of: attr.trimmedRange)
+      // TODO: Only provide these code lenses for `.executable` targets.
+      // It doesn't makes sense to provide them for library/plugin targets, even if they have a `@main` attribute.
+      // Until this is addressed `CodeLensTests.testCodeLensRequestWithInvalidProduct` will fail.
+      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

@@ -312,6 +312,7 @@ extension SwiftLanguageService {
             supportsCodeActions: true
           )
         ),
+        codeLensProvider: CodeLensOptions(resolveProvider: true),
         colorProvider: .bool(true),
         foldingRangeProvider: .bool(true),
         executeCommandProvider: ExecuteCommandOptions(
@@ -922,6 +923,12 @@ extension SwiftLanguageService {
     return Array(hints)
   }
 
+  public func codeLens(_ req: CodeLensRequest) async throws -> [CodeLens] {
+    let snapshot = try documentManager.latestSnapshot(req.textDocument.uri)
+    let syntaxTreeManager = SyntaxTreeManager()
+    return await SwiftCodeLensScanner.findCodeLenses(in: snapshot, syntaxTreeManager: syntaxTreeManager)
+  }
+
   public func documentDiagnostic(_ req: DocumentDiagnosticsRequest) async throws -> DocumentDiagnosticReport {
     do {
       await semanticIndexManager?.prepareFileForEditorFunctionality(req.textDocument.uri)

Tests/SourceKitLSPTests/CodeLensTests.swift

@@ -0,0 +1,115 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 SourceKitLSP
+import XCTest
+
+final class CodeLensTests: XCTestCase {
+  // MARK: - Helpers
+
+  func performCodeLensRequest(text: String, range: Range<Position>? = nil) async throws -> [CodeLens]? {
+    let testClient = try await TestSourceKitLSPClient()
+    let uri = DocumentURI(for: .swift)
+
+    testClient.openDocument(text, uri: uri)
+
+    let request = CodeLensRequest(textDocument: TextDocumentIdentifier(uri))
+    return try await testClient.send(request)
+  }
+
+  // MARK: - Tests
+
+  func testEmpty() async throws {
+    let text = ""
+    let lenses = try await performCodeLensRequest(text: text)
+    XCTAssertEqual(lenses, [])
+  }
+
+  func testSuccessfulCodeLensRequest() async throws {
+    let packageManifest = """
+      let package = Package(
+        name: "MyExecutable",
+        products: [.executable(name: "MyExecutable")],
+        targets: [.target(name: "MyExecutable")]
+      )
+      """
+
+    let project = try await SwiftPMTestProject(
+      files: [:],
+      manifest: packageManifest
+    )
+    let uri = DocumentURI(for: .swift)
+
+    let positions = project.testClient.openDocument(
+      """
+      1️⃣@main2️⃣
+      struct MyApp {
+        public static func main() {}
+      }
+      """,
+      uri: uri
+    )
+
+    let response = try await project.testClient.send(
+      CodeLensRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+
+    XCTAssertEqual(
+      try XCTUnwrap(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)
+        ),
+      ]
+    )
+  }
+
+  func testCodeLensRequestWithInvalidProduct() async throws {
+    let packageManifest = """
+      let package = Package(
+        name: "MyLibrary",
+        products: [.library(name: "MyLibrary")],
+        targets: [.target(name: "MyLibrary")]
+      )
+      """
+
+    let project = try await SwiftPMTestProject(
+      files: [:],
+      manifest: packageManifest
+    )
+    let uri = DocumentURI(for: .swift)
+
+    project.testClient.openDocument(
+      """
+      1️⃣@main2️⃣
+      struct MyApp {
+        public static func main() {}
+      }
+      """,
+      uri: uri
+    )
+
+    let response = try await project.testClient.send(
+      CodeLensRequest(textDocument: TextDocumentIdentifier(uri))
+    )
+
+    XCTAssertEqual(try XCTUnwrap(response), [])
+  }
+}