Add inlay chaining hints support (#220)

* feat(rust-analyzer-api): update

* feat(ctx): inlay hints

* feat(ctx): inlay hints works

* feat(ctx): chaining hints works

close #177

* fix(cmds): SSR param changed

* chore(doc): docs for CocRustChainingHint

README.md

@@ -18,6 +18,7 @@ This extension is configured using a jsonc file. You can open this configuration
 - `rust-analyzer.updates.channel`: Use `stable` or `nightly` updates, default: `stable`
 - `rust-analyzer.diagnostics.enable`: Whether to show native rust-analyzer diagnostics, default: `true`
 - `rust-analyzer.lruCapacity`: Number of syntax trees rust-analyzer keeps in memory, default: `null`
+- `rust-analyzer.inlayHints.chainingHints`: Whether to show inlay type hints for method chains, **Neovim Only**, default `true`
 - `rust-analyzer.files.watcher`: Controls file watching implementation, default: `client`
 - `rust-analyzer.files.exclude`: Paths to exclude from analysis, default: `[]`
 - `rust-analyzer.notifications.workspaceLoaded`: Whether to show `workspace loaded` message, default: `true`
@@ -62,6 +63,10 @@ You can use these commands by `:CocCommand XYZ`.
 - `rust-analyzer.serverVersion`: Show current Rust Analyzer server version
 - `rust-analyzer.upgrade`: Download latest `rust-analyzer` from [GitHub release](https://github.com/rust-analyzer/rust-analyzer/releases)
 
+## Highlight Group
+
+- `CocRustChainingHint`: highlight name for `chainingHints`, default link to `CocHintVirtualText`
+
 ## License
 
 MIT

package.json

@@ -196,6 +196,11 @@
           "default": true,
           "markdownDescription": "Check all targets and tests (will be passed as `--all-targets`)"
         },
+        "rust-analyzer.inlayHints.chainingHints": {
+          "type": "boolean",
+          "default": true,
+          "description": "Whether to show inlay type hints for method chains. *Neovim Only*"
+        },
         "rust-analyzer.completion.addCallParenthesis": {
           "type": "boolean",
           "default": true,

src/cmds/ssr.ts

@@ -15,7 +15,12 @@ export function ssr(ctx: Ctx): Cmd {
       return;
     }
 
-    const change = await ctx.client.sendRequest(ra.ssr, { arg: input });
+    const param: ra.SsrParams = {
+      query: input,
+      parseOnly: false,
+    };
+
+    const change = await ctx.client.sendRequest(ra.ssr, param);
     await applySourceChange(change);
   };
 }

src/config.ts

@@ -44,6 +44,13 @@ export class Config {
     return this.cfg.get<null | string>('serverPath')!;
   }
 
+  get inlayHints() {
+    const hasVirtualText = workspace.isNvim && workspace.nvim.hasFunction('nvim_buf_set_virtual_text');
+    return {
+      chainingHints: hasVirtualText && this.cfg.get<boolean>('inlayHints.chainingHints'),
+    };
+  }
+
   get checkOnSave() {
     return {
       command: this.cfg.get<string>('checkOnSave.command')!,

src/ctx.ts

@@ -3,12 +3,17 @@ import executable from 'executable';
 import { existsSync } from 'fs';
 import { homedir } from 'os';
 import { join } from 'path';
-import { Disposable, WorkDoneProgress } from 'vscode-languageserver-protocol';
+import { CancellationToken, Disposable, ErrorCodes, RequestType, TextDocument, WorkDoneProgress } from 'vscode-languageserver-protocol';
 import { createClient } from './client';
 import { Config } from './config';
 import { downloadServer, getLatestRelease } from './downloader';
 import { StatusDisplay } from './status_display';
 
+export type RustDocument = TextDocument & { languageId: 'rust' };
+export function isRustDocument(document: TextDocument): document is RustDocument {
+  return document.languageId === 'rust';
+}
+
 export type Cmd = (...args: any[]) => unknown;
 
 export class Ctx {
@@ -49,6 +54,10 @@ export class Ctx {
     return this.extCtx.subscriptions;
   }
 
+  pushCleanup(d: Disposable) {
+    this.extCtx.subscriptions.push(d);
+  }
+
   resolveBin(): string | undefined {
     // 1. from config, custom server path
     // 2. bundled
@@ -112,4 +121,29 @@ export class Ctx {
       commands.executeCommand('vscode.open', 'https://github.com/rust-analyzer/rust-analyzer/releases').catch(() => {});
     }
   }
+
+  sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms));
+
+  async sendRequestWithRetry<TParam, TRet>(reqType: RequestType<TParam, TRet, unknown>, param: TParam, token?: CancellationToken): Promise<TRet> {
+    for (const delay of [2, 4, 6, 8, 10, null]) {
+      try {
+        return await (token ? this.client.sendRequest(reqType, param, token) : this.client.sendRequest(reqType, param));
+      } catch (error) {
+        if (delay === null) {
+          throw error;
+        }
+
+        if (error.code === ErrorCodes.RequestCancelled) {
+          throw error;
+        }
+
+        if (error.code !== ErrorCodes.ContentModified) {
+          throw error;
+        }
+
+        await this.sleep(10 * (1 << delay));
+      }
+    }
+    throw 'unreachable';
+  }
 }

src/index.ts

@@ -4,6 +4,7 @@ import * as cmds from './cmds';
 import { Config } from './config';
 import { Ctx } from './ctx';
 import { downloadServer } from './downloader';
+import { activateInlayHints } from './inlay_hints';
 
 export async function activate(context: ExtensionContext): Promise<void> {
   const config = new Config();
@@ -33,6 +34,8 @@ export async function activate(context: ExtensionContext): Promise<void> {
 
   await ctx.startServer();
 
+  activateInlayHints(ctx);
+
   ctx.registerCommand('analyzerStatus', cmds.analyzerStatus);
   ctx.registerCommand('applySourceChange', cmds.applySourceChange);
   ctx.registerCommand('selectAndApplySourceChange', cmds.selectAndApplySourceChange);

src/inlay_hints.ts

@@ -0,0 +1,141 @@
+import { Disposable, workspace } from 'coc.nvim';
+import { CancellationTokenSource } from 'vscode-languageserver-protocol';
+import { Ctx, isRustDocument, RustDocument } from './ctx';
+import * as ra from './rust-analyzer-api';
+
+interface InlaysDecorations {
+  type: ra.InlayHint[];
+  param: ra.InlayHint[];
+  chaining: ra.InlayHint[];
+}
+
+interface RustSourceFile {
+  /**
+   * Source of the token to cancel in-flight inlay hints request if any.
+   */
+  inlaysRequest: null | CancellationTokenSource;
+  /**
+   * Last applied decorations.
+   */
+  cachedDecorations: null | InlaysDecorations;
+
+  document: RustDocument;
+}
+
+class HintsUpdater implements Disposable {
+  private sourceFiles = new Map<string, RustSourceFile>(); // map Uri -> RustSourceFile
+  private readonly disposables: Disposable[] = [];
+  private chainingHintNS = workspace.createNameSpace('rust-chaining-hint');
+
+  constructor(private readonly ctx: Ctx) {
+    workspace.onDidChangeTextDocument(
+      async (e) => {
+        const doc = workspace.getDocument(e.bufnr);
+        if (isRustDocument(doc.textDocument)) {
+          this.syncCacheAndRenderHints();
+        }
+      },
+      this,
+      this.disposables
+    );
+
+    // Set up initial cache shape
+    workspace.documents.forEach((doc) => {
+      if (isRustDocument(doc.textDocument)) {
+        doc.buffer.clearNamespace(this.chainingHintNS);
+        this.sourceFiles.set(doc.uri, { document: doc.textDocument, inlaysRequest: null, cachedDecorations: null });
+      }
+    });
+
+    this.syncCacheAndRenderHints();
+  }
+
+  dispose() {
+    this.sourceFiles.forEach((file) => file.inlaysRequest?.cancel());
+    this.disposables.forEach((d) => d.dispose());
+  }
+
+  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 doc of workspace.documents) {
+          if (doc.uri === uri) {
+            this.renderDecorations(file.cachedDecorations);
+          }
+        }
+      })
+    );
+  }
+
+  private async renderDecorations(decorations: InlaysDecorations) {
+    const doc = await workspace.document;
+    if (!doc) return;
+
+    doc.buffer.clearNamespace(this.chainingHintNS);
+    for (const item of decorations.chaining) {
+      doc.buffer.setVirtualText(this.chainingHintNS, item.range.end.line, [[item.label, 'CocRustChainingHint']], {}).logError();
+    }
+  }
+
+  private hintsToDecorations(hints: ra.InlayHint[]): InlaysDecorations {
+    const decorations: InlaysDecorations = { type: [], param: [], chaining: [] };
+    for (const hint of hints) {
+      // ChainingHint only now
+      if (hint.kind === ra.InlayHint.Kind.ChainingHint) {
+        decorations.chaining.push(hint);
+      }
+    }
+
+    return decorations;
+  }
+
+  private async fetchHints(file: RustSourceFile): Promise<null | ra.InlayHint[]> {
+    file.inlaysRequest?.cancel();
+
+    const tokenSource = new CancellationTokenSource();
+    file.inlaysRequest = tokenSource;
+
+    const param = { textDocument: { uri: file.document.uri.toString() } };
+    return this.ctx
+      .sendRequestWithRetry(ra.inlayHints, param, tokenSource.token)
+      .catch(() => null)
+      .finally(() => {
+        if (file.inlaysRequest === tokenSource) {
+          file.inlaysRequest = null;
+        }
+      });
+  }
+}
+
+export function activateInlayHints(ctx: Ctx) {
+  const maybeUpdater = {
+    updater: null as null | HintsUpdater,
+    async onConfigChange() {
+      if (!ctx.config.inlayHints.chainingHints) {
+        return this.dispose();
+      }
+
+      await ctx.sleep(100);
+      await workspace.nvim.command('hi default link CocRustChainingHint CocHintVirtualText');
+      if (this.updater) {
+        this.updater.syncCacheAndRenderHints();
+      } else {
+        this.updater = new HintsUpdater(ctx);
+      }
+    },
+    dispose() {
+      this.updater?.dispose();
+      this.updater = null;
+    },
+  };
+
+  ctx.pushCleanup(maybeUpdater);
+
+  workspace.onDidChangeConfiguration(maybeUpdater.onConfigChange, maybeUpdater, ctx.subscriptions);
+  maybeUpdater.onConfigChange();
+}

src/rust-analyzer-api.ts

@@ -71,32 +71,36 @@ export interface Runnable {
   label: string;
   bin: string;
   args: Vec<string>;
+  extraArgs: Vec<string>;
   env: FxHashMap<string, string>;
   cwd: Option<string>;
 }
 export const runnables = request<RunnablesParams, Vec<Runnable>>('runnables');
 
-export type InlayHint = InlayHint.TypeHint | InlayHint.ParamHint;
+export type InlayHint = InlayHint.TypeHint | InlayHint.ParamHint | InlayHint.ChainingHint;
 
 export namespace InlayHint {
   export const enum Kind {
     TypeHint = 'TypeHint',
     ParamHint = 'ParameterHint',
+    ChainingHint = 'ChainingHint',
   }
   interface Common {
     range: Range;
     label: string;
   }
   export type TypeHint = Common & { kind: Kind.TypeHint };
   export type ParamHint = Common & { kind: Kind.ParamHint };
+  export type ChainingHint = Common & { kind: Kind.ChainingHint };
 }
 export interface InlayHintsParams {
   textDocument: TextDocumentIdentifier;
 }
 export const inlayHints = request<InlayHintsParams, Vec<InlayHint>>('inlayHints');
 
 export interface SsrParams {
-  arg: string;
+  query: string;
+  parseOnly: boolean;
 }
 export const ssr = request<SsrParams, SourceChange>('ssr');