compiler: JIT zig fmt

See #19063

4 files changed, 573 insertions, 6 deletions

diff --git a/lib/std/zig.zig b/lib/std/zig.zig
index 9085b23de1..a195889d77 100644
--- a/lib/std/zig.zig
+++ b/lib/std/zig.zig
@@ -1,6 +1,3 @@
-/// Implementation of `zig fmt`.
-pub const fmt = @import("zig/fmt.zig");
-
 pub const ErrorBundle = @import("zig/ErrorBundle.zig");
 pub const Server = @import("zig/Server.zig");
 pub const Client = @import("zig/Client.zig");
@@ -30,6 +27,36 @@ pub const c_translation = @import("zig/c_translation.zig");
 pub const SrcHasher = std.crypto.hash.Blake3;
 pub const SrcHash = [16]u8;
 
+pub const Color = enum {
+    /// Determine whether stderr is a terminal or not automatically.
+    auto,
+    /// Assume stderr is not a terminal.
+    off,
+    /// Assume stderr is a terminal.
+    on,
+
+    pub fn get_tty_conf(color: Color) std.io.tty.Config {
+        return switch (color) {
+            .auto => std.io.tty.detectConfig(std.io.getStdErr()),
+            .on => .escape_codes,
+            .off => .no_color,
+        };
+    }
+
+    pub fn renderOptions(color: Color) std.zig.ErrorBundle.RenderOptions {
+        const ttyconf = get_tty_conf(color);
+        return .{
+            .ttyconf = ttyconf,
+            .include_source_line = ttyconf != .no_color,
+            .include_reference_trace = ttyconf != .no_color,
+        };
+    }
+};
+
+/// There are many assumptions in the entire codebase that Zig source files can
+/// be byte-indexed with a u32 integer.
+pub const max_src_size = std.math.maxInt(u32);
+
 pub fn hashSrc(src: []const u8) SrcHash {
     var out: SrcHash = undefined;
     SrcHasher.hash(src, &out, .{});
@@ -801,6 +828,78 @@ test isValidId {
     try std.testing.expect(isValidId("i386"));
 }
 
+pub fn readSourceFileToEndAlloc(
+    allocator: Allocator,
+    input: std.fs.File,
+    size_hint: ?usize,
+) ![:0]u8 {
+    const source_code = input.readToEndAllocOptions(
+        allocator,
+        max_src_size,
+        size_hint,
+        @alignOf(u16),
+        0,
+    ) catch |err| switch (err) {
+        error.ConnectionResetByPeer => unreachable,
+        error.ConnectionTimedOut => unreachable,
+        error.NotOpenForReading => unreachable,
+        else => |e| return e,
+    };
+    errdefer allocator.free(source_code);
+
+    // Detect unsupported file types with their Byte Order Mark
+    const unsupported_boms = [_][]const u8{
+        "\xff\xfe\x00\x00", // UTF-32 little endian
+        "\xfe\xff\x00\x00", // UTF-32 big endian
+        "\xfe\xff", // UTF-16 big endian
+    };
+    for (unsupported_boms) |bom| {
+        if (std.mem.startsWith(u8, source_code, bom)) {
+            return error.UnsupportedEncoding;
+        }
+    }
+
+    // If the file starts with a UTF-16 little endian BOM, translate it to UTF-8
+    if (std.mem.startsWith(u8, source_code, "\xff\xfe")) {
+        const source_code_utf16_le = std.mem.bytesAsSlice(u16, source_code);
+        const source_code_utf8 = std.unicode.utf16LeToUtf8AllocZ(allocator, source_code_utf16_le) catch |err| switch (err) {
+            error.DanglingSurrogateHalf => error.UnsupportedEncoding,
+            error.ExpectedSecondSurrogateHalf => error.UnsupportedEncoding,
+            error.UnexpectedSecondSurrogateHalf => error.UnsupportedEncoding,
+            else => |e| return e,
+        };
+
+        allocator.free(source_code);
+        return source_code_utf8;
+    }
+
+    return source_code;
+}
+
+pub fn printAstErrorsToStderr(gpa: Allocator, tree: Ast, path: []const u8, color: Color) !void {
+    var wip_errors: std.zig.ErrorBundle.Wip = undefined;
+    try wip_errors.init(gpa);
+    defer wip_errors.deinit();
+
+    try putAstErrorsIntoBundle(gpa, tree, path, &wip_errors);
+
+    var error_bundle = try wip_errors.toOwnedBundle("");
+    defer error_bundle.deinit(gpa);
+    error_bundle.renderToStdErr(color.renderOptions());
+}
+
+pub fn putAstErrorsIntoBundle(
+    gpa: Allocator,
+    tree: Ast,
+    path: []const u8,
+    wip_errors: *std.zig.ErrorBundle.Wip,
+) Allocator.Error!void {
+    var zir = try AstGen.generate(gpa, tree);
+    defer zir.deinit(gpa);
+
+    try wip_errors.addZirErrorMessages(zir, tree, tree.source, path);
+}
+
 test {
     _ = Ast;
     _ = AstRlAnnotate;
@@ -808,9 +907,12 @@ test {
     _ = Client;
     _ = ErrorBundle;
     _ = Server;
-    _ = fmt;
     _ = number_literal;
     _ = primitives;
     _ = string_literal;
     _ = system;
+
+    // This is not standard library API; it is the standalone executable
+    // implementation of `zig fmt`.
+    _ = @import("zig/fmt.zig");
 }
diff --git a/lib/std/zig/Ast.zig b/lib/std/zig/Ast.zig
index d4e393bdf0..6f4afbbe4f 100644
--- a/lib/std/zig/Ast.zig
+++ b/lib/std/zig/Ast.zig
@@ -32,6 +32,12 @@ pub const Location = struct {
     line_end: usize,
 };
 
+pub const Span = struct {
+    start: u32,
+    end: u32,
+    main: u32,
+};
+
 pub fn deinit(tree: *Ast, gpa: Allocator) void {
     tree.tokens.deinit(gpa);
     tree.nodes.deinit(gpa);
@@ -3533,6 +3539,39 @@ pub const Node = struct {
     };
 };
 
+pub fn nodeToSpan(tree: *const Ast, node: u32) Span {
+    return tokensToSpan(
+        tree,
+        tree.firstToken(node),
+        tree.lastToken(node),
+        tree.nodes.items(.main_token)[node],
+    );
+}
+
+pub fn tokenToSpan(tree: *const Ast, token: Ast.TokenIndex) Span {
+    return tokensToSpan(tree, token, token, token);
+}
+
+pub fn tokensToSpan(tree: *const Ast, start: Ast.TokenIndex, end: Ast.TokenIndex, main: Ast.TokenIndex) Span {
+    const token_starts = tree.tokens.items(.start);
+    var start_tok = start;
+    var end_tok = end;
+
+    if (tree.tokensOnSameLine(start, end)) {
+        // do nothing
+    } else if (tree.tokensOnSameLine(start, main)) {
+        end_tok = main;
+    } else if (tree.tokensOnSameLine(main, end)) {
+        start_tok = main;
+    } else {
+        start_tok = main;
+        end_tok = main;
+    }
+    const start_off = token_starts[start_tok];
+    const end_off = token_starts[end_tok] + @as(u32, @intCast(tree.tokenSlice(end_tok).len));
+    return Span{ .start = start_off, .end = end_off, .main = token_starts[main] };
+}
+
 const std = @import("../std.zig");
 const assert = std.debug.assert;
 const testing = std.testing;
@@ -3544,5 +3583,6 @@ const Parse = @import("Parse.zig");
 const private_render = @import("./render.zig");
 
 test {
-    testing.refAllDecls(@This());
+    _ = Parse;
+    _ = private_render;
 }
diff --git a/lib/std/zig/ErrorBundle.zig b/lib/std/zig/ErrorBundle.zig
index ff47e3794b..013d447ab1 100644
--- a/lib/std/zig/ErrorBundle.zig
+++ b/lib/std/zig/ErrorBundle.zig
@@ -459,6 +459,90 @@ pub const Wip = struct {
         return @intCast(wip.extra.items.len - notes_len);
     }
 
+    pub fn addZirErrorMessages(
+        eb: *ErrorBundle.Wip,
+        zir: std.zig.Zir,
+        tree: std.zig.Ast,
+        source: [:0]const u8,
+        src_path: []const u8,
+    ) !void {
+        const Zir = std.zig.Zir;
+        const payload_index = zir.extra[@intFromEnum(Zir.ExtraIndex.compile_errors)];
+        assert(payload_index != 0);
+
+        const header = zir.extraData(Zir.Inst.CompileErrors, payload_index);
+        const items_len = header.data.items_len;
+        var extra_index = header.end;
+        for (0..items_len) |_| {
+            const item = zir.extraData(Zir.Inst.CompileErrors.Item, extra_index);
+            extra_index = item.end;
+            const err_span = blk: {
+                if (item.data.node != 0) {
+                    break :blk tree.nodeToSpan(item.data.node);
+                }
+                const token_starts = tree.tokens.items(.start);
+                const start = token_starts[item.data.token] + item.data.byte_offset;
+                const end = start + @as(u32, @intCast(tree.tokenSlice(item.data.token).len)) - item.data.byte_offset;
+                break :blk std.zig.Ast.Span{ .start = start, .end = end, .main = start };
+            };
+            const err_loc = std.zig.findLineColumn(source, err_span.main);
+
+            {
+                const msg = zir.nullTerminatedString(item.data.msg);
+                try eb.addRootErrorMessage(.{
+                    .msg = try eb.addString(msg),
+                    .src_loc = try eb.addSourceLocation(.{
+                        .src_path = try eb.addString(src_path),
+                        .span_start = err_span.start,
+                        .span_main = err_span.main,
+                        .span_end = err_span.end,
+                        .line = @intCast(err_loc.line),
+                        .column = @intCast(err_loc.column),
+                        .source_line = try eb.addString(err_loc.source_line),
+                    }),
+                    .notes_len = item.data.notesLen(zir),
+                });
+            }
+
+            if (item.data.notes != 0) {
+                const notes_start = try eb.reserveNotes(item.data.notes);
+                const block = zir.extraData(Zir.Inst.Block, item.data.notes);
+                const body = zir.extra[block.end..][0..block.data.body_len];
+                for (notes_start.., body) |note_i, body_elem| {
+                    const note_item = zir.extraData(Zir.Inst.CompileErrors.Item, body_elem);
+                    const msg = zir.nullTerminatedString(note_item.data.msg);
+                    const span = blk: {
+                        if (note_item.data.node != 0) {
+                            break :blk tree.nodeToSpan(note_item.data.node);
+                        }
+                        const token_starts = tree.tokens.items(.start);
+                        const start = token_starts[note_item.data.token] + note_item.data.byte_offset;
+                        const end = start + @as(u32, @intCast(tree.tokenSlice(note_item.data.token).len)) - item.data.byte_offset;
+                        break :blk std.zig.Ast.Span{ .start = start, .end = end, .main = start };
+                    };
+                    const loc = std.zig.findLineColumn(source, span.main);
+
+                    eb.extra.items[note_i] = @intFromEnum(try eb.addErrorMessage(.{
+                        .msg = try eb.addString(msg),
+                        .src_loc = try eb.addSourceLocation(.{
+                            .src_path = try eb.addString(src_path),
+                            .span_start = span.start,
+                            .span_main = span.main,
+                            .span_end = span.end,
+                            .line = @intCast(loc.line),
+                            .column = @intCast(loc.column),
+                            .source_line = if (loc.eql(err_loc))
+                                0
+                            else
+                                try eb.addString(loc.source_line),
+                        }),
+                        .notes_len = 0, // TODO rework this function to be recursive
+                    }));
+                }
+            }
+        }
+    }
+
     fn addOtherMessage(wip: *Wip, other: ErrorBundle, msg_index: MessageIndex) !MessageIndex {
         const other_msg = other.getErrorMessage(msg_index);
         const src_loc = try wip.addOtherSourceLocation(other, other_msg.src_loc);
diff --git a/lib/std/zig/fmt.zig b/lib/std/zig/fmt.zig
index f8841bfb5b..2fc04b7935 100644
--- a/lib/std/zig/fmt.zig
+++ b/lib/std/zig/fmt.zig
@@ -1 +1,342 @@
-const std = @import("../std.zig");
+const std = @import("std");
+const mem = std.mem;
+const fs = std.fs;
+const process = std.process;
+const Allocator = std.mem.Allocator;
+const warn = std.log.warn;
+const Color = std.zig.Color;
+
+const usage_fmt =
+    \\Usage: zig fmt [file]...
+    \\
+    \\   Formats the input files and modifies them in-place.
+    \\   Arguments can be files or directories, which are searched
+    \\   recursively.
+    \\
+    \\Options:
+    \\  -h, --help             Print this help and exit
+    \\  --color [auto|off|on]  Enable or disable colored error messages
+    \\  --stdin                Format code from stdin; output to stdout
+    \\  --check                List non-conforming files and exit with an error
+    \\                         if the list is non-empty
+    \\  --ast-check            Run zig ast-check on every file
+    \\  --exclude [file]       Exclude file or directory from formatting
+    \\
+    \\
+;
+
+const Fmt = struct {
+    seen: SeenMap,
+    any_error: bool,
+    check_ast: bool,
+    color: Color,
+    gpa: Allocator,
+    arena: Allocator,
+    out_buffer: std.ArrayList(u8),
+
+    const SeenMap = std.AutoHashMap(fs.File.INode, void);
+};
+
+pub fn main() !void {
+    var arena_instance = std.heap.ArenaAllocator.init(std.heap.page_allocator);
+    defer arena_instance.deinit();
+    const arena = arena_instance.allocator();
+    const gpa = arena;
+
+    const args = try process.argsAlloc(arena);
+
+    var color: Color = .auto;
+    var stdin_flag: bool = false;
+    var check_flag: bool = false;
+    var check_ast_flag: bool = false;
+    var input_files = std.ArrayList([]const u8).init(gpa);
+    defer input_files.deinit();
+    var excluded_files = std.ArrayList([]const u8).init(gpa);
+    defer excluded_files.deinit();
+
+    {
+        var i: usize = 1;
+        while (i < args.len) : (i += 1) {
+            const arg = args[i];
+            if (mem.startsWith(u8, arg, "-")) {
+                if (mem.eql(u8, arg, "-h") or mem.eql(u8, arg, "--help")) {
+                    const stdout = std.io.getStdOut().writer();
+                    try stdout.writeAll(usage_fmt);
+                    return process.cleanExit();
+                } else if (mem.eql(u8, arg, "--color")) {
+                    if (i + 1 >= args.len) {
+                        fatal("expected [auto|on|off] after --color", .{});
+                    }
+                    i += 1;
+                    const next_arg = args[i];
+                    color = std.meta.stringToEnum(Color, next_arg) orelse {
+                        fatal("expected [auto|on|off] after --color, found '{s}'", .{next_arg});
+                    };
+                } else if (mem.eql(u8, arg, "--stdin")) {
+                    stdin_flag = true;
+                } else if (mem.eql(u8, arg, "--check")) {
+                    check_flag = true;
+                } else if (mem.eql(u8, arg, "--ast-check")) {
+                    check_ast_flag = true;
+                } else if (mem.eql(u8, arg, "--exclude")) {
+                    if (i + 1 >= args.len) {
+                        fatal("expected parameter after --exclude", .{});
+                    }
+                    i += 1;
+                    const next_arg = args[i];
+                    try excluded_files.append(next_arg);
+                } else {
+                    fatal("unrecognized parameter: '{s}'", .{arg});
+                }
+            } else {
+                try input_files.append(arg);
+            }
+        }
+    }
+
+    if (stdin_flag) {
+        if (input_files.items.len != 0) {
+            fatal("cannot use --stdin with positional arguments", .{});
+        }
+
+        const stdin = std.io.getStdIn();
+        const source_code = std.zig.readSourceFileToEndAlloc(gpa, stdin, null) catch |err| {
+            fatal("unable to read stdin: {}", .{err});
+        };
+        defer gpa.free(source_code);
+
+        var tree = std.zig.Ast.parse(gpa, source_code, .zig) catch |err| {
+            fatal("error parsing stdin: {}", .{err});
+        };
+        defer tree.deinit(gpa);
+
+        if (check_ast_flag) {
+            var zir = try std.zig.AstGen.generate(gpa, tree);
+
+            if (zir.hasCompileErrors()) {
+                var wip_errors: std.zig.ErrorBundle.Wip = undefined;
+                try wip_errors.init(gpa);
+                defer wip_errors.deinit();
+                try wip_errors.addZirErrorMessages(zir, tree, source_code, "<stdin>");
+                var error_bundle = try wip_errors.toOwnedBundle("");
+                defer error_bundle.deinit(gpa);
+                error_bundle.renderToStdErr(color.renderOptions());
+                process.exit(2);
+            }
+        } else if (tree.errors.len != 0) {
+            try std.zig.printAstErrorsToStderr(gpa, tree, "<stdin>", color);
+            process.exit(2);
+        }
+        const formatted = try tree.render(gpa);
+        defer gpa.free(formatted);
+
+        if (check_flag) {
+            const code: u8 = @intFromBool(mem.eql(u8, formatted, source_code));
+            process.exit(code);
+        }
+
+        return std.io.getStdOut().writeAll(formatted);
+    }
+
+    if (input_files.items.len == 0) {
+        fatal("expected at least one source file argument", .{});
+    }
+
+    var fmt = Fmt{
+        .gpa = gpa,
+        .arena = arena,
+        .seen = Fmt.SeenMap.init(gpa),
+        .any_error = false,
+        .check_ast = check_ast_flag,
+        .color = color,
+        .out_buffer = std.ArrayList(u8).init(gpa),
+    };
+    defer fmt.seen.deinit();
+    defer fmt.out_buffer.deinit();
+
+    // Mark any excluded files/directories as already seen,
+    // so that they are skipped later during actual processing
+    for (excluded_files.items) |file_path| {
+        const stat = fs.cwd().statFile(file_path) catch |err| switch (err) {
+            error.FileNotFound => continue,
+            // On Windows, statFile does not work for directories
+            error.IsDir => dir: {
+                var dir = try fs.cwd().openDir(file_path, .{});
+                defer dir.close();
+                break :dir try dir.stat();
+            },
+            else => |e| return e,
+        };
+        try fmt.seen.put(stat.inode, {});
+    }
+
+    for (input_files.items) |file_path| {
+        try fmtPath(&fmt, file_path, check_flag, fs.cwd(), file_path);
+    }
+    if (fmt.any_error) {
+        process.exit(1);
+    }
+}
+
+const FmtError = error{
+    SystemResources,
+    OperationAborted,
+    IoPending,
+    BrokenPipe,
+    Unexpected,
+    WouldBlock,
+    FileClosed,
+    DestinationAddressRequired,
+    DiskQuota,
+    FileTooBig,
+    InputOutput,
+    NoSpaceLeft,
+    AccessDenied,
+    OutOfMemory,
+    RenameAcrossMountPoints,
+    ReadOnlyFileSystem,
+    LinkQuotaExceeded,
+    FileBusy,
+    EndOfStream,
+    Unseekable,
+    NotOpenForWriting,
+    UnsupportedEncoding,
+    ConnectionResetByPeer,
+    SocketNotConnected,
+    LockViolation,
+    NetNameDeleted,
+    InvalidArgument,
+} || fs.File.OpenError;
+
+fn fmtPath(fmt: *Fmt, file_path: []const u8, check_mode: bool, dir: fs.Dir, sub_path: []const u8) FmtError!void {
+    fmtPathFile(fmt, file_path, check_mode, dir, sub_path) catch |err| switch (err) {
+        error.IsDir, error.AccessDenied => return fmtPathDir(fmt, file_path, check_mode, dir, sub_path),
+        else => {
+            warn("unable to format '{s}': {s}", .{ file_path, @errorName(err) });
+            fmt.any_error = true;
+            return;
+        },
+    };
+}
+
+fn fmtPathDir(
+    fmt: *Fmt,
+    file_path: []const u8,
+    check_mode: bool,
+    parent_dir: fs.Dir,
+    parent_sub_path: []const u8,
+) FmtError!void {
+    var dir = try parent_dir.openDir(parent_sub_path, .{ .iterate = true });
+    defer dir.close();
+
+    const stat = try dir.stat();
+    if (try fmt.seen.fetchPut(stat.inode, {})) |_| return;
+
+    var dir_it = dir.iterate();
+    while (try dir_it.next()) |entry| {
+        const is_dir = entry.kind == .directory;
+
+        if (is_dir and (mem.eql(u8, entry.name, "zig-cache") or mem.eql(u8, entry.name, "zig-out"))) continue;
+
+        if (is_dir or entry.kind == .file and (mem.endsWith(u8, entry.name, ".zig") or mem.endsWith(u8, entry.name, ".zon"))) {
+            const full_path = try fs.path.join(fmt.gpa, &[_][]const u8{ file_path, entry.name });
+            defer fmt.gpa.free(full_path);
+
+            if (is_dir) {
+                try fmtPathDir(fmt, full_path, check_mode, dir, entry.name);
+            } else {
+                fmtPathFile(fmt, full_path, check_mode, dir, entry.name) catch |err| {
+                    warn("unable to format '{s}': {s}", .{ full_path, @errorName(err) });
+                    fmt.any_error = true;
+                    return;
+                };
+            }
+        }
+    }
+}
+
+fn fmtPathFile(
+    fmt: *Fmt,
+    file_path: []const u8,
+    check_mode: bool,
+    dir: fs.Dir,
+    sub_path: []const u8,
+) FmtError!void {
+    const source_file = try dir.openFile(sub_path, .{});
+    var file_closed = false;
+    errdefer if (!file_closed) source_file.close();
+
+    const stat = try source_file.stat();
+
+    if (stat.kind == .directory)
+        return error.IsDir;
+
+    const gpa = fmt.gpa;
+    const source_code = try std.zig.readSourceFileToEndAlloc(
+        gpa,
+        source_file,
+        std.math.cast(usize, stat.size) orelse return error.FileTooBig,
+    );
+    defer gpa.free(source_code);
+
+    source_file.close();
+    file_closed = true;
+
+    // Add to set after no longer possible to get error.IsDir.
+    if (try fmt.seen.fetchPut(stat.inode, {})) |_| return;
+
+    var tree = try std.zig.Ast.parse(gpa, source_code, .zig);
+    defer tree.deinit(gpa);
+
+    if (tree.errors.len != 0) {
+        try std.zig.printAstErrorsToStderr(gpa, tree, file_path, fmt.color);
+        fmt.any_error = true;
+        return;
+    }
+
+    if (fmt.check_ast) {
+        if (stat.size > std.zig.max_src_size)
+            return error.FileTooBig;
+
+        var zir = try std.zig.AstGen.generate(gpa, tree);
+        defer zir.deinit(gpa);
+
+        if (zir.hasCompileErrors()) {
+            var wip_errors: std.zig.ErrorBundle.Wip = undefined;
+            try wip_errors.init(gpa);
+            defer wip_errors.deinit();
+            try wip_errors.addZirErrorMessages(zir, tree, source_code, file_path);
+            var error_bundle = try wip_errors.toOwnedBundle("");
+            defer error_bundle.deinit(gpa);
+            error_bundle.renderToStdErr(fmt.color.renderOptions());
+            fmt.any_error = true;
+        }
+    }
+
+    // As a heuristic, we make enough capacity for the same as the input source.
+    fmt.out_buffer.shrinkRetainingCapacity(0);
+    try fmt.out_buffer.ensureTotalCapacity(source_code.len);
+
+    try tree.renderToArrayList(&fmt.out_buffer, .{});
+    if (mem.eql(u8, fmt.out_buffer.items, source_code))
+        return;
+
+    if (check_mode) {
+        const stdout = std.io.getStdOut().writer();
+        try stdout.print("{s}\n", .{file_path});
+        fmt.any_error = true;
+    } else {
+        var af = try dir.atomicFile(sub_path, .{ .mode = stat.mode });
+        defer af.deinit();
+
+        try af.file.writeAll(fmt.out_buffer.items);
+        try af.finish();
+        const stdout = std.io.getStdOut().writer();
+        try stdout.print("{s}\n", .{file_path});
+    }
+}
+
+fn fatal(comptime format: []const u8, args: anytype) noreturn {
+    std.log.err(format, args);
+    process.exit(1);
+}