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: "swift-macro-expansion://"
+
+- params: `GetReferenceDocumentParams`
+
+- result: `GetReferenceDocumentResponse`
+
+```ts
+export interface GetReferenceDocumentParams {
+    /**
+     * The `DocumentUri` of the custom scheme url for which content is required
+     */
+    uri: langclient.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: "swift-macro-expansion://"
+///
+/// - 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/Clang/ClangLanguageService.swift

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

@@ -250,6 +250,8 @@ public 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/SourceKitLSPServer.swift

@@ -731,6 +731,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>:
@@ -1633,6 +1635,25 @@ extension SourceKitLSPServer {
     return try await languageService.executeCommand(executeCommand)
   }
 
+  func getReferenceDocument(_ req: GetReferenceDocumentRequest) async throws -> GetReferenceDocumentResponse {
+    guard let givenURL = URL(string: req.uri.stringValue) else {
+      throw ResponseError.unknown("Invalid URL")
+    }
+
+    let filePath = "file:/" + givenURL.pathComponents.dropLast().joined(separator: "/")
+    let sourceFileURI = try DocumentURI(string: filePath)
+
+    guard let workspace = await workspaceForDocument(uri: sourceFileURI) else {
+      throw ResponseError.workspaceNotOpen(sourceFileURI)
+    }
+
+    guard let languageService = workspace.documentService.value[sourceFileURI] else {
+      throw ResponseError.unknown("No Language Service for URI: \(sourceFileURI)")
+    }
+
+    return try await languageService.getReferenceDocument(req)
+  }
+
   func codeAction(
     _ req: CodeActionRequest,
     workspace: Workspace,

Sources/SourceKitLSP/Swift/MacroExpansion.swift

@@ -74,50 +74,38 @@ extension SwiftLanguageService {
     var macroExpansionFilePaths: [URL] = []
     for macroEdit in expansion.edits {
       if let bufferName = macroEdit.bufferName {
-        // buffer name without ".swift"
-        let macroExpansionBufferDirectoryName =
-          bufferName.hasSuffix(".swift")
-          ? String(bufferName.dropLast(6))
-          : bufferName
-
-        let macroExpansionBufferDirectoryURL = self.generatedMacroExpansionsPath
-          .appendingPathComponent(macroExpansionBufferDirectoryName)
-
-        completeExpansionDirectoryName += "\(bufferName)-"
-
-        do {
-          try FileManager.default.createDirectory(
-            at: macroExpansionBufferDirectoryURL,
-            withIntermediateDirectories: true
-          )
-        } catch {
-          throw ResponseError.unknown(
-            "Failed to create directory for macro expansion buffer at path: \(macroExpansionBufferDirectoryURL.path)"
-          )
+        // Custom Scheme for Macro Expansions:
+        // `swift-macro-expansion://[path_to_source_file]/LaCb-LcCd.swift?lowerL=&lowerC=&upperL=&upperC=&bufferName=`
+        //  ---------------------   --------------------- --------------- ------------------------------- -----------
+        //           1                       2                3                   4                     5
+        //
+        // 1. Scheme
+        // 2. The path to the source file
+        // 3. Position range of macro after expansion (used for display purpose in VS Code)
+        // 4. Selection range of the cursor
+        // 5. Buffer Name of the required Macro Expansion
+
+        guard var components = URLComponents(string: sourceFileURL.absoluteString) else {
+          throw ResponseError.unknown("Invalid URL")
         }
-
-        // name of the source file
-        let macroExpansionFileName = sourceFileURL.deletingPathExtension().lastPathComponent
-
-        // github permalink notation for position range
-        let macroExpansionPositionRangeIndicator =
-          "L\(macroEdit.range.lowerBound.line + 1)C\(macroEdit.range.lowerBound.utf16index + 1)-L\(macroEdit.range.upperBound.line + 1)C\(macroEdit.range.upperBound.utf16index + 1)"
-
-        let macroExpansionFilePath =
-          macroExpansionBufferDirectoryURL
-          .appendingPathComponent(
-            "\(macroExpansionFileName)_\(macroExpansionPositionRangeIndicator).\(sourceFileURL.pathExtension)"
-          )
-
-        do {
-          try macroEdit.newText.write(to: macroExpansionFilePath, atomically: true, encoding: .utf8)
-        } catch {
-          throw ResponseError.unknown(
-            "Unable to write macro expansion to file path: \"\(macroExpansionFilePath.path)\""
-          )
+        components.scheme = "swift-macro-expansion"
+        components.path +=
+          "/L\(macroEdit.range.lowerBound.line + 1)C\(macroEdit.range.lowerBound.utf16index + 1)-L\(macroEdit.range.upperBound.line + 1)C\(macroEdit.range.upperBound.utf16index + 1).swift"
+        let queryItems = [
+          URLQueryItem(name: "lowerL", value: String(expandMacroCommand.positionRange.lowerBound.line)),
+          URLQueryItem(name: "lowerC", value: String(expandMacroCommand.positionRange.lowerBound.utf16index)),
+          URLQueryItem(name: "upperL", value: String(expandMacroCommand.positionRange.upperBound.line)),
+          URLQueryItem(name: "upperC", value: String(expandMacroCommand.positionRange.upperBound.utf16index)),
+          URLQueryItem(name: "bufferName", value: bufferName),
+        ]
+        components.queryItems = queryItems
+
+        guard let url = components.url else {
+          throw ResponseError.unknown("Failed to construct URL")
         }
+        macroExpansionFilePaths.append(url)
 
-        macroExpansionFilePaths.append(macroExpansionFilePath)
+        completeExpansionDirectoryName += "\(bufferName)-"
 
         let editContent =
           """
@@ -208,4 +196,29 @@ extension SwiftLanguageService {
       }
     }
   }
+
+  func expandMacro(uri: DocumentURI, selectionRange: Range<Position>, bufferName: String) async throws -> String {
+    guard let sourceKitLSPServer else {
+      // `SourceKitLSPServer` has been destructed. We are tearing down the
+      // language server. Nothing left to do.
+      throw ResponseError.unknown("Connection to the editor closed")
+    }
+
+    let expandMacroCommand = ExpandMacroCommand(
+      positionRange: selectionRange,
+      textDocument: TextDocumentIdentifier(uri)
+    )
+
+    let expansion = try await self.refactoring(expandMacroCommand)
+
+    let macroExpansionEdit = expansion.edits.filter {
+      $0.bufferName == bufferName
+    }.only
+
+    guard let content = macroExpansionEdit?.newText else {
+      throw ResponseError.unknown("Macro Expansion Edit doesn't Exist")
+    }
+
+    return content
+  }
 }

Sources/SourceKitLSP/Swift/SwiftLanguageService.swift

@@ -968,6 +968,48 @@ extension SwiftLanguageService {
 
     return nil
   }
+
+  public func getReferenceDocument(_ req: GetReferenceDocumentRequest) async throws -> GetReferenceDocumentResponse {
+    switch req.uri.scheme {
+    case "swift-macro-expansion":
+      guard let givenURL = URL(string: req.uri.stringValue) else {
+        throw ResponseError.unknown("Unable to parse given URL")
+      }
+
+      let filePath = "file:/" + givenURL.pathComponents.dropLast().joined(separator: "/")
+
+      guard let sourceFileURL = try DocumentURI(string: filePath).fileURL else {
+        throw ResponseError.unknown("Unable to obtain sourceFileURL")
+      }
+
+      guard let components = URLComponents(string: givenURL.absoluteString) else {
+        throw ResponseError.unknown("Unable to obtain components of url")
+      }
+
+      guard let lowerL = Int(components.queryItems?.first { $0.name == "lowerL" }?.value ?? ""),
+        let lowerC = Int(components.queryItems?.first { $0.name == "lowerC" }?.value ?? ""),
+        let upperL = Int(components.queryItems?.first { $0.name == "upperL" }?.value ?? ""),
+        let upperC = Int(components.queryItems?.first { $0.name == "upperC" }?.value ?? ""),
+        let bufferName = components.queryItems?.first { $0.name == "bufferName" }?.value
+      else {
+        throw ResponseError.unknown("Invalid query in URL")
+      }
+
+      let selectionRange = Position(line: lowerL, utf16index: lowerC)..<Position(line: upperL, utf16index: upperC)
+
+      let content = try await expandMacro(
+        uri: DocumentURI(sourceFileURL),
+        selectionRange: selectionRange,
+        bufferName: bufferName
+      )
+
+      return GetReferenceDocumentResponse(content: content)
+    default:
+      throw ResponseError.unknown("Invalid Scheme")
+    }
+
+    throw ResponseError.unknown("No Content to show")
+  }
 }
 
 extension SwiftLanguageService: SKDNotificationHandler {

Tests/SourceKitLSPTests/ExecuteCommandTests.swift

@@ -244,27 +244,39 @@ final class ExecuteCommandTests: XCTestCase {
 
           try await fulfillmentOfOrThrow([expectation])
 
-          let urls = try XCTUnwrap(
+          let urlStrings = try XCTUnwrap(
             peekDocumentsRequestURIs.value?.map {
               return try XCTUnwrap(
-                $0.fileURL,
+                $0.stringValue,
                 "Failed for position range between \(positionMarker.start) and \(positionMarker.end)"
               )
             },
             "Failed for position range between \(positionMarker.start) and \(positionMarker.end)"
           )
 
-          let filesContents = try urls.map { try String(contentsOf: $0, encoding: .utf8) }
+          var filesContents = [String]()
+          for url in urlStrings {
+            let result = try await project.testClient.send(GetReferenceDocumentRequest(uri: DocumentURI(string: url)))
+
+            filesContents.append(result.content)
+          }
 
           XCTAssertEqual(
             filesContents.only,
             "(1 + 2, \"1 + 2\")",
             "File doesn't contain macro expansion. Failed for position range between \(positionMarker.start) and \(positionMarker.end)"
           )
 
+          let urls = try urlStrings.map {
+            try XCTUnwrap(
+              URL(string: $0),
+              "Failed for position range between \(positionMarker.start) and \(positionMarker.end)"
+            )
+          }
+
           XCTAssertEqual(
             urls.only?.lastPathComponent,
-            "MyMacroClient_L5C3-L5C20.swift",
+            "L5C3-L5C20.swift",
             "Failed for position range between \(positionMarker.start) and \(positionMarker.end)"
           )
         } else {
@@ -417,17 +429,22 @@ final class ExecuteCommandTests: XCTestCase {
 
           try await fulfillmentOfOrThrow([expectation])
 
-          let urls = try XCTUnwrap(
+          let urlStrings = try XCTUnwrap(
             peekDocumentsRequestURIs.value?.map {
               return try XCTUnwrap(
-                $0.fileURL,
+                $0.stringValue,
                 "Failed for position range between \(positionMarker.start) and \(positionMarker.end)"
               )
             },
             "Failed for position range between \(positionMarker.start) and \(positionMarker.end)"
           )
 
-          let filesContents = try urls.map { try String(contentsOf: $0, encoding: .utf8) }
+          var filesContents = [String]()
+          for url in urlStrings {
+            let result = try await project.testClient.send(GetReferenceDocumentRequest(uri: DocumentURI(string: url)))
+
+            filesContents.append(result.content)
+          }
 
           XCTAssertEqual(
             filesContents,
@@ -439,12 +456,19 @@ final class ExecuteCommandTests: XCTestCase {
             "Files doesn't contain correct macro expansion. Failed for position range between \(positionMarker.start) and \(positionMarker.end)"
           )
 
+          let urls = try urlStrings.map {
+            try XCTUnwrap(
+              URL(string: $0),
+              "Failed for position range between \(positionMarker.start) and \(positionMarker.end)"
+            )
+          }
+
           XCTAssertEqual(
             urls.map { $0.lastPathComponent },
             [
-              "MyMacroClient_L7C3-L7C3.swift",
-              "MyMacroClient_L8C3-L8C3.swift",
-              "MyMacroClient_L9C1-L9C1.swift",
+              "L7C3-L7C3.swift",
+              "L8C3-L8C3.swift",
+              "L9C1-L9C1.swift",
             ],
             "Failed for position range between \(positionMarker.start) and \(positionMarker.end)"
           )