Implement multi-file doc highlights

_extension/package.json

@@ -12,8 +12,11 @@
         "type": "git",
         "url": "https://github.com/microsoft/typescript-go"
     },
+    "enabledApiProposals": [
+        "multiDocumentHighlightProvider"
+    ],
     "engines": {
-        "vscode": "^1.106.0"
+        "vscode": "^1.110.0"
     },
     "capabilities": {
         "untrustedWorkspaces": {
@@ -194,10 +197,11 @@
     },
     "dependencies": {
         "@vscode/extension-telemetry": "^1.5.1",
-        "vscode-languageclient": "^10.0.0-next.21"
+        "vscode-languageclient": "10.0.0-next.21",
+        "vscode-languageserver-protocol": "3.17.6-next.17"
     },
     "devDependencies": {
-        "@types/vscode": "~1.106.1",
+        "@types/vscode": "~1.110.0",
         "@vscode/vsce": "^3.7.1",
         "esbuild": "^0.27.4"
     }

_extension/src/client.ts

@@ -20,6 +20,7 @@ import {
     configurationMiddleware,
     sendNotificationMiddleware,
 } from "./configurationMiddleware";
+import { registerMultiDocumentHighlightFeature } from "./languageFeatures/documentHighlight";
 import { registerSourceDefinitionFeature } from "./languageFeatures/sourceDefinition";
 import { registerTagClosingFeature } from "./languageFeatures/tagClosing";
 import * as tr from "./telemetryReporting";
@@ -207,6 +208,7 @@ export class Client implements vscode.Disposable {
 
         this.disposables.push(
             serverTelemetryListener,
+            registerMultiDocumentHighlightFeature(this.documentSelector, this.client),
             registerSourceDefinitionFeature(this.client),
             registerTagClosingFeature("typescript", this.documentSelector, this.client),
             registerTagClosingFeature("javascript", this.documentSelector, this.client),

_extension/src/languageFeatures/documentHighlight.ts

@@ -0,0 +1,78 @@
+import * as vscode from "vscode";
+import { LanguageClient } from "vscode-languageclient/node";
+import type { DocumentHighlight } from "vscode-languageserver-protocol";
+
+const multiDocumentHighlightMethod = "custom/textDocument/multiDocumentHighlight";
+
+interface MultiDocumentHighlightParams {
+    textDocument: { uri: string; };
+    position: { line: number; character: number; };
+    filesToSearch: string[];
+}
+
+interface MultiDocumentHighlightItem {
+    uri: string;
+    highlights: DocumentHighlight[];
+}
+
+class MultiDocumentHighlightProvider implements vscode.MultiDocumentHighlightProvider {
+    constructor(private readonly client: LanguageClient) {}
+
+    async provideMultiDocumentHighlights(
+        document: vscode.TextDocument,
+        position: vscode.Position,
+        otherDocuments: vscode.TextDocument[],
+        token: vscode.CancellationToken,
+    ): Promise<vscode.MultiDocumentHighlight[]> {
+        const allFiles = [document, ...otherDocuments]
+            .map(doc => this.client.code2ProtocolConverter.asUri(doc.uri))
+            .filter(file => !!file);
+
+        if (allFiles.length === 0) {
+            return [];
+        }
+
+        const params: MultiDocumentHighlightParams = {
+            textDocument: this.client.code2ProtocolConverter.asTextDocumentIdentifier(document),
+            position: this.client.code2ProtocolConverter.asPosition(position),
+            filesToSearch: allFiles,
+        };
+
+        let response: MultiDocumentHighlightItem[] | null;
+        try {
+            response = await this.client.sendRequest<MultiDocumentHighlightItem[] | null>(multiDocumentHighlightMethod, params, token);
+        }
+        catch (error) {
+            return [];
+        }
+
+        if (!response || token.isCancellationRequested) {
+            return [];
+        }
+
+        // MultiDocumentHighlight is proposed API; guard against missing or changed constructor.
+        try {
+            return response.map(item =>
+                new vscode.MultiDocumentHighlight(
+                    vscode.Uri.parse(item.uri),
+                    item.highlights.map(h => this.client.protocol2CodeConverter.asDocumentHighlight(h)),
+                )
+            );
+        }
+        catch {
+            return [];
+        }
+    }
+}
+
+export function registerMultiDocumentHighlightFeature(
+    selector: vscode.DocumentSelector,
+    client: LanguageClient,
+): vscode.Disposable {
+    const capabilities = client.initializeResult?.capabilities as { customMultiDocumentHighlightProvider?: boolean; } | undefined;
+    // registerMultiDocumentHighlightProvider is proposed API; guard against it not being available.
+    if (!capabilities?.customMultiDocumentHighlightProvider || typeof vscode.languages.registerMultiDocumentHighlightProvider !== "function") {
+        return { dispose() {} };
+    }
+    return vscode.languages.registerMultiDocumentHighlightProvider(selector, new MultiDocumentHighlightProvider(client));
+}

_extension/src/vscode.proposed.multiDocumentHighlightProvider.d.ts

@@ -0,0 +1,58 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+declare module "vscode" {
+    /**
+     * Represents a collection of document highlights from multiple documents.
+     */
+    export class MultiDocumentHighlight {
+        /**
+         * The URI of the document containing the highlights.
+         */
+        uri: Uri;
+
+        /**
+         * The highlights for the document.
+         */
+        highlights: DocumentHighlight[];
+
+        /**
+         * Creates a new instance of MultiDocumentHighlight.
+         * @param uri The URI of the document containing the highlights.
+         * @param highlights The highlights for the document.
+         */
+        constructor(uri: Uri, highlights: DocumentHighlight[]);
+    }
+
+    export interface MultiDocumentHighlightProvider {
+        /**
+         * Provide a set of document highlights, like all occurrences of a variable or
+         * all exit-points of a function.
+         *
+         * @param document The document in which the command was invoked.
+         * @param position The position at which the command was invoked.
+         * @param otherDocuments An array of additional valid documents for which highlights should be provided.
+         * @param token A cancellation token.
+         * @returns A Map containing a mapping of the Uri of a document to the document highlights or a thenable that resolves to such. The lack of a result can be
+         * signaled by returning `undefined`, `null`, or an empty map.
+         */
+        provideMultiDocumentHighlights(document: TextDocument, position: Position, otherDocuments: TextDocument[], token: CancellationToken): ProviderResult<MultiDocumentHighlight[]>;
+    }
+
+    namespace languages {
+        /**
+         * Register a multi document highlight provider.
+         *
+         * Multiple providers can be registered for a language. In that case providers are sorted
+         * by their {@link languages.match score} and groups sequentially asked for document highlights.
+         * The process stops when a provider returns a `non-falsy` or `non-failure` result.
+         *
+         * @param selector A selector that defines the documents this provider is applicable to.
+         * @param provider A multi-document highlight provider.
+         * @returns A {@link Disposable} that unregisters this provider when being disposed.
+         */
+        export function registerMultiDocumentHighlightProvider(selector: DocumentSelector, provider: MultiDocumentHighlightProvider): Disposable;
+    }
+}

internal/fourslash/_scripts/convertFourslash.mts

@@ -1262,6 +1262,7 @@ function parseBaselineFindAllReferencesArgs(args: readonly ts.Expression[]): [Ve
 function parseBaselineDocumentHighlightsArgs(args: readonly ts.Expression[]): [VerifyBaselineDocumentHighlightsCmd] {
     const newArgs: string[] = [];
     let preferences: string | undefined;
+    let filesToSearch: string[] | undefined;
     for (const arg of args) {
         let strArg;
         if (strArg = getArrayLiteralExpression(arg)) {
@@ -1270,8 +1271,47 @@ function parseBaselineDocumentHighlightsArgs(args: readonly ts.Expression[]): [V
                 newArgs.push(newArg);
             }
         }
+        else if (ts.isCallExpression(arg) && arg.getText().includes("test.ranges()")) {
+            newArgs.push("ToAny(f.Ranges())...");
+        }
         else if (ts.isObjectLiteralExpression(arg)) {
-            // !!! todo when multiple files supported in lsp
+            for (const prop of arg.properties) {
+                if (ts.isPropertyAssignment(prop) && ts.isIdentifier(prop.name) && prop.name.text === "filesToSearch" && ts.isArrayLiteralExpression(prop.initializer)) {
+                    filesToSearch = [];
+                    for (const e of prop.initializer.elements) {
+                        if (ts.isStringLiteral(e)) {
+                            filesToSearch.push(JSON.stringify(e.text));
+                        }
+                        else if (ts.isPropertyAccessExpression(e) && e.name.text === "fileName") {
+                            // e.g. test.ranges()[0].fileName -> f.Ranges()[0].FileName()
+                            const obj = e.expression;
+                            if (ts.isElementAccessExpression(obj) && ts.isCallExpression(obj.expression) && obj.expression.getText().includes("ranges")) {
+                                const index = obj.argumentExpression?.getText();
+                                if (index !== undefined) {
+                                    filesToSearch.push(`f.Ranges()[${index}].FileName()`);
+                                    continue;
+                                }
+                            }
+                            // e.g. range.fileName where `const range = test.ranges()[0]`
+                            if (ts.isIdentifier(obj)) {
+                                const resolved = parseRangeVariable(obj);
+                                if (resolved) {
+                                    filesToSearch.push(`${resolved}.FileName()`);
+                                    continue;
+                                }
+                            }
+                            // Fallback: skip filesToSearch entirely
+                            filesToSearch = undefined;
+                            break;
+                        }
+                        else {
+                            // Unsupported expression; skip filesToSearch
+                            filesToSearch = undefined;
+                            break;
+                        }
+                    }
+                }
+            }
         }
         else {
             newArgs.push(parseBaselineMarkerOrRangeArg(arg));
@@ -1286,6 +1326,7 @@ function parseBaselineDocumentHighlightsArgs(args: readonly ts.Expression[]): [V
         kind: "verifyBaselineDocumentHighlights",
         args: newArgs,
         preferences: preferences ? preferences : "nil /*preferences*/",
+        filesToSearch,
     }];
 }
 
@@ -1841,6 +1882,10 @@ function parseRangeVariable(arg: ts.Identifier | ts.ElementAccessExpression): st
                 if (ts.isElementAccessExpression(arg)) {
                     return `f.Ranges()[${arg.argumentExpression!.getText()}]`;
                 }
+                // `const range = test.ranges()[0]` used directly as `range`
+                if (ts.isIdentifier(arg) && ts.isElementAccessExpression(decl.initializer) && ts.isCallExpression(decl.initializer.expression) && decl.initializer.argumentExpression) {
+                    return `f.Ranges()[${decl.initializer.argumentExpression.getText()}]`;
+                }
             }
             // `const cRanges = ranges.get("C")` or `const cRanges = test.rangesByText().get("C")`
             if (ts.isIdentifier(decl.name) && decl.name.text === argName && decl.initializer && ts.isCallExpression(decl.initializer)) {
@@ -3022,6 +3067,7 @@ interface VerifyBaselineDocumentHighlightsCmd {
     kind: "verifyBaselineDocumentHighlights";
     args: string[];
     preferences: string;
+    filesToSearch?: string[];
 }
 
 interface VerifyBaselineInlayHintsCmd {
@@ -3292,7 +3338,11 @@ function generateBaselineFindAllReferences({ markers, ranges }: VerifyBaselineFi
     return `f.VerifyBaselineFindAllReferences(t, ${markers.join(", ")})`;
 }
 
-function generateBaselineDocumentHighlights({ args, preferences }: VerifyBaselineDocumentHighlightsCmd): string {
+function generateBaselineDocumentHighlights({ args, preferences, filesToSearch }: VerifyBaselineDocumentHighlightsCmd): string {
+    if (filesToSearch) {
+        const filesGo = `[]string{${filesToSearch.join(", ")}}`;
+        return `f.VerifyBaselineDocumentHighlightsWithOptions(t, ${preferences}, ${filesGo}, ${args.join(", ")})`;
+    }
     return `f.VerifyBaselineDocumentHighlights(t, ${preferences}, ${args.join(", ")})`;
 }
 

internal/fourslash/_scripts/unparsedTests.txt

@@ -1460,10 +1460,6 @@ docCommentTemplateWithMultipleJSDoc1.ts parse error: "Unrecognized fourslash sta
 docCommentTemplateWithMultipleJSDoc2.ts parse error: "Unrecognized fourslash statement: verify.docCommentTemplateAt(...)"
 docCommentTemplateWithMultipleJSDoc3.ts parse error: "Unrecognized fourslash statement: verify.docCommentTemplateAt(...)"
 docCommentTemplateWithMultipleJSDocAndParameters.ts parse error: "Unrecognized fourslash statement: verify.docCommentTemplateAt(...)"
-documentHighlights_moduleImport_filesToSearch.ts parse error: "Unrecognized marker or range argument: test.ranges()"
-documentHighlights_moduleImport_filesToSearchWithInvalidFile.ts parse error: "Unrecognized marker or range argument: test.ranges()"
-documentHighlights_windowsPath.ts parse error: "Unrecognized marker or range argument: range"
-documentHighlights02.ts parse error: "Unrecognized marker or range argument: test.ranges()"
 duplicateClassModuleError0.ts parse error: "Unrecognized edit function: disableFormatting"
 editClearsJsDocCache.ts parse error: "Unrecognized edit function: replace"
 eval.ts parse error: "Unrecognized fourslash statement: verify.eval(...)"