Allow macro expansions to be viewed through GetReferenceDocumentRequest instead of storing in temporary files

Author: lokesh-tr

Documentation/LSP Extensions.md

@@ -473,3 +473,28 @@ export interface PeekDocumentsResult {
   success: boolean;
 }
 ```
+
+## `workspace/getReferenceDocument`
+
+Request from the client to the server asking for contents of a URI having a custom scheme.
+For example: "sourcekit-lsp:"
+
+- params: `GetReferenceDocumentParams`
+
+- result: `GetReferenceDocumentResponse`
+
+```ts
+export interface GetReferenceDocumentParams {
+  /**
+   * The `DocumentUri` of the custom scheme url for which content is required
+   */
+  uri: DocumentUri;
+}
+
+/**
+ * Response containing `content` of `GetReferenceDocumentRequest`
+ */
+export interface GetReferenceDocumentResult {
+  content: string;
+}
+```

Sources/LanguageServerProtocol/CMakeLists.txt

@@ -55,6 +55,7 @@ add_library(LanguageServerProtocol STATIC
   Requests/ExecuteCommandRequest.swift
   Requests/FoldingRangeRequest.swift
   Requests/FormattingRequests.swift
+  Requests/GetReferenceDocumentRequest.swift
   Requests/HoverRequest.swift
   Requests/ImplementationRequest.swift
   Requests/IndexedRenameRequest.swift

Sources/LanguageServerProtocol/Messages.swift

@@ -48,6 +48,7 @@ public let builtinRequests: [_RequestType.Type] = [
   DocumentTestsRequest.self,
   ExecuteCommandRequest.self,
   FoldingRangeRequest.self,
+  GetReferenceDocumentRequest.self,
   HoverRequest.self,
   ImplementationRequest.self,
   InitializeRequest.self,

Sources/LanguageServerProtocol/Requests/GetReferenceDocumentRequest.swift

@@ -0,0 +1,42 @@
+//===----------------------------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+/// Request from the client to the server asking for contents of a URI having a custom scheme **(LSP Extension)**
+/// For example: "sourcekit-lsp:"
+///
+/// - Parameters:
+///   - uri: The `DocumentUri` of the custom scheme url for which content is required
+///
+/// - Returns: `GetReferenceDocumentResponse` which contains the `content` to be displayed.
+///
+/// ### LSP Extension
+///
+/// This request is an extension to LSP supported by SourceKit-LSP.
+public struct GetReferenceDocumentRequest: RequestType {
+  public static let method: String = "workspace/getReferenceDocument"
+  public typealias Response = GetReferenceDocumentResponse
+
+  public var uri: DocumentURI
+
+  public init(uri: DocumentURI) {
+    self.uri = uri
+  }
+}
+
+/// Response containing `content` of `GetReferenceDocumentRequest`
+public struct GetReferenceDocumentResponse: ResponseType {
+  public var content: String
+
+  public init(content: String) {
+    self.content = content
+  }
+}

Sources/SourceKitLSP/CMakeLists.txt

@@ -47,10 +47,12 @@ target_sources(SourceKitLSP PRIVATE
   Swift/ExpandMacroCommand.swift
   Swift/FoldingRange.swift
   Swift/MacroExpansion.swift
+  Swift/MacroExpansionReferenceDocumentURLData.swift
   Swift/OpenInterface.swift
   Swift/RefactoringResponse.swift
   Swift/RefactoringEdit.swift
   Swift/RefactorCommand.swift
+  Swift/ReferenceDocumentURL.swift
   Swift/RelatedIdentifiers.swift
   Swift/RewriteSourceKitPlaceholders.swift
   Swift/SemanticRefactorCommand.swift

Sources/SourceKitLSP/Clang/ClangLanguageService.swift

@@ -661,6 +661,10 @@ extension ClangLanguageService {
   func executeCommand(_ req: ExecuteCommandRequest) async throws -> LSPAny? {
     return try await forwardRequestToClangd(req)
   }
+
+  func getReferenceDocument(_ req: GetReferenceDocumentRequest) async throws -> GetReferenceDocumentResponse {
+    throw ResponseError.unknown("unsupported method")
+  }
 }
 
 /// Clang build settings derived from a `FileBuildSettingsChange`.

Sources/SourceKitLSP/LanguageService.swift

@@ -251,6 +251,8 @@ package protocol LanguageService: AnyObject, Sendable {
 
   func executeCommand(_ req: ExecuteCommandRequest) async throws -> LSPAny?
 
+  func getReferenceDocument(_ req: GetReferenceDocumentRequest) async throws -> GetReferenceDocumentResponse
+
   /// Perform a syntactic scan of the file at the given URI for test cases and test classes.
   ///
   /// This is used as a fallback to show the test cases in a file if the index for a given file is not up-to-date.

Sources/SourceKitLSP/Rename.swift

@@ -687,7 +687,7 @@ extension SourceKitLSPServer {
     guard let workspace = await workspaceForDocument(uri: uri) else {
       throw ResponseError.workspaceNotOpen(uri)
     }
-    guard let primaryFileLanguageService = workspace.documentService.value[uri] else {
+    guard let primaryFileLanguageService = workspace.documentService(for: uri) else {
       return nil
     }
 

Sources/SourceKitLSP/SourceKitLSPServer.swift

@@ -356,7 +356,7 @@ package actor SourceKitLSPServer {
 
     // This should be created as soon as we receive an open call, even if the document
     // isn't yet ready.
-    guard let languageService = workspace.documentService.value[doc] else {
+    guard let languageService = workspace.documentService(for: doc) else {
       return
     }
 
@@ -376,7 +376,7 @@ package actor SourceKitLSPServer {
       guard let workspace = await self.workspaceForDocument(uri: request.textDocument.uri) else {
         throw ResponseError.workspaceNotOpen(request.textDocument.uri)
       }
-      guard let languageService = workspace.documentService.value[doc] else {
+      guard let languageService = workspace.documentService(for: doc) else {
         throw ResponseError.unknown("No language service for '\(request.textDocument.uri)' found")
       }
       return try await requestHandler(request, workspace, languageService)
@@ -399,7 +399,7 @@ package actor SourceKitLSPServer {
       guard let workspace = await self.workspaceForDocument(uri: documentUri) else {
         continue
       }
-      guard workspace.documentService.value[documentUri] === languageService else {
+      guard workspace.documentService(for: documentUri) === languageService else {
         continue
       }
       guard let snapshot = try? self.documentManager.latestSnapshot(documentUri) else {
@@ -517,7 +517,7 @@ package actor SourceKitLSPServer {
     _ language: Language,
     in workspace: Workspace
   ) async -> LanguageService? {
-    if let service = workspace.documentService.value[uri] {
+    if let service = workspace.documentService(for: uri) {
       return service
     }
 
@@ -535,17 +535,7 @@ package actor SourceKitLSPServer {
       """
     )
 
-    return workspace.documentService.withLock { documentService in
-      if let concurrentlySetService = documentService[uri] {
-        // Since we await the construction of `service`, another call to this
-        // function might have happened and raced us, setting
-        // `workspace.documentServices[uri]`. If this is the case, return the
-        // existing value and discard the service that we just retrieved.
-        return concurrentlySetService
-      }
-      documentService[uri] = service
-      return service
-    }
+    return workspace.documentService(for: uri, newLanguageService: service)
   }
 }
 
@@ -732,6 +722,8 @@ extension SourceKitLSPServer: MessageHandler {
       await request.reply { try await executeCommand(request.params) }
     case let request as RequestAndReply<FoldingRangeRequest>:
       await self.handleRequest(for: request, requestHandler: self.foldingRange)
+    case let request as RequestAndReply<GetReferenceDocumentRequest>:
+      await request.reply { try await getReferenceDocument(request.params) }
     case let request as RequestAndReply<HoverRequest>:
       await self.handleRequest(for: request, requestHandler: self.hover)
     case let request as RequestAndReply<ImplementationRequest>:
@@ -803,7 +795,7 @@ extension SourceKitLSPServer: BuildSystemDelegate {
         continue
       }
 
-      guard let service = await self.workspaceForDocument(uri: uri)?.documentService.value[uri] else {
+      guard let service = await self.workspaceForDocument(uri: uri)?.documentService(for: uri) else {
         continue
       }
 
@@ -827,7 +819,7 @@ extension SourceKitLSPServer: BuildSystemDelegate {
       }
       for uri in self.affectedOpenDocumentsForChangeSet(changedFilesForWorkspace, self.documentManager) {
         logger.log("Dependencies updated for opened file \(uri.forLogging)")
-        if let service = workspace.documentService.value[uri] {
+        if let service = workspace.documentService(for: uri) {
           await service.documentDependenciesUpdated(uri)
         }
       }
@@ -962,6 +954,8 @@ 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`.
+    //
+    // Similarly, for "workspace/getReferenceDocument".
     var clientCapabilities = req.capabilities
     if case .dictionary(let initializationOptions) = req.initializationOptions {
       if let peekDocuments = initializationOptions["workspace/peekDocuments"] {
@@ -973,6 +967,15 @@ extension SourceKitLSPServer {
         }
       }
 
+      if let getReferenceDocument = initializationOptions["workspace/getReferenceDocument"] {
+        if case .dictionary(var experimentalCapabilities) = clientCapabilities.experimental {
+          experimentalCapabilities["workspace/getReferenceDocument"] = getReferenceDocument
+          clientCapabilities.experimental = .dictionary(experimentalCapabilities)
+        } else {
+          clientCapabilities.experimental = .dictionary(["workspace/getReferenceDocument": getReferenceDocument])
+        }
+      }
+
       // 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"],
@@ -1139,6 +1142,7 @@ extension SourceKitLSPServer {
         "workspace/tests": .dictionary(["version": .int(2)]),
         "textDocument/tests": .dictionary(["version": .int(2)]),
         "workspace/triggerReindex": .dictionary(["version": .int(1)]),
+        "workspace/getReferenceDocument": .dictionary(["version": .int(1)]),
       ])
     )
   }
@@ -1338,7 +1342,7 @@ extension SourceKitLSPServer {
       )
       return
     }
-    await workspace.documentService.value[uri]?.reopenDocument(notification)
+    await workspace.documentService(for: uri)?.reopenDocument(notification)
   }
 
   func closeDocument(_ notification: DidCloseTextDocumentNotification, workspace: Workspace) async {
@@ -1352,7 +1356,7 @@ extension SourceKitLSPServer {
 
     await workspace.buildSystemManager.unregisterForChangeNotifications(for: uri)
 
-    await workspace.documentService.value[uri]?.closeDocument(notification)
+    await workspace.documentService(for: uri)?.closeDocument(notification)
   }
 
   func changeDocument(_ notification: DidChangeTextDocumentNotification) async {
@@ -1378,7 +1382,7 @@ extension SourceKitLSPServer {
       // Already logged failure
       return
     }
-    await workspace.documentService.value[uri]?.changeDocument(
+    await workspace.documentService(for: uri)?.changeDocument(
       notification,
       preEditSnapshot: preEditSnapshot,
       postEditSnapshot: postEditSnapshot,
@@ -1641,7 +1645,7 @@ extension SourceKitLSPServer {
     guard let workspace = await workspaceForDocument(uri: uri) else {
       throw ResponseError.workspaceNotOpen(uri)
     }
-    guard let languageService = workspace.documentService.value[uri] else {
+    guard let languageService = workspace.documentService(for: uri) else {
       return nil
     }
 
@@ -1652,6 +1656,21 @@ extension SourceKitLSPServer {
     return try await languageService.executeCommand(executeCommand)
   }
 
+  func getReferenceDocument(_ req: GetReferenceDocumentRequest) async throws -> GetReferenceDocumentResponse {
+    let referenceDocumentURL = try ReferenceDocumentURL(from: req.uri)
+    let sourceFileURI = referenceDocumentURL.sourceDocument()
+
+    guard let workspace = await workspaceForDocument(uri: sourceFileURI) else {
+      throw ResponseError.workspaceNotOpen(sourceFileURI)
+    }
+
+    guard let languageService = workspace.documentService(for: sourceFileURI) else {
+      throw ResponseError.unknown("No Language Service for URI: \(sourceFileURI)")
+    }
+
+    return try await languageService.getReferenceDocument(req)
+  }
+
   func codeAction(
     _ req: CodeActionRequest,
     workspace: Workspace,