Merge #3378

3378: vscode: redesign inlay hints to be capable of handling multiple editors for one source file r=Veetaha a=Veetaha

Fixes: #3008 (inlay corruption with multiple editors for one file).
Fixes: #3319 (unnecessary requests for inlay hints when switching unrelated source files or output/log/console windows)
Also, I don't know how, but the problem described in #3057 doesn't appear for me anymore (maybe it was some fix on the server-side, idk), the inlay hints are displaying right away. Though the last time I checked this it was caused by the server returning an empty array of hints and responding with a very big latency, I am not sure that this redesign actually fixed #3057....

We didn't handle the case when one rust source file is open in multiple editors in vscode (e.g. by manually adding another editor for the same file or by opening an inline git diff view or just any kind of preview within the same file).

The git diff preview is actually quite special because it causes memory leaks in vscode (https://github.com/microsoft/vscode/issues/91782). It is not removed from `visibleEditors` once it is closed. However, this bug doesn't affect the inlay hints anymore, because we don't issue a request and set inlay hints for each editor in isolation. Editors are grouped by their respective files and we issue requests only for files and then update all duplicate editors using the results (so we just update the decorations for already closed git diff preview read-only editors).

Also, note on a hack I had to use. `vscode.TextEdtior` identity is not actually public, its `id` field is not exposed to us. I created a dedicated upstream issue for this (https://github.com/microsoft/vscode/issues/91788).

Regarding #3319: the newly designed hints client doesn't issue requests for type hints when switching the visible editors if it has them already cached (though it does rerender things anyway, but this could be optimized in future if so wanted).

<details>
<summary>Before</summary>

![bug_demo](https://user-images.githubusercontent.com/36276403/75613171-3cd0d480-5b33-11ea-9066-954fb2fb18a5.gif)


</details>

<details>
<summary> After </summary>

![multi-cursor-replace](https://user-images.githubusercontent.com/36276403/75612710-d5b12100-5b2e-11ea-99ba-214b4219e6d3.gif)

</details>

Co-authored-by: Veetaha <[email protected]>

4 files changed, 202 insertions, 134 deletions

diff --git a/editors/code/src/ctx.ts b/editors/code/src/ctx.ts
index b4e983a0c..25ef38aed 100644
--- a/editors/code/src/ctx.ts
+++ b/editors/code/src/ctx.ts
@@ -3,7 +3,7 @@ import * as lc from 'vscode-languageclient';
 
 import { Config } from './config';
 import { createClient } from './client';
-import { isRustDocument } from './util';
+import { isRustEditor, RustEditor } from './util';
 
 export class Ctx {
     private constructor(
@@ -22,17 +22,15 @@ export class Ctx {
         return res;
     }
 
-    get activeRustEditor(): vscode.TextEditor | undefined {
+    get activeRustEditor(): RustEditor | undefined {
         const editor = vscode.window.activeTextEditor;
-        return editor && isRustDocument(editor.document)
+        return editor && isRustEditor(editor)
             ? editor
             : undefined;
     }
 
-    get visibleRustEditors(): vscode.TextEditor[] {
-        return vscode.window.visibleTextEditors.filter(
-            editor => isRustDocument(editor.document),
-        );
+    get visibleRustEditors(): RustEditor[] {
+        return vscode.window.visibleTextEditors.filter(isRustEditor);
     }
 
     registerCommand(name: string, factory: (ctx: Ctx) => Cmd) {
diff --git a/editors/code/src/inlay_hints.ts b/editors/code/src/inlay_hints.ts
index 08d3a64a7..e1a82e03e 100644
--- a/editors/code/src/inlay_hints.ts
+++ b/editors/code/src/inlay_hints.ts
@@ -1,156 +1,214 @@
+import * as lc from "vscode-languageclient";
 import * as vscode from 'vscode';
 import * as ra from './rust-analyzer-api';
 
-import { Ctx } from './ctx';
-import { log, sendRequestWithRetry, isRustDocument } from './util';
+import { Ctx, Disposable } from './ctx';
+import { sendRequestWithRetry, isRustDocument, RustDocument, RustEditor } from './util';
 
-export function activateInlayHints(ctx: Ctx) {
-    const hintsUpdater = new HintsUpdater(ctx);
-    vscode.window.onDidChangeVisibleTextEditors(
-        async _ => hintsUpdater.refresh(),
-        null,
-        ctx.subscriptions
-    );
 
-    vscode.workspace.onDidChangeTextDocument(
-        async event => {
-            if (event.contentChanges.length === 0) return;
-            if (!isRustDocument(event.document)) return;
-            await hintsUpdater.refresh();
+export function activateInlayHints(ctx: Ctx) {
+    const maybeUpdater = {
+        updater: null as null | HintsUpdater,
+        onConfigChange() {
+            if (!ctx.config.displayInlayHints) {
+                return this.dispose();
+            }
+            if (!this.updater) this.updater = new HintsUpdater(ctx);
         },
-        null,
-        ctx.subscriptions
-    );
+        dispose() {
+            this.updater?.dispose();
+            this.updater = null;
+        }
+    };
+
+    ctx.pushCleanup(maybeUpdater);
 
     vscode.workspace.onDidChangeConfiguration(
-        async _ => hintsUpdater.setEnabled(ctx.config.displayInlayHints),
-        null,
-        ctx.subscriptions
+        maybeUpdater.onConfigChange, maybeUpdater, ctx.subscriptions
     );
 
-    ctx.pushCleanup({
-        dispose() {
-            hintsUpdater.clear();
+    maybeUpdater.onConfigChange();
+}
+
+
+const typeHints = {
+    decorationType: vscode.window.createTextEditorDecorationType({
+        after: {
+            color: new vscode.ThemeColor('rust_analyzer.inlayHint'),
+            fontStyle: "normal",
         }
-    });
+    }),
 
-    // XXX: we don't await this, thus Promise rejections won't be handled, but
-    // this should never throw in fact...
-    void hintsUpdater.setEnabled(ctx.config.displayInlayHints);
-}
+    toDecoration(hint: ra.InlayHint.TypeHint, conv: lc.Protocol2CodeConverter): vscode.DecorationOptions {
+        return {
+            range: conv.asRange(hint.range),
+            renderOptions: { after: { contentText: `: ${hint.label}` } }
+        };
+    }
+};
+
+const paramHints = {
+    decorationType: vscode.window.createTextEditorDecorationType({
+        before: {
+            color: new vscode.ThemeColor('rust_analyzer.inlayHint'),
+            fontStyle: "normal",
+        }
+    }),
 
-const typeHintDecorationType = vscode.window.createTextEditorDecorationType({
-    after: {
-        color: new vscode.ThemeColor('rust_analyzer.inlayHint'),
-        fontStyle: "normal",
-    },
-});
-
-const parameterHintDecorationType = vscode.window.createTextEditorDecorationType({
-    before: {
-        color: new vscode.ThemeColor('rust_analyzer.inlayHint'),
-        fontStyle: "normal",
-    },
-});
-
-class HintsUpdater {
-    private pending = new Map<string, vscode.CancellationTokenSource>();
-    private ctx: Ctx;
-    private enabled: boolean;
-
-    constructor(ctx: Ctx) {
-        this.ctx = ctx;
-        this.enabled = false;
+    toDecoration(hint: ra.InlayHint.ParamHint, conv: lc.Protocol2CodeConverter): vscode.DecorationOptions {
+        return {
+            range: conv.asRange(hint.range),
+            renderOptions: { before: { contentText: `${hint.label}: ` } }
+        };
     }
+};
 
-    async setEnabled(enabled: boolean): Promise<void> {
-        log.debug({ enabled, prev: this.enabled });
+class HintsUpdater implements Disposable {
+    private sourceFiles = new Map<string, RustSourceFile>(); // map Uri -> RustSourceFile
+    private readonly disposables: Disposable[] = [];
 
-        if (this.enabled === enabled) return;
-        this.enabled = enabled;
+    constructor(private readonly ctx: Ctx) {
+        vscode.window.onDidChangeVisibleTextEditors(
+            this.onDidChangeVisibleTextEditors,
+            this,
+            this.disposables
+        );
 
-        if (this.enabled) {
-            return await this.refresh();
-        } else {
-            return this.clear();
-        }
+        vscode.workspace.onDidChangeTextDocument(
+            this.onDidChangeTextDocument,
+            this,
+            this.disposables
+        );
+
+        // Set up initial cache shape
+        ctx.visibleRustEditors.forEach(editor => this.sourceFiles.set(
+            editor.document.uri.toString(),
+            {
+                document: editor.document,
+                inlaysRequest: null,
+                cachedDecorations: null
+            }
+        ));
+
+        this.syncCacheAndRenderHints();
     }
 
-    clear() {
-        this.ctx.visibleRustEditors.forEach(it => {
-            this.setTypeDecorations(it, []);
-            this.setParameterDecorations(it, []);
-        });
+    dispose() {
+        this.sourceFiles.forEach(file => file.inlaysRequest?.cancel());
+        this.ctx.visibleRustEditors.forEach(editor => this.renderDecorations(editor, { param: [], type: [] }));
+        this.disposables.forEach(d => d.dispose());
+    }
+
+    onDidChangeTextDocument({ contentChanges, document }: vscode.TextDocumentChangeEvent) {
+        if (contentChanges.length === 0 || !isRustDocument(document)) return;
+        this.syncCacheAndRenderHints();
     }
 
-    async refresh() {
-        if (!this.enabled) return;
-        await Promise.all(this.ctx.visibleRustEditors.map(it => this.refreshEditor(it)));
+    private syncCacheAndRenderHints() {
+        // FIXME: make inlayHints request pass an array of files?
+        this.sourceFiles.forEach((file, uri) => this.fetchHints(file).then(hints => {
+            if (!hints) return;
+
+            file.cachedDecorations = this.hintsToDecorations(hints);
+
+            for (const editor of this.ctx.visibleRustEditors) {
+                if (editor.document.uri.toString() === uri) {
+                    this.renderDecorations(editor, file.cachedDecorations);
+                }
+            }
+        }));
     }
 
-    private async refreshEditor(editor: vscode.TextEditor): Promise<void> {
-        const newHints = await this.queryHints(editor.document.uri.toString());
-        if (newHints == null) return;
-
-        const newTypeDecorations = newHints
-            .filter(hint => hint.kind === ra.InlayKind.TypeHint)
-            .map(hint => ({
-                range: this.ctx.client.protocol2CodeConverter.asRange(hint.range),
-                renderOptions: {
-                    after: {
-                        contentText: `: ${hint.label}`,
-                    },
-                },
-            }));
-        this.setTypeDecorations(editor, newTypeDecorations);
-
-        const newParameterDecorations = newHints
-            .filter(hint => hint.kind === ra.InlayKind.ParameterHint)
-            .map(hint => ({
-                range: this.ctx.client.protocol2CodeConverter.asRange(hint.range),
-                renderOptions: {
-                    before: {
-                        contentText: `${hint.label}: `,
-                    },
-                },
-            }));
-        this.setParameterDecorations(editor, newParameterDecorations);
+    onDidChangeVisibleTextEditors() {
+        const newSourceFiles = new Map<string, RustSourceFile>();
+
+        // Rerendering all, even up-to-date editors for simplicity
+        this.ctx.visibleRustEditors.forEach(async editor => {
+            const uri = editor.document.uri.toString();
+            const file = this.sourceFiles.get(uri) ?? {
+                document: editor.document,
+                inlaysRequest: null,
+                cachedDecorations: null
+            };
+            newSourceFiles.set(uri, file);
+
+            // No text documents changed, so we may try to use the cache
+            if (!file.cachedDecorations) {
+                file.inlaysRequest?.cancel();
+
+                const hints = await this.fetchHints(file);
+                if (!hints) return;
+
+                file.cachedDecorations = this.hintsToDecorations(hints);
+            }
+
+            this.renderDecorations(editor, file.cachedDecorations);
+        });
+
+        // Cancel requests for no longer visible (disposed) source files
+        this.sourceFiles.forEach((file, uri) => {
+            if (!newSourceFiles.has(uri)) file.inlaysRequest?.cancel();
+        });
+
+        this.sourceFiles = newSourceFiles;
     }
 
-    private setTypeDecorations(
-        editor: vscode.TextEditor,
-        decorations: vscode.DecorationOptions[],
-    ) {
-        editor.setDecorations(
-            typeHintDecorationType,
-            this.enabled ? decorations : [],
-        );
+    private renderDecorations(editor: RustEditor, decorations: InlaysDecorations) {
+        editor.setDecorations(typeHints.decorationType, decorations.type);
+        editor.setDecorations(paramHints.decorationType, decorations.param);
     }
 
-    private setParameterDecorations(
-        editor: vscode.TextEditor,
-        decorations: vscode.DecorationOptions[],
-    ) {
-        editor.setDecorations(
-            parameterHintDecorationType,
-            this.enabled ? decorations : [],
-        );
+    private hintsToDecorations(hints: ra.InlayHint[]): InlaysDecorations {
+        const decorations: InlaysDecorations = { type: [], param: [] };
+        const conv = this.ctx.client.protocol2CodeConverter;
+
+        for (const hint of hints) {
+            switch (hint.kind) {
+                case ra.InlayHint.Kind.TypeHint: {
+                    decorations.type.push(typeHints.toDecoration(hint, conv));
+                    continue;
+                }
+                case ra.InlayHint.Kind.ParamHint: {
+                    decorations.param.push(paramHints.toDecoration(hint, conv));
+                    continue;
+                }
+            }
+        }
+        return decorations;
     }
 
-    private async queryHints(documentUri: string): Promise<ra.InlayHint[] | null> {
-        this.pending.get(documentUri)?.cancel();
+    private async fetchHints(file: RustSourceFile): Promise<null | ra.InlayHint[]> {
+        file.inlaysRequest?.cancel();
 
         const tokenSource = new vscode.CancellationTokenSource();
-        this.pending.set(documentUri, tokenSource);
+        file.inlaysRequest = tokenSource;
 
-        const request = { textDocument: { uri: documentUri } };
+        const request = { textDocument: { uri: file.document.uri.toString() } };
 
         return sendRequestWithRetry(this.ctx.client, ra.inlayHints, request, tokenSource.token)
             .catch(_ => null)
             .finally(() => {
-                if (!tokenSource.token.isCancellationRequested) {
-                    this.pending.delete(documentUri);
+                if (file.inlaysRequest === tokenSource) {
+                    file.inlaysRequest = null;
                 }
             });
     }
 }
+
+interface InlaysDecorations {
+    type: vscode.DecorationOptions[];
+    param: vscode.DecorationOptions[];
+}
+
+interface RustSourceFile {
+    /*
+    * Source of the token to cancel in-flight inlay hints request if any.
+    */
+    inlaysRequest: null | vscode.CancellationTokenSource;
+    /**
+    * Last applied decorations.
+    */
+    cachedDecorations: null | InlaysDecorations;
+
+    document: RustDocument;
+}
diff --git a/editors/code/src/rust-analyzer-api.ts b/editors/code/src/rust-analyzer-api.ts
index c5a010e94..bd6e3ada0 100644
--- a/editors/code/src/rust-analyzer-api.ts
+++ b/editors/code/src/rust-analyzer-api.ts
@@ -86,14 +86,20 @@ export interface Runnable {
 export const runnables = request<RunnablesParams, Vec<Runnable>>("runnables");
 
 
-export const enum InlayKind {
-    TypeHint = "TypeHint",
-    ParameterHint = "ParameterHint",
-}
-export interface InlayHint {
-    range: lc.Range;
-    kind: InlayKind;
-    label: string;
+
+export type InlayHint = InlayHint.TypeHint | InlayHint.ParamHint;
+
+export namespace InlayHint {
+    export const enum Kind {
+        TypeHint = "TypeHint",
+        ParamHint = "ParameterHint",
+    }
+    interface Common {
+        range: lc.Range;
+        label: string;
+    }
+    export type TypeHint = Common & { kind: Kind.TypeHint };
+    export type ParamHint = Common & { kind: Kind.ParamHint };
 }
 export interface InlayHintsParams {
     textDocument: lc.TextDocumentIdentifier;
diff --git a/editors/code/src/util.ts b/editors/code/src/util.ts
index 7c95769bb..95a5f1227 100644
--- a/editors/code/src/util.ts
+++ b/editors/code/src/util.ts
@@ -1,7 +1,6 @@
 import * as lc from "vscode-languageclient";
 import * as vscode from "vscode";
 import { strict as nativeAssert } from "assert";
-import { TextDocument } from "vscode";
 
 export function assert(condition: boolean, explanation: string): asserts condition {
     try {
@@ -67,9 +66,16 @@ function sleep(ms: number) {
     return new Promise(resolve => setTimeout(resolve, ms));
 }
 
-export function isRustDocument(document: TextDocument) {
+export type RustDocument = vscode.TextDocument & { languageId: "rust" };
+export type RustEditor = vscode.TextEditor & { document: RustDocument; id: string };
+
+export function isRustDocument(document: vscode.TextDocument): document is RustDocument {
     return document.languageId === 'rust'
         // SCM diff views have the same URI as the on-disk document but not the same content
         && document.uri.scheme !== 'git'
         && document.uri.scheme !== 'svn';
-}
\ No newline at end of file
+}
+
+export function isRustEditor(editor: vscode.TextEditor): editor is RustEditor {
+    return isRustDocument(editor.document);
+}