Faster codemirror search match counting (#612)

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

apps/yaak-client/components/core/Editor/searchMatchCount.test.ts

@@ -0,0 +1,202 @@
+import { SearchQuery } from "@codemirror/search";
+import { EditorState } from "@codemirror/state";
+import { describe, expect, test } from "vite-plus/test";
+import {
+  currentMatch,
+  literalSearch,
+  MAX_COUNT,
+  MatchCounter,
+  normalizeDoc,
+  normalizeSearch,
+  scanNormalized,
+  scanQuery,
+} from "./searchMatchCount";
+
+type QueryConfig = ConstructorParameters<typeof SearchQuery>[0];
+
+const stateOf = (doc: string) => EditorState.create({ doc });
+
+/**
+ * The matches the counter finds, having checked them against the search panel's own cursor.
+ *
+ * The cursor decides which ranges the editor highlights and which one `find next` lands on, so
+ * a count that doesn't agree with it is a wrong count, however fast it was to produce.
+ */
+function matchesOf(doc: string, config: QueryConfig) {
+  const state = stateOf(doc);
+  const query = new SearchQuery(config);
+  const matches = new MatchCounter().matches(state, query);
+  expect(matches).toEqual(scanQuery(state, query));
+  return matches;
+}
+
+const countOf = (doc: string, config: QueryConfig) => matchesOf(doc, config).length;
+
+describe("counting", () => {
+  test("counts every match, whatever the case", () => {
+    expect(countOf("one Two three two", { search: "two" })).toBe(2);
+    expect(countOf("one Two three two", { search: "two", caseSensitive: true })).toBe(1);
+  });
+
+  test("skips matches overlapping an earlier one", () => {
+    expect(countOf("aaaaa", { search: "aa" })).toBe(2);
+    expect(countOf("ababa", { search: "aba" })).toBe(1);
+  });
+
+  test("treats a query as text, not as a pattern", () => {
+    expect(countOf("a.b axb", { search: "a.b" })).toBe(1);
+  });
+
+  test("unquotes escapes unless the query is literal", () => {
+    expect(countOf("one\ntwo\nthree", { search: "\\n" })).toBe(2);
+    expect(countOf("one\\ntwo", { search: "\\n", literal: true })).toBe(1);
+  });
+
+  test("counts regexp and whole word queries through the cursor", () => {
+    expect(literalSearch(new SearchQuery({ search: "a", regexp: true }))).toBe(null);
+    expect(literalSearch(new SearchQuery({ search: "a", wholeWord: true }))).toBe(null);
+    expect(countOf("a1 b2 c3", { search: "[a-z]\\d", regexp: true })).toBe(3);
+    expect(countOf("cat cats cat", { search: "cat", wholeWord: true })).toBe(2);
+  });
+
+  test("stops counting at the cap", () => {
+    expect(countOf("x".repeat(MAX_COUNT + 100), { search: "x" })).toBe(MAX_COUNT + 1);
+  });
+
+  test("reports where the matches are", () => {
+    expect(matchesOf("ab..ab", { search: "ab" })).toEqual([
+      { from: 0, to: 2 },
+      { from: 4, to: 6 },
+    ]);
+  });
+
+  test("finds nothing to match with an empty needle", () => {
+    expect(scanNormalized(normalizeDoc("abc", false), "")).toEqual([]);
+  });
+});
+
+describe("normalization", () => {
+  test("finds what a character decomposes into", () => {
+    // The é is one character holding an `e`, and the match covers the whole of it
+    expect(matchesOf("café", { search: "e" })).toEqual([{ from: 3, to: 4 }]);
+    expect(matchesOf("file", { search: "fi" })).toEqual([{ from: 0, to: 1 }]);
+    expect(matchesOf("a…b", { search: "..." })).toEqual([{ from: 1, to: 2 }]);
+    expect(countOf("one two", { search: "one two" })).toBe(1);
+    expect(countOf("ｆｕｌｌ width", { search: "full" })).toBe(1);
+  });
+
+  test("matches a decomposed query against composed text, and the reverse", () => {
+    expect(countOf("café", { search: "café" })).toBe(1);
+    expect(countOf("café", { search: "café" })).toBe(1);
+    expect(countOf("café", { search: "café" })).toBe(1);
+  });
+
+  test("keeps offsets straight after an expansion", () => {
+    expect(matchesOf("é.é.end", { search: "end" })).toEqual([{ from: 4, to: 7 }]);
+    expect(matchesOf("fififi stop", { search: "stop" })).toEqual([{ from: 4, to: 8 }]);
+  });
+
+  test("normalizes the query whole, the document by character", () => {
+    expect(normalizeSearch("CAFÉ", false)).toBe("café");
+    expect(normalizeSearch("CAFÉ", true)).toBe("CAFÉ");
+    // Whole-string NFKD would fold this to a final sigma, which the cursor never does
+    expect(normalizeDoc("ΟΔΟΣ", false).text).toBe("οδοσ");
+  });
+
+  test("leaves a document that normalizes to itself untouched", () => {
+    const { text, expansions } = normalizeDoc("plain 日本 🎉 text", false);
+    expect(text).toBe("plain 日本 🎉 text");
+    expect(expansions).toEqual([]);
+  });
+});
+
+describe("current match", () => {
+  const matches = [
+    { from: 0, to: 2 },
+    { from: 4, to: 6 },
+    { from: 8, to: 10 },
+  ];
+
+  test("counts from one, and reports 0 off a match", () => {
+    expect(currentMatch(matches, { from: 4, to: 6 })).toBe(2);
+    expect(currentMatch(matches, { from: 8, to: 10 })).toBe(3);
+    expect(currentMatch(matches, { from: 5, to: 5 })).toBe(2);
+    expect(currentMatch(matches, { from: 2, to: 3 })).toBe(0);
+    expect(currentMatch(matches, { from: 4, to: 7 })).toBe(0);
+    expect(currentMatch([], { from: 0, to: 0 })).toBe(0);
+  });
+
+  test("moving the selection doesn't scan again", () => {
+    const state = stateOf("a1 b2 c3");
+    const query = new SearchQuery({ search: "\\d", regexp: true });
+    const counter = new MatchCounter();
+    const found = counter.matches(state, query);
+
+    // The document a selection-only transaction leaves behind is the one already scanned
+    const moved = state.update({ selection: { anchor: 4, head: 5 } }).state;
+    expect(counter.matches(moved, query)).toBe(found);
+    expect(currentMatch(found, moved.selection.main)).toBe(2);
+  });
+});
+
+/**
+ * The mapping from normalized offsets back to document offsets is the part of this that can go
+ * quietly wrong, and only on input nobody thinks to write a case for. So generate the input.
+ */
+describe("against the cursor, on awkward text", () => {
+  const ALPHABET = [
+    ..."abcABC .\\\n".split(""),
+    "é",
+    "é",
+    "fi",
+    "…",
+    " ",
+    "İ",
+    "Σ",
+    "ς",
+    "日",
+    "🎉",
+    "Ⅻ",
+    "ｆ",
+    "①",
+    "́",
+  ];
+
+  /** Seeded, so a failure is the same failure next run */
+  function random(seed: number) {
+    let state = seed;
+    return () => {
+      state = (state * 1664525 + 1013904223) >>> 0;
+      return state / 2 ** 32;
+    };
+  }
+
+  for (const caseSensitive of [false, true]) {
+    test(`agrees on every generated document (caseSensitive: ${caseSensitive})`, () => {
+      const next = random(caseSensitive ? 20260831 : 7);
+
+      for (let round = 0; round < 400; round++) {
+        const doc = Array.from(
+          { length: 2 + Math.floor(next() * 60) },
+          () => ALPHABET[Math.floor(next() * ALPHABET.length)]!,
+        ).join("");
+
+        // Half the queries are lifted out of the document, so matches are actually found
+        const start = Math.floor(next() * doc.length);
+        const search =
+          next() < 0.5
+            ? doc.slice(start, start + 1 + Math.floor(next() * 3))
+            : Array.from(
+                { length: 1 + Math.floor(next() * 2) },
+                () => ALPHABET[Math.floor(next() * ALPHABET.length)]!,
+              ).join("");
+        if (search === "") continue;
+
+        const state = stateOf(doc);
+        const query = new SearchQuery({ search, caseSensitive });
+        const where = `doc=${JSON.stringify(doc)} search=${JSON.stringify(search)}`;
+        expect(new MatchCounter().matches(state, query), where).toEqual(scanQuery(state, query));
+      }
+    });
+  }
+});

apps/yaak-client/components/core/Editor/searchMatchCount.ts

@@ -1,7 +1,232 @@
-import { getSearchQuery, searchPanelOpen } from "@codemirror/search";
-import type { Extension } from "@codemirror/state";
+import { getSearchQuery, type SearchQuery, searchPanelOpen } from "@codemirror/search";
+import type { EditorState, Extension, Text } from "@codemirror/state";
 import { type EditorView, ViewPlugin, type ViewUpdate } from "@codemirror/view";
 
+/** Matches are counted no further than this, since an exact total stops being useful long before */
+export const MAX_COUNT = 9999;
+
+/** What normalizing rewrites: anything outside ASCII, plus the case it folds */
+const REWRITTEN = /\P{ASCII}|[A-Z]+/gu;
+const REWRITTEN_CASE_SENSITIVE = /\P{ASCII}/gu;
+
+export interface Match {
+  from: number;
+  to: number;
+}
+
+/** A character whose normalized form is a different length, shifting every offset past it */
+interface Expansion {
+  normFrom: number;
+  normTo: number;
+  docFrom: number;
+  docTo: number;
+}
+
+/** A document as SearchCursor compares it, with what's needed to get back to real offsets */
+export interface NormalizedDoc {
+  text: string;
+  expansions: Expansion[];
+}
+
+/**
+ * Rewrites a document the way SearchCursor does — NFKD, then a case fold unless the search is
+ * case-sensitive — in a single pass rather than one call per code point.
+ *
+ * The cursor spends 90% of its time asking ICU about one character at a time, which is what
+ * makes counting matches in a large response slow. Doing it a character at a time still matters
+ * for the result, since it keeps NFKD from reordering marks across characters, so the
+ * granularity stays and only the repeated work goes: each distinct character is normalized once
+ * and the answer reused, and ASCII runs never reach ICU at all.
+ */
+export function normalizeDoc(text: string, caseSensitive: boolean): NormalizedDoc {
+  const rewritten = new Map<string, string>();
+  const expansions: Expansion[] = [];
+  let shift = 0;
+
+  const normalized = text.replace(
+    caseSensitive ? REWRITTEN_CASE_SENSITIVE : REWRITTEN,
+    (chunk: string, at: number) => {
+      // An ASCII run only ever folds case, which can't change its length
+      if (chunk.charCodeAt(0) < 0x80) return chunk.toLowerCase();
+
+      let out = rewritten.get(chunk);
+      if (out === undefined) {
+        out = chunk.normalize("NFKD");
+        if (!caseSensitive) out = out.toLowerCase();
+        rewritten.set(chunk, out);
+      }
+
+      if (out.length !== chunk.length) {
+        expansions.push({
+          normFrom: at + shift,
+          normTo: at + shift + out.length,
+          docFrom: at,
+          docTo: at + chunk.length,
+        });
+        shift += out.length - chunk.length;
+      }
+
+      return out;
+    },
+  );
+
+  return { text: normalized, expansions };
+}
+
+/** The query as SearchCursor compares it, which it normalizes whole rather than by character */
+export function normalizeSearch(search: string, caseSensitive: boolean): string {
+  const normalized = search.normalize("NFKD");
+  return caseSensitive ? normalized : normalized.toLowerCase();
+}
+
+/** Every occurrence of `needle`, skipping matches that overlap an earlier one */
+export function scanNormalized(doc: NormalizedDoc, needle: string): Match[] {
+  const matches: Match[] = [];
+  if (needle === "") return matches;
+
+  const { text } = doc;
+  let pos = text.indexOf(needle);
+  while (pos >= 0) {
+    let end = pos + needle.length;
+    // However the query was cut, a match ends on a whole code point, as the cursor's do
+    if (isLowSurrogate(text.charCodeAt(end))) end++;
+
+    matches.push({ from: docStart(doc, pos), to: docEnd(doc, end) });
+    if (matches.length > MAX_COUNT) break;
+
+    pos = text.indexOf(needle, resumeAfter(doc, end));
+  }
+
+  return matches;
+}
+
+/** The same through the query's own cursor, which handles regexps and whole words */
+export function scanQuery(state: EditorState, query: SearchQuery): Match[] {
+  const matches: Match[] = [];
+  const cursor = query.getCursor(state);
+
+  for (let result = cursor.next(); !result.done; result = cursor.next()) {
+    matches.push({ from: result.value.from, to: result.value.to });
+    if (matches.length > MAX_COUNT) break;
+  }
+
+  return matches;
+}
+
+const isLowSurrogate = (code: number) => code >= 0xdc00 && code <= 0xdfff;
+
+/** The last character expansion beginning at or before `offset`, if there is one */
+function expansionAt({ expansions }: NormalizedDoc, offset: number): Expansion | null {
+  let low = 0;
+  let high = expansions.length - 1;
+  let found: Expansion | null = null;
+
+  while (low <= high) {
+    const mid = (low + high) >> 1;
+    if (expansions[mid]!.normFrom <= offset) {
+      found = expansions[mid]!;
+      low = mid + 1;
+    } else {
+      high = mid - 1;
+    }
+  }
+
+  return found;
+}
+
+function docStart(doc: NormalizedDoc, offset: number): number {
+  const expansion = expansionAt(doc, offset);
+  if (expansion == null) return offset;
+  // A match starting inside a character's expansion starts at the character
+  return offset < expansion.normTo
+    ? expansion.docFrom
+    : offset - (expansion.normTo - expansion.docTo);
+}
+
+/**
+ * Where scanning picks up after a match ending at `offset`.
+ *
+ * The cursor moves through the document a character at a time, so once a match ends inside a
+ * character's expansion the rest of that expansion is behind it — "…" holds three dots but only
+ * ever counts as one match of ".".
+ */
+function resumeAfter(doc: NormalizedDoc, offset: number): number {
+  const expansion = expansionAt(doc, offset);
+  return expansion != null && offset > expansion.normFrom && offset < expansion.normTo
+    ? expansion.normTo
+    : offset;
+}
+
+function docEnd(doc: NormalizedDoc, offset: number): number {
+  const expansion = expansionAt(doc, offset);
+  if (expansion == null) return offset;
+  if (offset <= expansion.normFrom) return expansion.docFrom;
+  // A match ending inside a character's expansion covers the whole character
+  return offset < expansion.normTo
+    ? expansion.docTo
+    : offset - (expansion.normTo - expansion.docTo);
+}
+
+/** Position of the match holding the selection, counting from one, or 0 when it isn't on one */
+export function currentMatch(matches: Match[], selection: { from: number; to: number }): number {
+  let index = 0;
+  for (const match of matches) {
+    index++;
+    if (match.from <= selection.from && match.to >= selection.to) return index;
+  }
+  return 0;
+}
+
+/** The text a plain query looks for, or null when only the cursor can answer it */
+export function literalSearch(query: SearchQuery): string | null {
+  if (query.regexp || query.wholeWord || query.test != null) return null;
+  // Mirrors SearchQuery's own unquoting, which the published type doesn't expose
+  return query.literal
+    ? query.search
+    : query.search.replace(/\\([nrt\\])/g, (_, ch) =>
+        ch === "n" ? "\n" : ch === "r" ? "\r" : ch === "t" ? "\t" : "\\",
+      );
+}
+
+/**
+ * Finds the matches for the search panel, keeping the normalized document and the matches it
+ * last found, so neither moving the selection nor typing another character starts over.
+ */
+export class MatchCounter {
+  private doc: { doc: Text; caseSensitive: boolean; normalized: NormalizedDoc } | null = null;
+  private last: { doc: Text; query: SearchQuery; matches: Match[] } | null = null;
+
+  matches(state: EditorState, query: SearchQuery): Match[] {
+    const last = this.last;
+    if (last != null && last.doc === state.doc && last.query.eq(query)) {
+      return last.matches;
+    }
+
+    const matches = this.scan(state, query);
+    this.last = { doc: state.doc, query, matches };
+    return matches;
+  }
+
+  private scan(state: EditorState, query: SearchQuery): Match[] {
+    const search = literalSearch(query);
+    if (search == null) return scanQuery(state, query);
+
+    const doc = this.normalizedDoc(state.doc, query.caseSensitive);
+    return scanNormalized(doc, normalizeSearch(search, query.caseSensitive));
+  }
+
+  private normalizedDoc(doc: Text, caseSensitive: boolean): NormalizedDoc {
+    const cached = this.doc;
+    if (cached != null && cached.doc === doc && cached.caseSensitive === caseSensitive) {
+      return cached.normalized;
+    }
+
+    const normalized = normalizeDoc(doc.toString(), caseSensitive);
+    this.doc = { doc, caseSensitive, normalized };
+    return normalized;
+  }
+}
+
 /**
  * A CodeMirror extension that displays the total number of search matches
  * inside the built-in search panel.
@@ -10,6 +235,7 @@ export function searchMatchCount(): Extension {
   return ViewPlugin.fromClass(
     class {
       private countEl: HTMLElement | null = null;
+      private counter = new MatchCounter();
 
       constructor(private view: EditorView) {
         this.updateCount();
@@ -38,38 +264,21 @@ export function searchMatchCount(): Extension {
         }
 
         this.ensureCountEl();
+        if (this.countEl == null) return;
 
         if (!query.search) {
-          if (this.countEl) {
-            this.countEl.textContent = "0/0";
-          }
+          this.countEl.textContent = "0/0";
           return;
         }
 
-        const selection = state.selection.main;
-        let count = 0;
-        let currentIndex = 0;
-        const MAX_COUNT = 9999;
-        const cursor = query.getCursor(state);
-        for (let result = cursor.next(); !result.done; result = cursor.next()) {
-          count++;
-          const match = result.value;
-          if (match.from <= selection.from && match.to >= selection.to) {
-            currentIndex = count;
-          }
-          if (count > MAX_COUNT) break;
-        }
-
-        if (this.countEl) {
-          if (count > MAX_COUNT) {
-            this.countEl.textContent = `${MAX_COUNT}+`;
-          } else if (count === 0) {
-            this.countEl.textContent = "0/0";
-          } else if (currentIndex > 0) {
-            this.countEl.textContent = `${currentIndex}/${count}`;
-          } else {
-            this.countEl.textContent = `0/${count}`;
-          }
+        const matches = this.counter.matches(state, query);
+        if (matches.length > MAX_COUNT) {
+          this.countEl.textContent = `${MAX_COUNT}+`;
+        } else if (matches.length === 0) {
+          this.countEl.textContent = "0/0";
+        } else {
+          const current = currentMatch(matches, state.selection.main);
+          this.countEl.textContent = `${current}/${matches.length}`;
         }
       }