improve formatting integration

package.json

@@ -116,7 +116,7 @@
         "zig.formattingProvider": {
           "scope": "resource",
           "type": "string",
-          "description": "Whether to enable formatting (requires restarting editor)",
+          "description": "Whether to enable formatting",
           "enum": [
             "off",
             "extension",
@@ -129,10 +129,10 @@
           ],
           "enumDescriptions": [
             "Disable formatting",
-            "Use extension formatting",
-            "Use ZLS formatting (not recommended as zls's formatting is slower)"
+            "Provide formatting by directly invoking `zig fmt`",
+            "Provide formatting by using ZLS (which matches `zig fmt`)"
           ],
-          "default": "extension"
+          "default": "zls"
         },
         "zig.zls.debugLog": {
           "scope": "resource",

src/extension.ts

@@ -1,25 +1,16 @@
 import vscode from "vscode";
 
-import { ZigFormatProvider, ZigRangeFormatProvider } from "./zigFormat";
 import { activate as activateZls, deactivate as deactivateZls } from "./zls";
 import ZigCompilerProvider from "./zigCompilerProvider";
+import { registerDocumentFormatting } from "./zigFormat";
 import { setupZig } from "./zigSetup";
 
-const ZIG_MODE: vscode.DocumentFilter = { language: "zig", scheme: "file" };
-
 export async function activate(context: vscode.ExtensionContext) {
     await setupZig(context).finally(() => {
         const compiler = new ZigCompilerProvider();
         compiler.activate(context.subscriptions);
 
-        if (vscode.workspace.getConfiguration("zig").get<string>("formattingProvider") === "extension") {
-            context.subscriptions.push(
-                vscode.languages.registerDocumentFormattingEditProvider(ZIG_MODE, new ZigFormatProvider()),
-            );
-            context.subscriptions.push(
-                vscode.languages.registerDocumentRangeFormattingEditProvider(ZIG_MODE, new ZigRangeFormatProvider()),
-            );
-        }
+        context.subscriptions.push(registerDocumentFormatting());
 
         void activateZls(context);
     });

src/zigFormat.ts

@@ -1,31 +1,81 @@
 import vscode from "vscode";
 
 import childProcess from "child_process";
+import util from "util";
 
 import { getZigPath } from "./zigUtil";
 
-export class ZigFormatProvider implements vscode.DocumentFormattingEditProvider {
-    provideDocumentFormattingEdits(document: vscode.TextDocument): Promise<vscode.TextEdit[] | null> {
-        return Promise.resolve(zigFormat(document));
-    }
+const execFile = util.promisify(childProcess.execFile);
+const ZIG_MODE: vscode.DocumentSelector = { language: "zig" };
+
+export function registerDocumentFormatting(): vscode.Disposable {
+    let registeredFormatter: vscode.Disposable | null = null;
+
+    preCompileZigFmt();
+    vscode.workspace.onDidChangeConfiguration((change: vscode.ConfigurationChangeEvent) => {
+        if (
+            change.affectsConfiguration("zig.path", undefined) ||
+            change.affectsConfiguration("zig.formattingProvider", undefined)
+        ) {
+            preCompileZigFmt();
+        }
+    });
+
+    const onformattingProviderChange = () => {
+        if (vscode.workspace.getConfiguration("zig").get<string>("formattingProvider") === "off") {
+            // Unregister the formatting provider
+            if (registeredFormatter !== null) registeredFormatter.dispose();
+            registeredFormatter = null;
+        } else {
+            // register the formatting provider
+            registeredFormatter ??= vscode.languages.registerDocumentRangeFormattingEditProvider(ZIG_MODE, {
+                provideDocumentRangeFormattingEdits,
+            });
+        }
+    };
+
+    onformattingProviderChange();
+    const registeredDidChangeEvent = vscode.workspace.onDidChangeConfiguration(onformattingProviderChange);
+
+    return {
+        dispose: () => {
+            registeredDidChangeEvent.dispose();
+            if (registeredFormatter !== null) registeredFormatter.dispose();
+        },
+    };
 }
 
-// Same as full document formatter for now
-export class ZigRangeFormatProvider implements vscode.DocumentRangeFormattingEditProvider {
-    provideDocumentRangeFormattingEdits(document: vscode.TextDocument): Promise<vscode.TextEdit[] | null> {
-        return Promise.resolve(zigFormat(document));
-    }
+/** Ensures that `zig fmt` has been JIT compiled. */
+function preCompileZigFmt() {
+    // This pre-compiles even if "zig.formattingProvider" is "zls".
+    if (vscode.workspace.getConfiguration("zig").get<string>("formattingProvider") === "off") return;
+
+    childProcess.execFile(getZigPath(), ["fmt", "--help"], {
+        timeout: 60000, // 60 seconds (this is a very high value because 'zig fmt' is just in time compiled)
+    });
 }
 
-function zigFormat(document: vscode.TextDocument): vscode.TextEdit[] | null {
+async function provideDocumentRangeFormattingEdits(
+    document: vscode.TextDocument,
+    range: vscode.Range,
+    options: vscode.FormattingOptions,
+    token: vscode.CancellationToken,
+): Promise<vscode.TextEdit[] | null> {
     const zigPath = getZigPath();
 
-    const stdout = childProcess.execFileSync(zigPath, ["fmt", "--stdin"], {
-        input: document.getText(),
+    const abortController = new AbortController();
+    token.onCancellationRequested(() => {
+        abortController.abort();
+    });
+
+    const promise = execFile(zigPath, ["fmt", "--stdin"], {
         maxBuffer: 10 * 1024 * 1024, // 10MB
-        encoding: "utf8",
+        signal: abortController.signal,
         timeout: 60000, // 60 seconds (this is a very high value because 'zig fmt' is just in time compiled)
     });
+    promise.child.stdin?.end(document.getText());
+
+    const { stdout } = await promise;
 
     if (stdout.length === 0) return null;
     const lastLineId = document.lineCount - 1;

src/zls.ts

@@ -5,6 +5,7 @@ import fs from "fs";
 import {
     CancellationToken,
     ConfigurationParams,
+    DocumentSelector,
     LSPAny,
     LanguageClient,
     LanguageClientOptions,
@@ -30,6 +31,11 @@ import {
 let outputChannel: vscode.OutputChannel;
 export let client: LanguageClient | null = null;
 
+const ZIG_MODE: DocumentSelector = [
+    { language: "zig", scheme: "file" },
+    { language: "zig", scheme: "untitled" },
+];
+
 async function startClient() {
     const configuration = vscode.workspace.getConfiguration("zig.zls");
     const debugLog = configuration.get<boolean>("debugLog", false);
@@ -43,7 +49,7 @@ async function startClient() {
 
     // Options to control the language client
     const clientOptions: LanguageClientOptions = {
-        documentSelector: [{ scheme: "file", language: "zig" }],
+        documentSelector: ZIG_MODE,
         outputChannel,
         middleware: {
             workspace: {
@@ -329,7 +335,7 @@ async function installVersion(context: vscode.ExtensionContext, version: semver.
                 }
                 throw err;
             }
-            
+
             await stopClient();
 
             const installDir = vscode.Uri.joinPath(context.globalStorageUri, "zls_install");
@@ -407,6 +413,14 @@ export async function activate(context: vscode.ExtensionContext) {
                     await startClient();
                 }
             }
+            if (client && change.affectsConfiguration("zig.formattingProvider", undefined)) {
+                client.getFeature("textDocument/formatting").dispose();
+                if (vscode.workspace.getConfiguration("zig").get<string>("formattingProvider") === "zls") {
+                    client
+                        .getFeature("textDocument/formatting")
+                        .initialize(client.initializeResult?.capabilities ?? {}, ZIG_MODE);
+                }
+            }
         }),
     );