fix(editor): make fold all reliable for large files

src/cm/foldingCommands.js

@@ -6,13 +6,20 @@ import {
 	foldNodeProp,
 	foldService,
 	foldState,
+	forceParsing,
+	language,
+	StreamLanguage,
 	syntaxTree,
+	syntaxTreeAvailable,
 	unfoldEffect,
 } from "@codemirror/language";
 import { StateEffect } from "@codemirror/state";
 
-const FULL_PARSE_BUDGET_MS = 100;
-const INCOMPLETE_PARSE_MARGIN = 50;
+const INITIAL_PARSE_BUDGET_MS = 12;
+const IDLE_PARSE_BUDGET_MS = 10;
+const MIN_IDLE_PARSE_BUDGET_MS = 2;
+const IDLE_CALLBACK_TIMEOUT_MS = 50;
+const pendingFoldAllJobs = new WeakMap();
 
 function findServiceFold(state, line) {
 	for (const service of state.facet(foldService)) {
@@ -22,68 +29,104 @@ function findServiceFold(state, line) {
 	return null;
 }
 
-function addSyntaxFolds(state, tree, line, addRange) {
-	if (!tree || tree.length < line.to) return;
+function scheduleIdleWork(callback) {
+	if (
+		typeof window !== "undefined" &&
+		typeof window.requestIdleCallback === "function"
+	) {
+		return {
+			kind: "idle",
+			id: window.requestIdleCallback(callback, {
+				timeout: IDLE_CALLBACK_TIMEOUT_MS,
+			}),
+		};
+	}
 
-	for (let iter = tree.resolveStack(line.to, 1); iter; iter = iter.next) {
-		const node = iter.node;
-		if (node.to <= line.to || node.from > line.to) continue;
-		const lastChild = node.lastChild;
-		if (
-			tree.length !== state.doc.length &&
-			node.to >= tree.length - INCOMPLETE_PARSE_MARGIN &&
-			lastChild?.to === node.to &&
-			lastChild.type.isError
-		) {
-			continue;
-		}
+	return {
+		kind: "timeout",
+		id: setTimeout(() => callback(null), 0),
+	};
+}
 
-		const fold = node.type.prop(foldNodeProp);
-		if (!fold) continue;
+function cancelIdleWork(handle) {
+	if (!handle) return;
+	if (
+		handle.kind === "idle" &&
+		typeof window !== "undefined" &&
+		typeof window.cancelIdleCallback === "function"
+	) {
+		window.cancelIdleCallback(handle.id);
+	} else {
+		clearTimeout(handle.id);
+	}
+}
 
-		const range = fold(node, state);
-		if (
-			range &&
-			range.from >= line.from &&
-			range.from <= line.to &&
-			range.to > line.to
-		) {
-			addRange(range);
-		}
+function cancelPendingFoldAll(view) {
+	const job = pendingFoldAllJobs.get(view);
+	if (!job) return;
+	job.cancelled = true;
+	cancelIdleWork(job.handle);
+	pendingFoldAllJobs.delete(view);
+}
+
+function getIdleParseBudget(deadline) {
+	if (!deadline || typeof deadline.timeRemaining !== "function") {
+		return IDLE_PARSE_BUDGET_MS;
 	}
+	return Math.max(
+		MIN_IDLE_PARSE_BUDGET_MS,
+		Math.min(IDLE_PARSE_BUDGET_MS, Math.floor(deadline.timeRemaining())),
+	);
 }
 
-/**
- * Fold every foldable block, including nested blocks. CodeMirror's built-in
- * foldAll intentionally skips nested ranges after finding a top-level fold.
- */
-export function foldAllCodeBlocks(view) {
-	const { state } = view;
-	const tree =
-		ensureSyntaxTree(state, state.doc.length, FULL_PARSE_BUDGET_MS) ??
-		syntaxTree(state);
+function addSyntaxFolds(state, tree, addRange) {
+	if (!tree || tree.length < state.doc.length) return;
+
+	tree.iterate({
+		enter(ref) {
+			const node = ref.node;
+			const fold = node.type.prop(foldNodeProp);
+			if (!fold) return;
+
+			const range = fold(node, state);
+			if (!range) return;
 
+			const line = state.doc.lineAt(range.from);
+			if (range.from <= line.to && range.to > line.to) addRange(range);
+		},
+	});
+}
+
+function collectFoldEffects(state, tree) {
 	const existing = new Set();
 	foldedRanges(state).between(0, state.doc.length, (from, to) => {
 		existing.add(`${from}:${to}`);
 	});
 
-	const effects = [];
+	const ranges = [];
 	const discovered = new Set();
 	const addRange = (range) => {
 		const id = `${range.from}:${range.to}`;
 		if (existing.has(id) || discovered.has(id)) return;
 		discovered.add(id);
-		effects.push(foldEffect.of(range));
+		ranges.push(range);
 	};
 
 	for (let lineNumber = 1; lineNumber <= state.doc.lines; lineNumber += 1) {
 		const line = state.doc.line(lineNumber);
 		const serviceRange = findServiceFold(state, line);
 		if (serviceRange) addRange(serviceRange);
-		else addSyntaxFolds(state, tree, line, addRange);
 	}
+	addSyntaxFolds(state, tree, addRange);
+
+	return ranges
+		.sort((a, b) => a.from - b.from || a.to - b.to)
+		.map((range) => foldEffect.of(range));
+}
 
+function applyAllFolds(view, tree) {
+	const { state } = view;
+	const effects = collectFoldEffects(state, tree);
 	if (!effects.length) return false;
 	if (!state.field(foldState, false)) {
 		// Install the state field first so it can consume the fold effects in
@@ -94,8 +137,75 @@ export function foldAllCodeBlocks(view) {
 	return true;
 }
 
+function scheduleProgressiveFoldAll(view, document) {
+	const job = {
+		cancelled: false,
+		document,
+		handle: null,
+	};
+	pendingFoldAllJobs.set(view, job);
+
+	const continueParsing = (deadline) => {
+		job.handle = null;
+		if (
+			job.cancelled ||
+			pendingFoldAllJobs.get(view) !== job ||
+			view.state.doc !== document
+		) {
+			pendingFoldAllJobs.delete(view);
+			return;
+		}
+
+		try {
+			const complete =
+				syntaxTreeAvailable(view.state, document.length) ||
+				forceParsing(view, document.length, getIdleParseBudget(deadline));
+
+			if (complete) {
+				pendingFoldAllJobs.delete(view);
+				applyAllFolds(view, syntaxTree(view.state));
+				return;
+			}
+		} catch (error) {
+			console.error("Failed to finish parsing for Fold all.", error);
+			pendingFoldAllJobs.delete(view);
+			return;
+		}
+
+		job.handle = scheduleIdleWork(continueParsing);
+	};
+
+	job.handle = scheduleIdleWork(continueParsing);
+}
+
+/**
+ * Fold every foldable block, including nested blocks. Small and already-parsed
+ * documents fold synchronously. Large documents finish parsing in short idle
+ * slices before all ranges are collected in one pass and one transaction.
+ */
+export function foldAllCodeBlocks(view) {
+	cancelPendingFoldAll(view);
+	const { state } = view;
+	const activeLanguage = state.facet(language);
+
+	if (!activeLanguage || activeLanguage instanceof StreamLanguage) {
+		return applyAllFolds(view, syntaxTree(state));
+	}
+
+	const tree = ensureSyntaxTree(
+		state,
+		state.doc.length,
+		INITIAL_PARSE_BUDGET_MS,
+	);
+	if (tree) return applyAllFolds(view, tree);
+
+	scheduleProgressiveFoldAll(view, state.doc);
+	return true;
+}
+
 /** Unfold every stored fold, including nested folds created above. */
 export function unfoldAllCodeBlocks(view) {
+	cancelPendingFoldAll(view);
 	const effects = [];
 	foldedRanges(view.state).between(0, view.state.doc.length, (from, to) => {
 		effects.push(unfoldEffect.of({ from, to }));

tests/unit/foldingCommands.test.js

@@ -0,0 +1,150 @@
+// @vitest-environment happy-dom
+
+import { javascript } from "@codemirror/lang-javascript";
+import { codeFolding, foldedRanges } from "@codemirror/language";
+import { EditorState } from "@codemirror/state";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+const parsing = vi.hoisted(() => ({
+	ensureSyntaxTree: vi.fn(),
+	forceParsing: vi.fn(),
+	syntaxTreeAvailable: vi.fn(),
+}));
+
+vi.mock("@codemirror/language", async (importOriginal) => ({
+	...(await importOriginal()),
+	ensureSyntaxTree: parsing.ensureSyntaxTree,
+	forceParsing: parsing.forceParsing,
+	syntaxTreeAvailable: parsing.syntaxTreeAvailable,
+}));
+
+const { foldAllCodeBlocks, unfoldAllCodeBlocks } = await import(
+	"cm/foldingCommands"
+);
+
+function createView() {
+	let state = EditorState.create({
+		doc: [
+			"function outer() {",
+			"\tif (ready) {",
+			'\t\tconsole.log("ready");',
+			"\t}",
+			"}",
+		].join("\n"),
+		extensions: [javascript(), codeFolding()],
+	});
+
+	return {
+		get state() {
+			return state;
+		},
+		dispatch(spec) {
+			state = state.update(spec).state;
+		},
+	};
+}
+
+function countFoldedRanges(state) {
+	let count = 0;
+	foldedRanges(state).between(0, state.doc.length, () => {
+		count += 1;
+	});
+	return count;
+}
+
+describe("progressive Fold all", () => {
+	let idleCallbacks;
+	let nextIdleId;
+
+	beforeEach(() => {
+		idleCallbacks = new Map();
+		nextIdleId = 1;
+		window.requestIdleCallback = vi.fn((callback) => {
+			const id = nextIdleId;
+			nextIdleId += 1;
+			idleCallbacks.set(id, callback);
+			return id;
+		});
+		window.cancelIdleCallback = vi.fn((id) => {
+			idleCallbacks.delete(id);
+		});
+
+		parsing.ensureSyntaxTree.mockReset().mockReturnValue(null);
+		parsing.forceParsing.mockReset();
+		parsing.syntaxTreeAvailable.mockReset().mockReturnValue(false);
+	});
+
+	afterEach(() => {
+		vi.restoreAllMocks();
+	});
+
+	function takeNextIdleCallback() {
+		const next = idleCallbacks.entries().next().value;
+		expect(next).toBeDefined();
+		const [id, callback] = next;
+		idleCallbacks.delete(id);
+		return callback;
+	}
+
+	function runNextIdleCallback() {
+		takeNextIdleCallback()({
+			didTimeout: false,
+			timeRemaining: () => 8,
+		});
+	}
+
+	it("folds after parsing completes across multiple idle slices", () => {
+		const view = createView();
+		parsing.forceParsing
+			.mockReturnValueOnce(false)
+			.mockReturnValueOnce(false)
+			.mockReturnValueOnce(true);
+
+		expect(foldAllCodeBlocks(view)).toBe(true);
+		expect(countFoldedRanges(view.state)).toBe(0);
+
+		runNextIdleCallback();
+		runNextIdleCallback();
+		expect(countFoldedRanges(view.state)).toBe(0);
+
+		runNextIdleCallback();
+		expect(parsing.forceParsing).toHaveBeenCalledTimes(3);
+		expect(countFoldedRanges(view.state)).toBeGreaterThan(0);
+		expect(idleCallbacks.size).toBe(0);
+	});
+
+	it("cancels pending parsing when the document changes", () => {
+		const view = createView();
+		parsing.forceParsing.mockReturnValue(false);
+
+		foldAllCodeBlocks(view);
+		runNextIdleCallback();
+		expect(parsing.forceParsing).toHaveBeenCalledTimes(1);
+
+		view.dispatch({ changes: { from: 0, insert: "// changed\n" } });
+		runNextIdleCallback();
+
+		expect(parsing.forceParsing).toHaveBeenCalledTimes(1);
+		expect(countFoldedRanges(view.state)).toBe(0);
+		expect(idleCallbacks.size).toBe(0);
+	});
+
+	it("cancels pending parsing when Unfold all runs", () => {
+		const view = createView();
+		parsing.forceParsing.mockReturnValue(false);
+
+		foldAllCodeBlocks(view);
+		const pendingCallback = takeNextIdleCallback();
+		expect(unfoldAllCodeBlocks(view)).toBe(false);
+
+		pendingCallback({
+			didTimeout: false,
+			timeRemaining: () => 8,
+		});
+
+		expect(window.cancelIdleCallback).toHaveBeenCalledTimes(1);
+		expect(parsing.forceParsing).not.toHaveBeenCalled();
+		expect(countFoldedRanges(view.state)).toBe(0);
+		expect(idleCallbacks.size).toBe(0);
+	});
+});