fix(editor): highlight all visible search matches

src/cm/mainEditorExtensions.ts

@@ -1,5 +1,6 @@
 import type { Extension } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
+import searchMatchHighlighter from "./searchMatchHighlighter";
 
 interface MainEditorExtensionOptions {
 	emmetExtensions?: Extension[];
@@ -51,7 +52,9 @@ export function createMainEditorExtensions(
 	pushExtension(extensions, options.multiCursorSelectionExtension);
 	pushExtension(extensions, options.touchSelectionUpdateExtension);
 	pushExtension(extensions, options.quickToolsModifierInputExtension);
-	pushExtension(extensions, options.searchExtension);
+	if (options.searchExtension != null) {
+		extensions.push(options.searchExtension, searchMatchHighlighter);
+	}
 	pushExtension(extensions, options.readOnlyExtension);
 
 	if (options.optionExtensions?.length) {

src/cm/searchMatchHighlighter.ts

@@ -0,0 +1,140 @@
+import {
+	getSearchQuery,
+	searchPanelOpen,
+	type SearchQuery,
+} from "@codemirror/search";
+import { RangeSetBuilder, type EditorState } from "@codemirror/state";
+import type { DecorationSet, ViewUpdate } from "@codemirror/view";
+import { Decoration, EditorView, ViewPlugin } from "@codemirror/view";
+
+interface DocRange {
+	from: number;
+	to: number;
+}
+
+// Keep the same regular-expression look-around used by CodeMirror's built-in
+// search highlighter. Literal searches only need enough context to catch a
+// match that crosses a viewport boundary.
+const REGEXP_SCAN_MARGIN = 250;
+
+const matchMark = Decoration.mark({ class: "cm-searchMatch" });
+const selectedMatchMark = Decoration.mark({
+	class: "cm-searchMatch cm-searchMatch-selected",
+});
+
+function scanMargin(query: SearchQuery): number {
+	return query.regexp ? REGEXP_SCAN_MARGIN : query.search.length;
+}
+
+function sameHighlightQuery(left: SearchQuery, right: SearchQuery): boolean {
+	return (
+		left.search === right.search &&
+		left.caseSensitive === right.caseSensitive &&
+		left.literal === right.literal &&
+		left.regexp === right.regexp &&
+		left.wholeWord === right.wholeWord &&
+		left.test === right.test
+	);
+}
+
+function scanRanges(
+	visibleRanges: readonly DocRange[],
+	margin: number,
+	docLength: number,
+): DocRange[] {
+	const ranges: DocRange[] = [];
+
+	for (const visible of visibleRanges) {
+		const from = Math.max(0, visible.from - margin);
+		const to = Math.min(docLength, visible.to + margin);
+		const previous = ranges[ranges.length - 1];
+
+		if (previous && from <= previous.to) {
+			previous.to = Math.max(previous.to, to);
+		} else {
+			ranges.push({ from, to });
+		}
+	}
+
+	return ranges;
+}
+
+/**
+ * Builds search marks only around rendered document ranges. Keeping this
+ * separate from the view plugin makes the viewport-only behavior testable.
+ */
+export function buildSearchMatchDecorations(
+	state: EditorState,
+	visibleRanges: readonly DocRange[],
+): DecorationSet {
+	const query = getSearchQuery(state);
+	if (!query.search || !query.valid || searchPanelOpen(state)) {
+		return Decoration.none;
+	}
+
+	const builder = new RangeSetBuilder<Decoration>();
+	const ranges = scanRanges(visibleRanges, scanMargin(query), state.doc.length);
+	const selectedRanges = new Map<number, Set<number>>();
+	for (const range of state.selection.ranges) {
+		let ends = selectedRanges.get(range.from);
+		if (!ends) selectedRanges.set(range.from, (ends = new Set()));
+		ends.add(range.to);
+	}
+
+	for (const { from, to } of ranges) {
+		const cursor = query.getCursor(state, from, to);
+		for (let result = cursor.next(); !result.done; result = cursor.next()) {
+			const match = result.value;
+			const selected = selectedRanges.get(match.from)?.has(match.to) === true;
+			builder.add(
+				match.from,
+				match.to,
+				selected ? selectedMatchMark : matchMark,
+			);
+		}
+	}
+
+	return builder.finish();
+}
+
+class SearchMatchHighlighterPlugin {
+	decorations: DecorationSet;
+
+	constructor(view: EditorView) {
+		this.decorations = buildSearchMatchDecorations(
+			view.state,
+			view.visibleRanges,
+		);
+	}
+
+	update(update: ViewUpdate): void {
+		const queryChanged = !sameHighlightQuery(
+			getSearchQuery(update.state),
+			getSearchQuery(update.startState),
+		);
+		const panelChanged =
+			searchPanelOpen(update.state) !== searchPanelOpen(update.startState);
+
+		if (
+			queryChanged ||
+			panelChanged ||
+			update.docChanged ||
+			update.selectionSet ||
+			update.viewportChanged
+		) {
+			this.decorations = buildSearchMatchDecorations(
+				update.state,
+				update.view.visibleRanges,
+			);
+		}
+	}
+}
+
+export const searchMatchHighlighter = ViewPlugin.fromClass(
+	SearchMatchHighlighterPlugin,
+	{
+		decorations: (plugin) => plugin.decorations,
+	},
+);
+
+export default searchMatchHighlighter;

tests/unit/searchMatchHighlighter.test.ts

@@ -0,0 +1,161 @@
+// @vitest-environment happy-dom
+
+import {
+	openSearchPanel,
+	search,
+	SearchQuery,
+	setSearchQuery,
+} from "@codemirror/search";
+import { EditorSelection, EditorState } from "@codemirror/state";
+import { type DecorationSet, EditorView } from "@codemirror/view";
+import searchMatchHighlighter, {
+	buildSearchMatchDecorations,
+} from "cm/searchMatchHighlighter";
+import { afterEach, describe, expect, it } from "vitest";
+
+interface HighlightRange {
+	from: number;
+	to: number;
+	className: string;
+}
+
+const views: EditorView[] = [];
+
+afterEach(() => {
+	while (views.length) views.pop()?.destroy();
+	document.body.replaceChildren();
+});
+
+function searchState(
+	doc: string,
+	searchText: string,
+	selection = EditorSelection.cursor(0),
+	options: Partial<ConstructorParameters<typeof SearchQuery>[0]> = {},
+): EditorState {
+	const state = EditorState.create({
+		doc,
+		selection,
+		extensions: search(),
+	});
+	return state.update({
+		effects: setSearchQuery.of(
+			new SearchQuery({ search: searchText, ...options }),
+		),
+	}).state;
+}
+
+function highlightRanges(
+	state: EditorState,
+	visibleRanges: readonly { from: number; to: number }[],
+): HighlightRange[] {
+	return decorationRanges(
+		buildSearchMatchDecorations(state, visibleRanges),
+		state.doc.length,
+	);
+}
+
+function decorationRanges(
+	decorations: DecorationSet,
+	docLength: number,
+): HighlightRange[] {
+	const ranges: HighlightRange[] = [];
+	decorations.between(0, docLength, (from, to, decoration) => {
+		ranges.push({
+			from,
+			to,
+			className: String(decoration.spec.class ?? ""),
+		});
+	});
+	return ranges;
+}
+
+describe("custom search match highlighter", () => {
+	it("reacts to custom query and selection updates without a native panel", () => {
+		const view = new EditorView({
+			state: EditorState.create({
+				doc: "foo foo",
+				extensions: [search(), searchMatchHighlighter],
+			}),
+			parent: document.body,
+		});
+		views.push(view);
+
+		view.dispatch({
+			effects: setSearchQuery.of(new SearchQuery({ search: "foo" })),
+		});
+		view.dispatch({ selection: EditorSelection.range(4, 7) });
+
+		const plugin = view.plugin(searchMatchHighlighter);
+		expect(plugin).not.toBeNull();
+		expect(
+			decorationRanges(plugin!.decorations, view.state.doc.length),
+		).toEqual([
+			{ from: 0, to: 3, className: "cm-searchMatch" },
+			{
+				from: 4,
+				to: 7,
+				className: "cm-searchMatch cm-searchMatch-selected",
+			},
+		]);
+	});
+
+	it("highlights every match in the rendered range and marks the selection", () => {
+		const state = searchState(
+			"foo foo foo",
+			"foo",
+			EditorSelection.range(4, 7),
+		);
+
+		expect(highlightRanges(state, [{ from: 0, to: state.doc.length }])).toEqual(
+			[
+				{ from: 0, to: 3, className: "cm-searchMatch" },
+				{
+					from: 4,
+					to: 7,
+					className: "cm-searchMatch cm-searchMatch-selected",
+				},
+				{ from: 8, to: 11, className: "cm-searchMatch" },
+			],
+		);
+	});
+
+	it("limits work to the viewport instead of scanning a large document", () => {
+		const padding = "x".repeat(100_000);
+		const doc = `needle${padding}needle${padding}needle`;
+		const middleMatch = 6 + padding.length;
+		const state = searchState(doc, "needle");
+
+		expect(
+			highlightRanges(state, [
+				{ from: middleMatch - 10, to: middleMatch + 16 },
+			]),
+		).toEqual([
+			{ from: middleMatch, to: middleMatch + 6, className: "cm-searchMatch" },
+		]);
+	});
+
+	it("merges nearby viewport windows so matches are not decorated twice", () => {
+		const state = searchState("foo foo foo", "foo");
+
+		expect(
+			highlightRanges(state, [
+				{ from: 0, to: 5 },
+				{ from: 6, to: 11 },
+			]),
+		).toHaveLength(3);
+	});
+
+	it("leaves highlighting to CodeMirror when its native panel is open", () => {
+		const view = new EditorView({
+			state: searchState("foo foo", "foo"),
+			parent: document.body,
+		});
+		views.push(view);
+
+		openSearchPanel(view);
+
+		expect(
+			highlightRanges(view.state, [{ from: 0, to: view.state.doc.length }]),
+		).toEqual([]);
+	});
+});