fix(editor): make usage highlighting follow the caret at once and paint multi-line usages as a block (#3220)

Usage highlighting at the caret waited a fixed 300 ms after every caret move, so clicking around lagged. A click now asks at once and keyboard moves wait 100 ms; with the document unchanged, a caret landing on usages found recently shows them immediately while they are refreshed. A highlight over several lines (a return statement's exit point) is painted as a block: from its start to the right edge, whole lines between, and the last line up to its end.

Co-authored-by: PLZY <41029728+PLZY@users.noreply.github.com>

src/editor/lsp/navigation.ts

@@ -10,7 +10,16 @@ import {
  * click), the multiple-targets chooser, and document highlights (occurrences
  * of the symbol at the caret, plus the sticky find.highlightUsages variant).
  */
-import { EditorSelection, Prec, StateEffect, StateField, type Extension, type Text } from "@codemirror/state";
+import {
+  countColumn,
+  EditorSelection,
+  Prec,
+  StateEffect,
+  StateField,
+  type Extension,
+  type Text,
+} from "@codemirror/state";
+import { LINE_INSET } from "@/editor/core/cm/lineInset";
 import {
   Decoration,
   EditorView,
@@ -1031,6 +1040,8 @@ function navigateFromPointer(view: EditorView, e: MouseEvent): boolean {
 
 interface HighlightState {
   deco: DecorationSet;
+  /** The highlighted usages (the decorations also hold whole-line backgrounds). */
+  ranges: readonly { from: number; to: number }[];
   sticky: boolean;
 }
 
@@ -1038,31 +1049,93 @@ type HighlightItem = { from: number; to: number; kind: Highlight["kind"] };
 
 const setHighlights = StateEffect.define<{ items: HighlightItem[]; sticky: boolean }>();
 
+const KIND_CLASS = {
+  text: "cm-lsp-occurrence",
+  read: "cm-lsp-occurrence cm-lsp-occurrence-read",
+  write: "cm-lsp-occurrence cm-lsp-occurrence-write",
+};
 const marks = {
-  text: Decoration.mark({ class: "cm-lsp-occurrence" }),
-  read: Decoration.mark({ class: "cm-lsp-occurrence cm-lsp-occurrence-read" }),
-  write: Decoration.mark({ class: "cm-lsp-occurrence cm-lsp-occurrence-write" }),
+  text: Decoration.mark({ class: KIND_CLASS.text }),
+  read: Decoration.mark({ class: KIND_CLASS.read }),
+  write: Decoration.mark({ class: KIND_CLASS.write }),
+};
+const fullLines = {
+  text: Decoration.line({ class: KIND_CLASS.text }),
+  read: Decoration.line({ class: KIND_CLASS.read }),
+  write: Decoration.line({ class: KIND_CLASS.write }),
 };
 
+/**
+ * Highlights as IDEA paints a text range: one on a single line marks its text; one over
+ * several lines (a `return` statement's exit point) is a block, from its start to the right
+ * edge, whole lines between, and its last line up to its end.
+ */
+function highlightDeco(doc: Text, tabSize: number, items: readonly HighlightItem[]): DecorationSet {
+  const out = [];
+  for (const i of items) {
+    if (i.from >= i.to) continue;
+    const first = doc.lineAt(i.from);
+    const last = doc.lineAt(i.to);
+    if (first.number === last.number) {
+      out.push(marks[i.kind].range(i.from, i.to));
+      continue;
+    }
+    const col = countColumn(first.text.slice(0, i.from - first.from), tabSize);
+    const tail = Decoration.line({
+      class: i.kind === "write" ? "cm-lsp-occurrence-tail cm-lsp-occurrence-tail-write" : "cm-lsp-occurrence-tail",
+      attributes: { style: `--lsp-occurrence-from: calc(${LINE_INSET}px + ${col}ch)` },
+    });
+    out.push(tail.range(first.from));
+    for (let n = first.number + 1; n < last.number; n++) out.push(fullLines[i.kind].range(doc.line(n).from));
+    if (i.to > last.from) out.push(marks[i.kind].range(last.from, i.to));
+  }
+  return Decoration.set(out, true);
+}
+
 const highlightField = StateField.define<HighlightState>({
-  create: () => ({ deco: Decoration.none, sticky: false }),
+  create: () => ({ deco: Decoration.none, ranges: [], sticky: false }),
   update(value, tr) {
     for (const e of tr.effects)
-      if (e.is(setHighlights))
+      if (e.is(setHighlights)) {
+        const items = e.value.items.filter((i) => i.from < i.to);
         return {
-          deco: Decoration.set(
-            e.value.items.filter((i) => i.from < i.to).map((i) => marks[i.kind].range(i.from, i.to)),
-            true,
-          ),
+          deco: highlightDeco(tr.state.doc, tr.state.tabSize, items),
+          ranges: items.map(({ from, to }) => ({ from, to })),
           sticky: e.value.sticky,
         };
+      }
     // Kept (mapped) until recomputed; ⌘⇧F7 highlights survive edits until Escape.
-    if (tr.docChanged) return { deco: value.deco.map(tr.changes), sticky: value.sticky };
+    if (tr.docChanged)
+      return {
+        deco: value.deco.map(tr.changes),
+        ranges: value.ranges.map((r) => ({ from: tr.changes.mapPos(r.from, 1), to: tr.changes.mapPos(r.to, -1) })),
+        sticky: value.sticky,
+      };
     return value;
   },
   provide: (f) => EditorView.decorations.from(f, (v) => v.deco),
 });
 
+/**
+ * The latest usages found per editor, for its current document: a caret landing in one shows
+ * them at once while the server is asked again, as IDEA's highlighting follows the caret.
+ */
+const recentHighlights = new WeakMap<EditorView, { doc: Text; sets: HighlightItem[][] }>();
+const RECENT_SETS = 8;
+
+function rememberHighlights(view: EditorView, doc: Text, items: HighlightItem[]) {
+  if (!items.length) return;
+  let recent = recentHighlights.get(view);
+  if (recent?.doc !== doc) recentHighlights.set(view, (recent = { doc, sets: [] }));
+  recent.sets = [items, ...recent.sets].slice(0, RECENT_SETS);
+}
+
+function recentHighlightsAt(view: EditorView, doc: Text, head: number): HighlightItem[] | null {
+  const recent = recentHighlights.get(view);
+  if (recent?.doc !== doc) return null;
+  return recent.sets.find((set) => set.some((i) => i.from <= head && head <= i.to)) ?? null;
+}
+
 /** IDEA's status bar text while ⌘⇧F7 highlights are shown. */
 export const useUsagesStatus = create<{ text: string | null; owner: EditorView | null }>(() => ({
   text: null,
@@ -1092,6 +1165,7 @@ async function fetchHighlights(view: EditorView, sticky: boolean) {
   const stale = () => view.state.doc !== state.doc || view.state.selection.main.head !== head;
   const apply = (items: HighlightItem[]) => {
     view.dispatch({ effects: setHighlights.of({ items, sticky: sticky && items.length > 0 }) });
+    rememberHighlights(view, state.doc, items);
     if (!sticky) return;
     if (!items.length) {
       if (useUsagesStatus.getState().owner === view) setUsagesStatus(null, null);
@@ -1129,7 +1203,7 @@ async function fetchHighlights(view: EditorView, sticky: boolean) {
   if (lspError) reportLspError(view, lspError, sticky);
   else if (sticky) showHint(view, NO_SERVER);
   // Drop stale marks from an earlier caret position.
-  if (view.state.field(highlightField).deco.size) apply([]);
+  if (view.state.field(highlightField).ranges.length) apply([]);
 }
 
 async function syntaxHighlights(view: EditorView, doc: EditorDocContext, head: number) {
@@ -1148,10 +1222,7 @@ async function syntaxHighlights(view: EditorView, doc: EditorDocContext, head: n
 export function highlightUsages(view: EditorView) {
   const st = view.state.field(highlightField, false);
   const head = view.state.selection.main.head;
-  let onUsage = false;
-  st?.deco.between(head, head, (from, to) => {
-    if (from <= head && head <= to) onUsage = true;
-  });
+  const onUsage = !!st?.ranges.some(({ from, to }) => from <= head && head <= to);
   if (st?.sticky && onUsage) {
     clearHighlights(view);
     return;
@@ -1165,8 +1236,7 @@ export function highlightUsages(view: EditorView) {
  */
 export function gotoHighlightedUsage(view: EditorView, dir: 1 | -1): boolean {
   const st = view.state.field(highlightField, false);
-  const ranges: { from: number; to: number }[] = [];
-  st?.deco.between(0, view.state.doc.length, (from, to) => void ranges.push({ from, to }));
+  const ranges = [...(st?.ranges ?? [])].sort((a, b) => a.from - b.from);
   const head = view.state.selection.main.head;
   const current = ranges.findIndex((r) => r.from <= head && head <= r.to);
   if (ranges.length < 2 || current < 0) return false;
@@ -1178,7 +1248,8 @@ export function gotoHighlightedUsage(view: EditorView, dir: 1 | -1): boolean {
 
 const highlightPlugin = ViewPlugin.fromClass(
   class {
-    private request = debounce(() => void fetchHighlights(this.view, false), 300);
+    // Keyboard moves wait a little, so holding an arrow key does not ask for every line.
+    private request = debounce(() => void fetchHighlights(this.view, false), 100);
     constructor(private view: EditorView) {}
     update(u: ViewUpdate) {
       // Recomputed after caret moves and, like IDEA's daemon, after edits.
@@ -1190,7 +1261,7 @@ const highlightPlugin = ViewPlugin.fromClass(
       }
       const clear = () => {
         this.request.cancel();
-        if (st.deco.size)
+        if (st.ranges.length)
           queueMicrotask(() => this.view.dispatch({ effects: setHighlights.of({ items: [], sticky: false }) }));
       };
       // Code Editing › "Usages of element at caret".
@@ -1198,7 +1269,20 @@ const highlightPlugin = ViewPlugin.fromClass(
         clear();
         return;
       }
-      this.request();
+      const doc = u.state.doc;
+      const head = u.state.selection.main.head;
+      const known = u.docChanged ? null : recentHighlightsAt(this.view, doc, head);
+      if (known)
+        queueMicrotask(() => {
+          const now = this.view.state;
+          if (now.doc === doc && now.selection.main.head === head)
+            this.view.dispatch({ effects: setHighlights.of({ items: known, sticky: false }) });
+        });
+      // A click asks at once (IDEA highlights as the caret lands).
+      if (u.transactions.some((tr) => tr.isUserEvent("select.pointer"))) {
+        this.request.cancel();
+        void fetchHighlights(this.view, false);
+      } else this.request();
     }
     destroy() {
       this.request.cancel();

src/editor/lsp/theme.ts

@@ -288,11 +288,18 @@ export const lspTheme = EditorView.theme({
     textUnderlineOffset: "2px",
     cursor: "pointer",
   },
-  ".cm-lsp-occurrence": {
-    backgroundColor: "var(--bg-identifier-under-caret, color-mix(in srgb, var(--accent) 14%, transparent))",
+  ".cm-lsp-occurrence, .cm-lsp-occurrence-tail": {
+    "--lsp-occurrence-bg": "var(--bg-identifier-under-caret, color-mix(in srgb, var(--accent) 14%, transparent))",
   },
-  ".cm-lsp-occurrence-write": {
-    backgroundColor: "var(--bg-identifier-under-caret-write, color-mix(in srgb, var(--danger) 12%, transparent))",
+  ".cm-lsp-occurrence-write, .cm-lsp-occurrence-tail-write": {
+    "--lsp-occurrence-bg": "var(--bg-identifier-under-caret-write, color-mix(in srgb, var(--danger) 12%, transparent))",
+  },
+  ".cm-lsp-occurrence": { backgroundColor: "var(--lsp-occurrence-bg)" },
+  // First line of a highlight over several lines: from its start column to the right edge,
+  // over the line's own background (the caret row's).
+  ".cm-lsp-occurrence-tail": {
+    backgroundImage:
+      "linear-gradient(to right, transparent var(--lsp-occurrence-from), var(--lsp-occurrence-bg) var(--lsp-occurrence-from))",
   },
   ".cm-lsp-nav-info": { padding: "4px 8px" },
   // Chooser highlight (tempHighlight.ts): the expression / occurrences a choice would take.

src/editor/lsp/usageHighlight.test.ts

@@ -0,0 +1,125 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import type { Highlight, IdeError, Range } from "@/bindings";
+
+type R<T> = Promise<{ status: "ok"; data: T } | { status: "error"; error: IdeError }>;
+const ok = <T>(data: T) => Promise.resolve({ status: "ok" as const, data });
+
+const cmd = {
+  lspHover: vi.fn((): R<unknown> => ok(null)),
+  lspDocumentLinks: vi.fn(() => ok([])),
+  cancelRequest: vi.fn(() => ok(null)),
+  lspDocumentHighlights: vi.fn((): R<Highlight[]> => ok([])),
+};
+
+vi.mock("@/platform/ipc", async (orig) => {
+  const actual = await orig<typeof import("@/platform/ipc")>();
+  return { ...actual, isTauri: () => true, commands: { ...actual.commands, ...cmd } };
+});
+
+const { navigationExtension, highlightUsages } = await import("./navigation");
+const { lspDoc } = await import("./lsp");
+
+const range = (line: number, col: number, endLine: number, endCol: number): Range => ({
+  start: { line, col },
+  end: { line: endLine, col: endCol },
+});
+
+// 0: int f() {          4: }
+// 1:     return a      5: int g(String a) {
+// 2:         .b()      6:     return a;
+// 3:         .c();     7: }
+const DOC = "int f() {\n    return a\n        .b()\n        .c();\n}\nint g(String a) {\n    return a;\n}";
+
+const views: EditorView[] = [];
+function editor(anchor = 0) {
+  const view = new EditorView({
+    state: EditorState.create({
+      doc: DOC,
+      selection: { anchor },
+      extensions: [
+        lspDoc.of({ doc: 7, path: "/p/A.java", language: "java", uiMode: null, readOnly: false }),
+        navigationExtension(),
+      ],
+    }),
+    parent: document.body,
+  });
+  views.push(view);
+  return view;
+}
+const lines = (v: EditorView) => [...v.contentDOM.querySelectorAll<HTMLElement>(".cm-line")];
+const at = (line: number, col: number) => DOC.split("\n").slice(0, line).join("\n").length + (line ? 1 : 0) + col;
+const click = (v: EditorView, pos: number) => v.dispatch({ selection: { anchor: pos }, userEvent: "select.pointer" });
+
+beforeEach(() => {
+  cmd.lspDocumentHighlights.mockReset().mockImplementation(() => ok([]));
+});
+afterEach(() => views.splice(0).forEach((v) => v.destroy()));
+
+describe("highlights of the element at the caret", () => {
+  it("paints a highlight over several lines as one block, as IDEA paints a text range", async () => {
+    cmd.lspDocumentHighlights.mockImplementation(() => ok([{ range: range(1, 4, 3, 13), kind: "text" }]));
+    const v = editor(at(1, 4));
+    highlightUsages(v);
+    await vi.waitFor(() => expect(lines(v)[1]!.classList.contains("cm-lsp-occurrence-tail")).toBe(true));
+    // From `return` (column 4) to the right edge, the whole middle line, the last line up to `;`.
+    expect(lines(v)[1]!.style.getPropertyValue("--lsp-occurrence-from")).toBe("calc(2.5px + 4ch)");
+    expect(lines(v)[2]!.classList.contains("cm-lsp-occurrence")).toBe(true);
+    const last = lines(v)[3]!.querySelector(".cm-lsp-occurrence")!;
+    expect(last.textContent).toBe("        .c();");
+    expect(lines(v)[0]!.className).not.toContain("cm-lsp-occurrence");
+    expect(lines(v)[4]!.className).not.toContain("cm-lsp-occurrence");
+  });
+
+  it("keeps a single-line highlight on its text", async () => {
+    cmd.lspDocumentHighlights.mockImplementation(() => ok([{ range: range(6, 11, 6, 12), kind: "read" }]));
+    const v = editor(at(6, 11));
+    highlightUsages(v);
+    await vi.waitFor(() => expect(v.contentDOM.querySelectorAll(".cm-lsp-occurrence-read")).toHaveLength(1));
+    expect(lines(v)[6]!.classList.contains("cm-lsp-occurrence")).toBe(false);
+  });
+
+  it("asks at once after a click, and after a pause for keyboard moves", async () => {
+    const v = editor(0);
+    click(v, at(6, 11));
+    await vi.waitFor(() => expect(cmd.lspDocumentHighlights).toHaveBeenCalledTimes(1));
+    v.dispatch({ selection: { anchor: at(6, 4) }, userEvent: "select" });
+    await new Promise((r) => setTimeout(r, 20));
+    expect(cmd.lspDocumentHighlights).toHaveBeenCalledTimes(1);
+    await vi.waitFor(() => expect(cmd.lspDocumentHighlights).toHaveBeenCalledTimes(2));
+  });
+
+  it("shows the usages found before at once when the caret comes back to one", async () => {
+    const a = [
+      { range: range(5, 13, 5, 14), kind: "write" as const },
+      { range: range(6, 11, 6, 12), kind: "read" as const },
+    ];
+    cmd.lspDocumentHighlights.mockImplementationOnce(() => ok(a));
+    const v = editor(0);
+    click(v, at(6, 11));
+    await vi.waitFor(() => expect(v.contentDOM.querySelectorAll(".cm-lsp-occurrence")).toHaveLength(2));
+    click(v, at(1, 4));
+    await vi.waitFor(() => expect(v.contentDOM.querySelectorAll(".cm-lsp-occurrence")).toHaveLength(0));
+    // The server is slow now: the remembered usages of `a` show without waiting for it.
+    cmd.lspDocumentHighlights.mockImplementation(() => new Promise(() => {}));
+    click(v, at(5, 13));
+    await vi.waitFor(() => expect(v.contentDOM.querySelectorAll(".cm-lsp-occurrence")).toHaveLength(2));
+  });
+
+  it("forgets them after an edit", async () => {
+    cmd.lspDocumentHighlights.mockImplementationOnce(() => ok([{ range: range(6, 11, 6, 12), kind: "read" }]));
+    const v = editor(0);
+    click(v, at(6, 11));
+    await vi.waitFor(() => expect(v.contentDOM.querySelectorAll(".cm-lsp-occurrence")).toHaveLength(1));
+    v.dispatch({ changes: { from: 0, insert: "\n" } });
+    click(v, 0);
+    await vi.waitFor(() => expect(v.contentDOM.querySelectorAll(".cm-lsp-occurrence")).toHaveLength(0));
+    // Back on `a` with a slow server: the results of the document before the edit are not reused.
+    cmd.lspDocumentHighlights.mockImplementation(() => new Promise(() => {}));
+    click(v, at(6, 11) + 1);
+    await vi.waitFor(() => expect(cmd.lspDocumentHighlights).toHaveBeenCalledTimes(3));
+    await new Promise((r) => setTimeout(r, 20));
+    expect(v.contentDOM.querySelectorAll(".cm-lsp-occurrence")).toHaveLength(0);
+  });
+});