1 files changed, 136 insertions, 222 deletions

diff --git a/editors/code/src/inlay_hints.ts b/editors/code/src/inlay_hints.ts
index 161b34037..6d084362d 100644
--- a/editors/code/src/inlay_hints.ts
+++ b/editors/code/src/inlay_hints.ts
@@ -2,48 +2,32 @@ import * as lc from "vscode-languageclient";
 import * as vscode from 'vscode';
 import * as ra from './rust-analyzer-api';
 
-import { Ctx } from './ctx';
-import { sendRequestWithRetry, assert } from './util';
+import { Ctx, Disposable } from './ctx';
+import { sendRequestWithRetry, isRustDocument, RustDocument, RustEditor, log } from './util';
 
-export function activateInlayHints(ctx: Ctx) {
-    const hintsUpdater = new HintsUpdater(ctx.client);
-
-    vscode.window.onDidChangeVisibleTextEditors(
-        () => hintsUpdater.refreshVisibleRustEditors(),
-        null,
-        ctx.subscriptions
-    );
-
-    vscode.workspace.onDidChangeTextDocument(
-        ({ contentChanges, document }) => {
-            if (contentChanges.length === 0) return;
-            if (!isRustTextDocument(document)) return;
 
-            hintsUpdater.forceRefreshVisibleRustEditors();
+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 = HintsUpdater.create(ctx);
         },
-        null,
-        ctx.subscriptions
-    );
+        dispose() {
+            this.updater?.dispose();
+            this.updater = null;
+        }
+    };
+
+    ctx.pushCleanup(maybeUpdater);
 
     vscode.workspace.onDidChangeConfiguration(
-        async _ => {
-            // FIXME: ctx.config may have not been refreshed at this point of time, i.e.
-            // it's on onDidChangeConfiguration() handler may've not executed yet
-            // (order of invokation is unspecified)
-            // To fix this we should expose an event emitter from our `Config` itself.
-            await hintsUpdater.setEnabled(ctx.config.displayInlayHints);
-        },
-        null,
-        ctx.subscriptions
+        maybeUpdater.onConfigChange, maybeUpdater, ctx.subscriptions
     );
 
-    ctx.pushCleanup({
-        dispose() {
-            hintsUpdater.clearHints();
-        }
-    });
-
-    hintsUpdater.setEnabled(ctx.config.displayInlayHints);
+    maybeUpdater.onConfigChange();
 }
 
 
@@ -79,239 +63,169 @@ const paramHints = {
     }
 };
 
-class HintsUpdater {
-    private sourceFiles = new RustSourceFiles();
-    private enabled = false;
-
-    constructor(readonly client: lc.LanguageClient) { }
+class HintsUpdater implements Disposable {
+    private sourceFiles = new Map<string, RustSourceFile>(); // map Uri -> RustSourceFile
+    private readonly disposables: Disposable[] = [];
 
-    setEnabled(enabled: boolean) {
-        if (this.enabled === enabled) return;
-        this.enabled = enabled;
+    private constructor(readonly ctx: Ctx) { }
 
-        if (this.enabled) {
-            this.refreshVisibleRustEditors();
-        } else {
-            this.clearHints();
-        }
-    }
-
-    clearHints() {
-        for (const file of this.sourceFiles) {
-            file.inlaysRequest?.cancel();
-            file.renderHints([], this.client.protocol2CodeConverter);
-        }
-    }
-
-    forceRefreshVisibleRustEditors() {
-        if (!this.enabled) return;
+    static create(ctx: Ctx) {
+        const self = new HintsUpdater(ctx);
 
-        for (const file of this.sourceFiles) {
-            void file.fetchAndRenderHints(this.client);
-        }
-    }
-
-    refreshVisibleRustEditors() {
-        if (!this.enabled) return;
+        vscode.window.onDidChangeVisibleTextEditors(
+            self.onDidChangeVisibleTextEditors,
+            self,
+            self.disposables
+        );
 
-        const visibleSourceFiles = this.sourceFiles.drainEditors(
-            vscode.window.visibleTextEditors.filter(isRustTextEditor)
+        vscode.workspace.onDidChangeTextDocument(
+            self.onDidChangeTextDocument,
+            self,
+            self.disposables
         );
 
-        // Cancel requests for source files whose editors were disposed (leftovers after drain).
-        for (const { inlaysRequest } of this.sourceFiles) inlaysRequest?.cancel();
+        // Set up initial cache shape
+        ctx.visibleRustEditors.forEach(editor => self.sourceFiles.set(
+            editor.document.uri.toString(), {
+                document: editor.document,
+                inlaysRequest: null,
+                cachedDecorations: null
+            }
+        ));
 
-        this.sourceFiles = visibleSourceFiles;
+        self.syncCacheAndRenderHints();
 
-        for (const file of this.sourceFiles) {
-            if (!file.rerenderHints()) {
-                void file.fetchAndRenderHints(this.client);
-            }
-        }
+        return self;
     }
-}
-
 
-/**
- * This class encapsulates a map of file uris to respective inlay hints
- * request cancellation token source (cts) and an array of editors.
- * E.g.
- * ```
- * {
- *    file1.rs -> (cts, (typeDecor, paramDecor), [editor1, editor2])
- *                  ^-- there is a cts to cancel the in-flight request
- *    file2.rs -> (cts, null, [editor3])
- *                       ^-- no decorations are applied to this source file yet
- *    file3.rs -> (null, (typeDecor, paramDecor), [editor4])
- * }                ^-- there is no inflight request
- * ```
- *
- * Invariants: each stored source file has at least 1 editor.
- */
-class RustSourceFiles {
-    private files = new Map<string, RustSourceFile>();
+    dispose() {
+        this.sourceFiles.forEach(file => file.inlaysRequest?.cancel());
+        this.ctx.visibleRustEditors.forEach(editor => this.renderDecorations(editor, { param: [], type: [] }));
+        this.disposables.forEach(d => d.dispose());
+    }
 
-    /**
-     * Removes `editors` from `this` source files and puts them into a returned
-     * source files object. cts and decorations are moved to the returned source files.
-     */
-    drainEditors(editors: RustTextEditor[]): RustSourceFiles {
-        const result = new RustSourceFiles;
+    onDidChangeTextDocument({contentChanges, document}: vscode.TextDocumentChangeEvent) {
+        if (contentChanges.length === 0 || !isRustDocument(document)) return;
+        log.debug(`[inlays]: changed text doc!`);
+        this.syncCacheAndRenderHints();
+    }
 
-        for (const editor of editors) {
-            const oldFile = this.removeEditor(editor);
-            const newFile = result.addEditor(editor);
+    private syncCacheAndRenderHints() {
+        // FIXME: make inlayHints request pass an array of files?
+        this.sourceFiles.forEach((file, uri) => this.fetchHints(file).then(hints => {
+            if (!hints) return;
 
-            if (oldFile) newFile.stealCacheFrom(oldFile);
-        }
+            file.cachedDecorations = this.hintsToDecorations(hints);
 
-        return result;
+            for (const editor of this.ctx.visibleRustEditors) {
+                if (editor.document.uri.toString() === uri) {
+                    this.renderDecorations(editor, file.cachedDecorations);
+                }
+            }
+        }));
     }
 
-    /**
-     * Remove the editor and if it was the only editor for a source file,
-     * the source file is removed altogether.
-     *
-     * @returns A reference to the source file for this editor or
-     *          null if no such source file was not found.
-     */
-    private removeEditor(editor: RustTextEditor): null | RustSourceFile {
-        const uri = editor.document.uri.toString();
+    onDidChangeVisibleTextEditors() {
+        log.debug(`[inlays]: changed visible text editors`);
+        const newSourceFiles = new Map<string, RustSourceFile>();
 
-        const file = this.files.get(uri);
-        if (!file) return null;
+        // 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);
 
-        const editorIndex = file.editors.findIndex(suspect => areEditorsEqual(suspect, editor));
+            // No text documents changed, so we may try to use the cache
+            if (!file.cachedDecorations) {
+                file.inlaysRequest?.cancel();
 
-        if (editorIndex >= 0) {
-            file.editors.splice(editorIndex, 1);
+                const hints = await this.fetchHints(file);
+                if (!hints) return;
 
-            if (file.editors.length === 0) this.files.delete(uri);
-        }
-
-        return file;
-    }
-
-    /**
-     * @returns A reference to an existing source file or newly created one for the editor.
-     */
-    private addEditor(editor: RustTextEditor): RustSourceFile {
-        const uri = editor.document.uri.toString();
-        const file = this.files.get(uri);
-
-        if (!file) {
-            const newFile = new RustSourceFile([editor]);
-            this.files.set(uri, newFile);
-            return newFile;
-        }
+                file.cachedDecorations = this.hintsToDecorations(hints);
+            }
 
-        if (!file.editors.find(suspect => areEditorsEqual(suspect, editor))) {
-            file.editors.push(editor);
-        }
-        return file;
-    }
+            this.renderDecorations(editor, file.cachedDecorations);
+        });
 
-    getSourceFile(uri: string): undefined | RustSourceFile {
-        return this.files.get(uri);
-    }
+        // Cancel requests for no longer visible (disposed) source files
+        this.sourceFiles.forEach((file, uri) => {
+            if (!newSourceFiles.has(uri)) file.inlaysRequest?.cancel();
+        });
 
-    [Symbol.iterator](): IterableIterator<RustSourceFile> {
-        return this.files.values();
-    }
-}
-class RustSourceFile {
-    constructor(
-        /**
-         * Editors for this source file (one text document may be opened in multiple editors).
-         * We keep this just an array, because most of the time we have 1 editor for 1 source file.
-         */
-        readonly editors: RustTextEditor[],
-        /**
-         * Source of the token to cancel in-flight inlay hints request if any.
-         */
-        public inlaysRequest: null | vscode.CancellationTokenSource = null,
-
-        public decorations: null | {
-            type: vscode.DecorationOptions[];
-            param: vscode.DecorationOptions[];
-        } = null
-    ) { }
-
-    stealCacheFrom(other: RustSourceFile) {
-        if (other.inlaysRequest) this.inlaysRequest = other.inlaysRequest;
-        if (other.decorations) this.decorations = other.decorations;
-
-        other.inlaysRequest = null;
-        other.decorations = null;
+        this.sourceFiles = newSourceFiles;
     }
 
-    rerenderHints(): boolean {
-        if (!this.decorations) return false;
-
-        for (const editor of this.editors) {
-            editor.setDecorations(typeHints.decorationType, this.decorations.type);
-            editor.setDecorations(paramHints.decorationType, this.decorations.param);
-        }
-        return true;
+    private renderDecorations(editor: RustEditor, decorations: InlaysDecorations) {
+        editor.setDecorations(typeHints.decorationType, decorations.type);
+        editor.setDecorations(paramHints.decorationType, decorations.param);
     }
 
-    renderHints(hints: ra.InlayHint[], conv: lc.Protocol2CodeConverter) {
-        this.decorations = { type: [], param: [] };
+    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: {
-                    this.decorations.type.push(typeHints.toDecoration(hint, conv));
+                    decorations.type.push(typeHints.toDecoration(hint, conv));
                     continue;
                 }
                 case ra.InlayHint.Kind.ParamHint: {
-                    this.decorations.param.push(paramHints.toDecoration(hint, conv));
+                    decorations.param.push(paramHints.toDecoration(hint, conv));
                     continue;
                 }
             }
         }
-        this.rerenderHints();
+        return decorations;
     }
 
-    async fetchAndRenderHints(client: lc.LanguageClient): Promise<void> {
-        this.inlaysRequest?.cancel();
+    lastReqId = 0;
+    private async fetchHints(file: RustSourceFile): Promise<null | ra.InlayHint[]> {
+        const reqId = ++this.lastReqId;
+
+        log.debug(`[inlays]: ${reqId} requesting`);
+        file.inlaysRequest?.cancel();
 
         const tokenSource = new vscode.CancellationTokenSource();
-        this.inlaysRequest = tokenSource;
-
-        const request = { textDocument: { uri: this.editors[0].document.uri.toString() } };
-
-        try {
-            const hints = await sendRequestWithRetry(client, ra.inlayHints, request, tokenSource.token);
-            this.renderHints(hints, client.protocol2CodeConverter);
-        } catch {
-            /* ignore */
-        } finally {
-            if (this.inlaysRequest === tokenSource) {
-                this.inlaysRequest = null;
-            }
-        }
+        file.inlaysRequest = tokenSource;
+
+        const request = { textDocument: { uri: file.document.uri.toString() } };
+
+        return sendRequestWithRetry(this.ctx.client, ra.inlayHints, request, tokenSource.token)
+            .catch(_ => {
+                log.debug(`[inlays]: ${reqId} err`);
+                return null;
+            })
+            .finally(() => {
+                if (file.inlaysRequest === tokenSource) {
+                    file.inlaysRequest = null;
+                    log.debug(`[inlays]: ${reqId} got response!`);
+                } else {
+                    log.debug(`[inlays]: ${reqId} cancelled!`);
+                }
+            })
     }
 }
 
-type RustTextDocument = vscode.TextDocument & { languageId: "rust" };
-type RustTextEditor = vscode.TextEditor & { document: RustTextDocument; id: string };
-
-function areEditorsEqual(a: RustTextEditor, b: RustTextEditor): boolean {
-    return a.id === b.id;
+interface InlaysDecorations {
+    type: vscode.DecorationOptions[];
+    param: vscode.DecorationOptions[];
 }
 
-function isRustTextEditor(suspect: vscode.TextEditor & { id?: unknown }): suspect is RustTextEditor {
-    // Dirty hack, we need to access private vscode editor id,
-    // see https://github.com/microsoft/vscode/issues/91788
-    assert(
-        typeof suspect.id === "string",
-        "Private text editor id is no longer available, please update the workaround!"
-    );
-
-    return isRustTextDocument(suspect.document);
-}
+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;
 
-function isRustTextDocument(suspect: vscode.TextDocument): suspect is RustTextDocument {
-    return suspect.languageId === "rust";
+    document: RustDocument
 }