Add Run/Debug CodeLens Support

Sources/LanguageServerProtocol/CMakeLists.txt

@@ -125,6 +125,7 @@ add_library(LanguageServerProtocol STATIC
   SupportTypes/SemanticTokenTypes.swift
   SupportTypes/ServerCapabilities.swift
   SupportTypes/StringOrMarkupContent.swift
+  SupportTypes/SupportedCodeLensCommand.swift
   SupportTypes/SymbolKind.swift
   SupportTypes/TestItem.swift
   SupportTypes/TextDocumentContentChangeEvent.swift

Sources/LanguageServerProtocol/SupportTypes/ClientCapabilities.swift

@@ -465,6 +465,25 @@ public struct TextDocumentClientCapabilities: Hashable, Codable, Sendable {
     }
   }
 
+  public struct CodeLens: Hashable, Codable, Sendable {
+
+    /// Whether the client supports dynamic registration of this request.
+    public var dynamicRegistration: Bool?
+
+    /// Dictionary of supported commands announced by the client.
+    /// The key is the CodeLens name recognized by SourceKit-LSP and the
+    /// value is the command as recognized by the client.
+    public var supportedCommands: [SupportedCodeLensCommand: String]?
+
+    public init(
+      dynamicRegistration: Bool? = nil,
+      supportedCommands: [SupportedCodeLensCommand: String] = [:]
+    ) {
+      self.dynamicRegistration = dynamicRegistration
+      self.supportedCommands = supportedCommands
+    }
+  }
+
   /// Capabilities specific to `textDocument/rename`.
   public struct Rename: Hashable, Codable, Sendable {
 
@@ -666,7 +685,7 @@ public struct TextDocumentClientCapabilities: Hashable, Codable, Sendable {
 
   public var codeAction: CodeAction? = nil
 
-  public var codeLens: DynamicRegistrationCapability? = nil
+  public var codeLens: CodeLens? = nil
 
   public var documentLink: DynamicRegistrationCapability? = nil
 
@@ -715,7 +734,7 @@ public struct TextDocumentClientCapabilities: Hashable, Codable, Sendable {
     documentHighlight: DynamicRegistrationCapability? = nil,
     documentSymbol: DocumentSymbol? = nil,
     codeAction: CodeAction? = nil,
-    codeLens: DynamicRegistrationCapability? = nil,
+    codeLens: CodeLens? = nil,
     documentLink: DynamicRegistrationCapability? = nil,
     colorProvider: DynamicRegistrationCapability? = nil,
     formatting: DynamicRegistrationCapability? = nil,

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/LanguageServerProtocol/SupportTypes/SupportedCodeLensCommand.swift

@@ -0,0 +1,29 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2024 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
+//
+//===----------------------------------------------------------------------===//
+
+/// Code lenses that LSP can annotate code with.
+///
+/// Clients provide these as keys to the `supportedCommands` dictionary supplied
+/// in the client's `InitializeRequest`.
+public struct SupportedCodeLensCommand: Codable, Hashable, RawRepresentable, Sendable {
+  public var rawValue: String
+
+  public init(rawValue: String) {
+    self.rawValue = rawValue
+  }
+
+  /// Lens to run the application
+  public static let run: Self = Self(rawValue: "swift.run")
+
+  /// Lens to debug the application
+  public static let debug: Self = Self(rawValue: "swift.debug")
+}

Sources/SourceKitLSP/CMakeLists.txt

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

@@ -80,6 +80,10 @@ package final actor CapabilityRegistry {
     clientCapabilities.textDocument?.publishDiagnostics?.codeDescriptionSupport == true
   }
 
+  public var supportedCodeLensCommands: [SupportedCodeLensCommand: String] {
+    clientCapabilities.textDocument?.codeLens?.supportedCommands ?? [:]
+  }
+
   /// Since LSP 3.17.0, diagnostics can be reported through pull-based requests in addition to the existing push-based
   /// publish notifications.
   ///
@@ -279,6 +283,7 @@ package final actor CapabilityRegistry {
     server: SourceKitLSPServer
   ) async {
     guard clientHasDynamicInlayHintRegistration else { return }
+
     await registerLanguageSpecificCapability(
       options: InlayHintRegistrationOptions(
         documentSelector: DocumentSelector(for: languages),
@@ -345,7 +350,7 @@ package 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 @@ package 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

@@ -700,6 +700,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>:
@@ -961,14 +963,30 @@ extension SourceKitLSPServer {
     // The below is a workaround for the vscode-swift extension since it cannot set client capabilities.
     // It passes "workspace/peekDocuments" through the `initializationOptions`.
     var clientCapabilities = req.capabilities
-    if case .dictionary(let initializationOptions) = req.initializationOptions,
-      let peekDocuments = initializationOptions["workspace/peekDocuments"]
-    {
-      if case .dictionary(var experimentalCapabilities) = clientCapabilities.experimental {
-        experimentalCapabilities["workspace/peekDocuments"] = peekDocuments
-        clientCapabilities.experimental = .dictionary(experimentalCapabilities)
-      } else {
-        clientCapabilities.experimental = .dictionary(["workspace/peekDocuments": peekDocuments])
+    if case .dictionary(let initializationOptions) = req.initializationOptions {
+      if let peekDocuments = initializationOptions["workspace/peekDocuments"] {
+        if case .dictionary(var experimentalCapabilities) = clientCapabilities.experimental {
+          experimentalCapabilities["workspace/peekDocuments"] = peekDocuments
+          clientCapabilities.experimental = .dictionary(experimentalCapabilities)
+        } else {
+          clientCapabilities.experimental = .dictionary(["workspace/peekDocuments": peekDocuments])
+        }
+      }
+
+      // The client announces what CodeLenses it supports, and the LSP will only return
+      // ones found in the supportedCommands dictionary.
+      if let codeLens = initializationOptions["textDocument/codeLens"],
+        case let .dictionary(codeLensConfig) = codeLens,
+        case let .dictionary(supportedCommands) = codeLensConfig["supportedCommands"]
+      {
+        let commandMap = supportedCommands.compactMap { (key, value) in
+          if case let .string(clientCommand) = value {
+            return (SupportedCodeLensCommand(rawValue: key), clientCommand)
+          }
+          return nil
+        }
+
+        clientCapabilities.textDocument?.codeLens?.supportedCommands = Dictionary(uniqueKeysWithValues: commandMap)
       }
     }
 
@@ -1100,6 +1118,7 @@ extension SourceKitLSPServer {
           supportsCodeActions: true
         )
       ),
+      codeLensProvider: CodeLensOptions(),
       documentFormattingProvider: .value(DocumentFormattingOptions(workDoneProgress: false)),
       renameProvider: .value(RenameOptions(prepareProvider: true)),
       colorProvider: .bool(true),
@@ -1642,6 +1661,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,86 @@
+//===----------------------------------------------------------------------===//
+//
+// 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 let snapshot: DocumentSnapshot
+
+  /// The collection of CodeLenses found in the document.
+  private var result: [CodeLens] = []
+
+  /// The map of supported commands and their client side command names
+  private let supportedCommands: [SupportedCodeLensCommand: String]
+
+  private init(snapshot: DocumentSnapshot, supportedCommands: [SupportedCodeLensCommand: String]) {
+    self.snapshot = snapshot
+    self.supportedCommands = supportedCommands
+    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,
+    supportedCommands: [SupportedCodeLensCommand: String]
+  ) async -> [CodeLens] {
+    guard snapshot.text.contains("@main") && !supportedCommands.isEmpty 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, supportedCommands: supportedCommands)
+    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)
+
+      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", command: runCommand, arguments: nil)
+          )
+        )
+      }
+
+      if let debugCommand = supportedCommands[SupportedCodeLensCommand.debug] {
+        self.result.append(
+          CodeLens(
+            range: range,
+            command: Command(title: "Debug", command: debugCommand, arguments: nil)
+          )
+        )
+      }
+    }
+  }
+}

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

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