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

Contributor Documentation/LSP Extensions.md

@@ -473,3 +473,30 @@ 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:"
+
+Enable the experimental client capability `"workspace/getReferenceDocument"` so that the server responds with reference document URLs for certain requests or commands whenever possible.
+
+- 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,44 @@
+//===----------------------------------------------------------------------===//
+//
+// 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.
+/// Enable the experimental client capability `"workspace/getReferenceDocument"` so that the server responds with
+/// reference document URLs for certain requests or commands whenever possible.
+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/LanguageServerProtocol/SupportTypes/DocumentURI.swift

@@ -36,6 +36,10 @@ public struct DocumentURI: Codable, Hashable, Sendable {
     }
   }
 
+  /// The URL representation of the URI. Note that this URL can have an arbitrary scheme and might
+  /// not represent a file URL.
+  public var arbitrarySchemeURL: URL { storage }
+
   /// The document's URL scheme, if present.
   public var scheme: String? {
     return storage.scheme

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

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

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

@@ -357,7 +357,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
     }
 
@@ -377,7 +377,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)
@@ -400,7 +400,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 {
@@ -518,7 +518,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
     }
 
@@ -536,17 +536,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.setDocumentService(for: uri, service)
   }
 }
 
@@ -733,6 +723,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>:
@@ -804,7 +796,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
       }
 
@@ -828,7 +820,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)
         }
       }
@@ -964,6 +956,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"] {
@@ -975,6 +969,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"],
@@ -1143,6 +1146,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)]),
       ])
     )
   }
@@ -1342,7 +1346,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 {
@@ -1356,7 +1360,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 {
@@ -1382,7 +1386,7 @@ extension SourceKitLSPServer {
       // Already logged failure
       return
     }
-    await workspace.documentService.value[uri]?.changeDocument(
+    await workspace.documentService(for: uri)?.changeDocument(
       notification,
       preEditSnapshot: preEditSnapshot,
       postEditSnapshot: postEditSnapshot,
@@ -1645,7 +1649,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
     }
 
@@ -1656,6 +1660,21 @@ extension SourceKitLSPServer {
     return try await languageService.executeCommand(executeCommand)
   }
 
+  func getReferenceDocument(_ req: GetReferenceDocumentRequest) async throws -> GetReferenceDocumentResponse {
+    let referenceDocumentURL = try ReferenceDocumentURL(from: req.uri)
+    let primaryFileURI = referenceDocumentURL.primaryFile
+
+    guard let workspace = await workspaceForDocument(uri: primaryFileURI) else {
+      throw ResponseError.workspaceNotOpen(primaryFileURI)
+    }
+
+    guard let languageService = workspace.documentService(for: primaryFileURI) else {
+      throw ResponseError.unknown("No Language Service for URI: \(primaryFileURI)")
+    }
+
+    return try await languageService.getReferenceDocument(req)
+  }
+
   func codeAction(
     _ req: CodeActionRequest,
     workspace: Workspace,