Retrieve macro expansions for PeekDocumentsRequest through GetReferenceDocumentRequest instead of showing temporary files

Author: lokesh-tr

src/sourcekit-lsp/LanguageClientManager.ts

@@ -28,6 +28,7 @@ import { DiagnosticsManager } from "../DiagnosticsManager";
 import { LSPLogger, LSPOutputChannel } from "./LSPOutputChannel";
 import { SwiftOutputChannel } from "../ui/SwiftOutputChannel";
 import { promptForDiagnostics } from "../commands/captureDiagnostics";
+import { activateGetReferenceDocument } from "./getReferenceDocument";
 
 interface SourceKitLogMessageParams extends langclient.LogMessageParams {
     logName?: string;
@@ -110,6 +111,7 @@ export class LanguageClientManager {
     private cancellationToken?: vscode.CancellationTokenSource;
     private legacyInlayHints?: vscode.Disposable;
     private peekDocuments?: vscode.Disposable;
+    private getReferenceDocument?: vscode.Disposable;
     private restartedPromise?: Promise<void>;
     private currentWorkspaceFolder?: vscode.Uri;
     private waitingOnRestartCount: number;
@@ -247,6 +249,7 @@ export class LanguageClientManager {
         this.cancellationToken?.dispose();
         this.legacyInlayHints?.dispose();
         this.peekDocuments?.dispose();
+        this.getReferenceDocument?.dispose();
         this.subscriptions.forEach(item => item.dispose());
         this.languageClient?.stop();
         this.namedOutputChannels.forEach(channel => channel.dispose());
@@ -397,6 +400,8 @@ export class LanguageClientManager {
         this.legacyInlayHints = undefined;
         this.peekDocuments?.dispose();
         this.peekDocuments = undefined;
+        this.getReferenceDocument?.dispose();
+        this.getReferenceDocument = undefined;
         if (client) {
             this.cancellationToken?.cancel();
             this.cancellationToken?.dispose();
@@ -573,6 +578,7 @@ export class LanguageClientManager {
             initializationFailedHandler: () => false,
             initializationOptions: {
                 "workspace/peekDocuments": true, // workaround for client capability to handle `PeekDocumentsRequest`
+                "workspace/getReferenceDocument": true, // the client can handle URIs with scheme `sourcekit-lsp:`
                 "textDocument/codeLens": {
                     supportedCommands: {
                         "swift.run": "swift.run",
@@ -636,6 +642,8 @@ export class LanguageClientManager {
                 }
 
                 this.peekDocuments = activatePeekDocuments(client);
+                this.getReferenceDocument = activateGetReferenceDocument(client);
+                this.workspaceContext.subscriptions.push(this.getReferenceDocument);
             })
             .catch(reason => {
                 this.workspaceContext.outputChannel.log(`${reason}`);

src/sourcekit-lsp/getReferenceDocument.ts

@@ -0,0 +1,40 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the VS Code Swift open source project
+//
+// Copyright (c) 2024 the VS Code Swift project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of VS Code Swift project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import * as vscode from "vscode";
+import * as langclient from "vscode-languageclient/node";
+import { GetReferenceDocumentParams, GetReferenceDocumentRequest } from "./lspExtensions";
+
+export function activateGetReferenceDocument(client: langclient.LanguageClient): vscode.Disposable {
+    const getReferenceDocument = vscode.workspace.registerTextDocumentContentProvider(
+        "sourcekit-lsp",
+        {
+            provideTextDocumentContent: async (uri, token) => {
+                const params: GetReferenceDocumentParams = {
+                    uri: uri.toString(true),
+                };
+
+                const result = await client.sendRequest(GetReferenceDocumentRequest, params, token);
+
+                if (result) {
+                    return result.content;
+                } else {
+                    return "Unable to retrieve reference document";
+                }
+            },
+        }
+    );
+
+    return getReferenceDocument;
+}

src/sourcekit-lsp/lspExtensions.ts

@@ -56,6 +56,31 @@ export const PeekDocumentsRequest = new langclient.RequestType<
     unknown
 >("workspace/peekDocuments");
 
+// Get Reference Document
+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;
+}
+
+/**
+ * Request from the client to the server asking for contents of a URI having a custom scheme
+ * For example: "sourcekit-lsp:"
+ */
+export const GetReferenceDocumentRequest = new langclient.RequestType<
+    GetReferenceDocumentParams,
+    GetReferenceDocumentResult,
+    unknown
+>("workspace/getReferenceDocument");
+
 // Inlay Hints (pre Swift 5.6)
 export interface LegacyInlayHintsParams {
     /**

src/sourcekit-lsp/peekDocuments.ts

@@ -21,11 +21,9 @@ export function activatePeekDocuments(client: langclient.LanguageClient): vscode
         PeekDocumentsRequest.method,
         async (params: PeekDocumentsParams) => {
             const locations = params.locations.map(uri => {
+                const url = new URL(uri);
                 const location = new vscode.Location(
-                    vscode.Uri.from({
-                        scheme: "file",
-                        path: new URL(uri).pathname,
-                    }),
+                    vscode.Uri.parse(url.href, true),
                     new vscode.Position(0, 0)
                 );