Simplify folding to not be extension dependent

Fixes #125

Author: celonymire

CHANGELOG.md

@@ -3,6 +3,7 @@
 ### Changed
 
 - Let the stress tester run indefinitely by default
+- Folding now doesn't depend on extensions and correctly folds the entire inserted content
 
 ### Fixed
 

README.md

@@ -370,8 +370,6 @@ Interactive testcases have a special badge to make them distinguishable. If ther
 ### 👜 Inserting Prewritten Code
 
 - Add the root directory of the templates to the settings
-- **NOTE**: Remove trailing newlines for fold to work (folding is optional via settings)
-  - Folding depends on VSCode support, which may require other extensions depending on the language.
 
 | ![Insert File Template Gif](media/insert_file_template.gif) |
 | :---------------------------------------------------------: |

src/extension/index.ts

@@ -11,6 +11,7 @@ import {
   resolveVariables,
 } from "./utils/vscode";
 import { initLogging } from "./utils/logging";
+import { TemplateFoldingProvider, type TemplateRange } from "./utils/folding";
 import JudgeViewProvider from "./providers/JudgeViewProvider";
 import StressViewProvider from "./providers/StressViewProvider";
 import PanelViewProvider from "./providers/PanelViewProvider";
@@ -20,9 +21,13 @@ import { createStatusBarItem } from "./statusBar";
 let judgeViewProvider: JudgeViewProvider;
 let stressViewProvider: StressViewProvider;
 let panelViewProvider: PanelViewProvider;
+let templateFoldingProvider: TemplateFoldingProvider;
 
 type Dependencies = Record<string, string[]>;
 
+// Track inserted template ranges by document URI for folding
+const templateRangesByUri = new Map<string, TemplateRange>();
+
 async function getTemplateContent(
   relativeFile: string,
   baseDirectory: string,
@@ -123,10 +128,27 @@ function registerDocumentContentProviders(context: vscode.ExtensionContext): voi
   context.subscriptions.push(
     vscode.workspace.onDidCloseTextDocument((document) => {
       ReadonlyStringProvider.cleanup(document.uri);
+      // Also clean up template ranges for closed documents
+      templateRangesByUri.delete(document.uri.toString());
     })
   );
 }
 
+function registerFoldingProvider(context: vscode.ExtensionContext): void {
+  // Create the folding provider with a callback to get template ranges
+  templateFoldingProvider = new TemplateFoldingProvider((documentUri) => {
+    return templateRangesByUri.get(documentUri);
+  });
+
+  // Register for all file scheme documents
+  context.subscriptions.push(
+    vscode.languages.registerFoldingRangeProvider({ scheme: "file" }, templateFoldingProvider)
+  );
+
+  // Dispose the provider when extension deactivates
+  context.subscriptions.push(templateFoldingProvider);
+}
+
 function registerCommands(context: vscode.ExtensionContext): void {
   registerRunSettingsCommands(context);
 
@@ -254,14 +276,37 @@ function registerCommands(context: vscode.ExtensionContext): void {
           return;
         }
 
-        const inserted = vscode.window.activeTextEditor?.edit((edit: vscode.TextEditorEdit) => {
-          if (vscode.window.activeTextEditor) {
-            edit.insert(vscode.window.activeTextEditor.selection.active, content);
-          }
+        const editor = vscode.window.activeTextEditor;
+        if (!editor) {
+          return;
+        }
+
+        const startLine = editor.selection.active.line;
+        const templateLineCount = content.split("\n").length - 1;
+
+        const inserted = await editor.edit((edit: vscode.TextEditorEdit) => {
+          edit.insert(editor.selection.active, content);
         });
+
+        if (!inserted) {
+          return;
+        }
+
         const foldTemplate = config.get<boolean>("foldFileTemplate")!;
-        if (inserted && foldTemplate) {
-          vscode.commands.executeCommand("editor.fold");
+        if (foldTemplate) {
+          // Track the template range for folding
+          const endLine = startLine + templateLineCount;
+          const documentUri = editor.document.uri.toString();
+          templateRangesByUri.set(documentUri, { startLine, endLine });
+
+          // Notify the folding provider that ranges have changed
+          templateFoldingProvider.notifyFoldingRangesChanged();
+
+          // Wait for VS Code's internal folding debounce (~200ms), then fold
+          setTimeout(async () => {
+            await vscode.commands.executeCommand("editor.fold", { levels: 1, direction: "down" });
+            templateRangesByUri.delete(documentUri);
+          }, 200);
         }
       })();
     })
@@ -361,6 +406,7 @@ export function activate(context: vscode.ExtensionContext): void {
   registerViewProviders(context);
   registerCommands(context);
   registerDocumentContentProviders(context);
+  registerFoldingProvider(context);
 
   createStatusBarItem(context);
 

src/extension/utils/folding.ts

@@ -0,0 +1,52 @@
+import * as vscode from "vscode";
+import { getLogger } from "./logging";
+
+export interface TemplateRange {
+  startLine: number;
+  endLine: number;
+}
+
+type GetTemplateRangesCallback = (documentUri: string) => TemplateRange | undefined;
+
+/**
+ * Custom folding provider that folds inserted template regions.
+ * Returns a single folding range for the entire template content.
+ */
+export class TemplateFoldingProvider implements vscode.FoldingRangeProvider {
+  private readonly _onDidChangeFoldingRangesEmitter = new vscode.EventEmitter<void>();
+  public readonly onDidChangeFoldingRanges = this._onDidChangeFoldingRangesEmitter.event;
+
+  constructor(private readonly getTemplateRanges: GetTemplateRangesCallback) {}
+
+  provideFoldingRanges(
+    document: vscode.TextDocument
+  ): vscode.ProviderResult<vscode.FoldingRange[]> {
+    const templateRange = this.getTemplateRanges(document.uri.toString());
+    if (!templateRange) {
+      return [];
+    }
+
+    const logger = getLogger("folding");
+    logger.debug(
+      `Providing folding range for ${document.uri.fsPath}: ${templateRange.startLine}-${templateRange.endLine}`
+    );
+
+    // Simply fold the entire inserted template region
+    return [new vscode.FoldingRange(templateRange.startLine, templateRange.endLine)];
+  }
+
+  /**
+   * Notify VS Code that folding ranges have changed.
+   * Call this after inserting a template to trigger re-evaluation.
+   */
+  notifyFoldingRangesChanged(): void {
+    this._onDidChangeFoldingRangesEmitter.fire();
+  }
+
+  /**
+   * Dispose of resources.
+   */
+  dispose(): void {
+    this._onDidChangeFoldingRangesEmitter.dispose();
+  }
+}