fix fold all command same line thing

Author: bajrangCoder

src/cm/foldingCommands.js

@@ -1,23 +1,67 @@
 import {
 	codeFolding,
 	ensureSyntaxTree,
-	foldable,
 	foldEffect,
 	foldedRanges,
+	foldNodeProp,
+	foldService,
 	foldState,
+	syntaxTree,
 	unfoldEffect,
 } from "@codemirror/language";
 import { StateEffect } from "@codemirror/state";
 
 const FULL_PARSE_BUDGET_MS = 100;
+const INCOMPLETE_PARSE_MARGIN = 50;
+
+function findServiceFold(state, line) {
+	for (const service of state.facet(foldService)) {
+		const range = service(state, line.from, line.to);
+		if (range) return range;
+	}
+	return null;
+}
+
+function addSyntaxFolds(state, tree, line, addRange) {
+	if (!tree || tree.length < line.to) return;
+
+	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;
+		}
+
+		const fold = node.type.prop(foldNodeProp);
+		if (!fold) continue;
+
+		const range = fold(node, state);
+		if (
+			range &&
+			range.from >= line.from &&
+			range.from <= line.to &&
+			range.to > line.to
+		) {
+			addRange(range);
+		}
+	}
+}
 
 /**
  * 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;
-	ensureSyntaxTree(state, state.doc.length, FULL_PARSE_BUDGET_MS);
+	const tree =
+		ensureSyntaxTree(state, state.doc.length, FULL_PARSE_BUDGET_MS) ??
+		syntaxTree(state);
 
 	const existing = new Set();
 	foldedRanges(state).between(0, state.doc.length, (from, to) => {
@@ -26,15 +70,18 @@ export function foldAllCodeBlocks(view) {
 
 	const effects = [];
 	const discovered = new Set();
-	for (let lineNumber = 1; lineNumber <= state.doc.lines; lineNumber += 1) {
-		const line = state.doc.line(lineNumber);
-		const range = foldable(state, line.from, line.to);
-		if (!range) continue;
-
+	const addRange = (range) => {
 		const id = `${range.from}:${range.to}`;
-		if (existing.has(id) || discovered.has(id)) continue;
+		if (existing.has(id) || discovered.has(id)) return;
 		discovered.add(id);
 		effects.push(foldEffect.of(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);
 	}
 
 	if (!effects.length) return false;

src/test/editor.tests.js

@@ -760,20 +760,17 @@ export async function runCodeMirrorTests(writeOutput) {
 	);
 
 	runner.test(
-		"Fold all includes nested blocks and unfold all clears them",
+		"Fold all includes same-line nested blocks and unfold all clears them",
 		async (test) => {
 			await withEditor(
 				test,
 				async (view) => {
 					test.assert(foldAllCodeBlocks(view), "Nested blocks should fold");
-					test.assert(
-						countFolds(view) >= 2,
-						"Fold all should retain both outer and nested folds",
-					);
+					test.assertEqual(countFolds(view), 2);
 					test.assert(unfoldAllCodeBlocks(view), "All folds should unfold");
 					test.assertEqual(countFolds(view), 0);
 				},
-				"function outer() {\n  if (true) {\n    console.log('nested');\n  }\n}",
+				"function outer() { if (true) {\n  console.log('nested');\n} }",
 				[javascript()],
 			);
 		},