compiler: introduce ZonGen and make `ast-check` run it for ZON inputs

Currently, `zig ast-check` fails on ZON files, because it tries to
interpret the file as Zig source code. This commit introduces a new
verification pass, `std.zig.ZonGen`, which applies to an AST in ZON
mode.

Like `AstGen`, this pass also converts the AST into a more helpful
format. Rather than a sequence of instructions like `Zir`, the output
format of `ZonGen` is a new datastructure called `Zoir`. This type is
essentially a simpler form of AST, containing only the information
required for consumers of ZON. It is also far more compact than
`std.zig.Ast`, with the size generally being comparable to the size of
the well-formatted source file.

The emitted `Zoir` is currently not used aside from the `-t` option to
`ast-check` which causes it to be dumped to stdout. However, in future,
it can be used for comptime `@import` of ZON files, as well as for
simpler handling of files like `build.zig.zon`, and even by other parts
of the Zig Standard Library.

Resolves: #22078

3 files changed, 323 insertions, 101 deletions

diff --git a/src/fmt.zig b/src/fmt.zig
index b545905367..9ab581cad4 100644
--- a/src/fmt.zig
+++ b/src/fmt.zig
@@ -13,6 +13,7 @@ const usage_fmt =
     \\                         if the list is non-empty
     \\  --ast-check            Run zig ast-check on every file
     \\  --exclude [file]       Exclude file or directory from formatting
+    \\  --zon                  Treat all input files as ZON, regardless of file extension
     \\
     \\
 ;
@@ -21,6 +22,7 @@ const Fmt = struct {
     seen: SeenMap,
     any_error: bool,
     check_ast: bool,
+    force_zon: bool,
     color: Color,
     gpa: Allocator,
     arena: Allocator,
@@ -35,9 +37,10 @@ pub fn run(
     args: []const []const u8,
 ) !void {
     var color: Color = .auto;
-    var stdin_flag: bool = false;
-    var check_flag: bool = false;
-    var check_ast_flag: bool = false;
+    var stdin_flag = false;
+    var check_flag = false;
+    var check_ast_flag = false;
+    var force_zon = false;
     var input_files = std.ArrayList([]const u8).init(gpa);
     defer input_files.deinit();
     var excluded_files = std.ArrayList([]const u8).init(gpa);
@@ -74,6 +77,8 @@ pub fn run(
                     i += 1;
                     const next_arg = args[i];
                     try excluded_files.append(next_arg);
+                } else if (mem.eql(u8, arg, "--zon")) {
+                    force_zon = true;
                 } else {
                     fatal("unrecognized parameter: '{s}'", .{arg});
                 }
@@ -94,23 +99,40 @@ pub fn run(
         };
         defer gpa.free(source_code);
 
-        var tree = std.zig.Ast.parse(gpa, source_code, .zig) catch |err| {
+        var tree = std.zig.Ast.parse(gpa, source_code, if (force_zon) .zon else .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);
+            if (!force_zon) {
+                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, "<stdin>");
+                    var error_bundle = try wip_errors.toOwnedBundle("");
+                    defer error_bundle.deinit(gpa);
+                    error_bundle.renderToStdErr(color.renderOptions());
+                    process.exit(2);
+                }
+            } else {
+                const zoir = try std.zig.ZonGen.generate(gpa, tree);
+                defer zoir.deinit(gpa);
+
+                if (zoir.hasCompileErrors()) {
+                    var wip_errors: std.zig.ErrorBundle.Wip = undefined;
+                    try wip_errors.init(gpa);
+                    defer wip_errors.deinit();
+                    try wip_errors.addZoirErrorMessages(zoir, 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);
@@ -131,12 +153,13 @@ pub fn run(
         fatal("expected at least one source file argument", .{});
     }
 
-    var fmt = Fmt{
+    var fmt: Fmt = .{
         .gpa = gpa,
         .arena = arena,
-        .seen = Fmt.SeenMap.init(gpa),
+        .seen = .init(gpa),
         .any_error = false,
         .check_ast = check_ast_flag,
+        .force_zon = force_zon,
         .color = color,
         .out_buffer = std.ArrayList(u8).init(gpa),
     };
@@ -276,7 +299,13 @@ fn fmtPathFile(
     // 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);
+    const mode: std.zig.Ast.Mode = mode: {
+        if (fmt.force_zon) break :mode .zon;
+        if (mem.endsWith(u8, sub_path, ".zon")) break :mode .zon;
+        break :mode .zig;
+    };
+
+    var tree = try std.zig.Ast.parse(gpa, source_code, mode);
     defer tree.deinit(gpa);
 
     if (tree.errors.len != 0) {
@@ -289,18 +318,37 @@ fn fmtPathFile(
         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;
+        switch (mode) {
+            .zig => {
+                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;
+                }
+            },
+            .zon => {
+                var zoir = try std.zig.ZonGen.generate(gpa, tree);
+                defer zoir.deinit(gpa);
+
+                if (zoir.hasCompileErrors()) {
+                    var wip_errors: std.zig.ErrorBundle.Wip = undefined;
+                    try wip_errors.init(gpa);
+                    defer wip_errors.deinit();
+                    try wip_errors.addZoirErrorMessages(zoir, 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;
+                }
+            },
         }
     }
 
diff --git a/src/main.zig b/src/main.zig
index 39fd3e6213..1d0f15f8fd 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -19,6 +19,7 @@ const Directory = std.Build.Cache.Directory;
 const EnvVar = std.zig.EnvVar;
 const LibCInstallation = std.zig.LibCInstallation;
 const AstGen = std.zig.AstGen;
+const ZonGen = std.zig.ZonGen;
 const Server = std.zig.Server;
 
 const tracy = @import("tracy.zig");
@@ -6007,15 +6008,16 @@ fn parseCodeModel(arg: []const u8) std.builtin.CodeModel {
 const usage_ast_check =
     \\Usage: zig ast-check [file]
     \\
-    \\    Given a .zig source file, reports any compile errors that can be
-    \\    ascertained on the basis of the source code alone, without target
-    \\    information or type checking.
+    \\    Given a .zig source file or .zon file, reports any compile errors
+    \\    that can be ascertained on the basis of the source code alone,
+    \\    without target information or type checking.
     \\
     \\    If [file] is omitted, stdin is used.
     \\
     \\Options:
     \\  -h, --help            Print this help and exit
     \\  --color [auto|off|on] Enable or disable colored error messages
+    \\  --zon                 Treat the input file as ZON, regardless of file extension
     \\  -t                    (debug option) Output ZIR in text form to stdout
     \\
     \\
@@ -6032,6 +6034,7 @@ fn cmdAstCheck(
 
     var color: Color = .auto;
     var want_output_text = false;
+    var force_zon = false;
     var zig_source_file: ?[]const u8 = null;
 
     var i: usize = 0;
@@ -6043,6 +6046,8 @@ fn cmdAstCheck(
                 return cleanExit();
             } else if (mem.eql(u8, arg, "-t")) {
                 want_output_text = true;
+            } else if (mem.eql(u8, arg, "--zon")) {
+                force_zon = true;
             } else if (mem.eql(u8, arg, "--color")) {
                 if (i + 1 >= args.len) {
                     fatal("expected [auto|on|off] after --color", .{});
@@ -6110,89 +6115,136 @@ fn cmdAstCheck(
         file.stat.size = source.len;
     }
 
+    const mode: Ast.Mode = mode: {
+        if (force_zon) break :mode .zon;
+        if (zig_source_file) |name| {
+            if (mem.endsWith(u8, name, ".zon")) {
+                break :mode .zon;
+            }
+        }
+        break :mode .zig;
+    };
+
     file.mod = try Package.Module.createLimited(arena, .{
         .root = Path.cwd(),
         .root_src_path = file.sub_file_path,
         .fully_qualified_name = "root",
     });
 
-    file.tree = try Ast.parse(gpa, file.source, .zig);
+    file.tree = try Ast.parse(gpa, file.source, mode);
     file.tree_loaded = true;
     defer file.tree.deinit(gpa);
 
-    file.zir = try AstGen.generate(gpa, file.tree);
-    file.zir_loaded = true;
-    defer file.zir.deinit(gpa);
+    switch (mode) {
+        .zig => {
+            file.zir = try AstGen.generate(gpa, file.tree);
+            file.zir_loaded = true;
+            defer file.zir.deinit(gpa);
+
+            if (file.zir.hasCompileErrors()) {
+                var wip_errors: std.zig.ErrorBundle.Wip = undefined;
+                try wip_errors.init(gpa);
+                defer wip_errors.deinit();
+                try Compilation.addZirErrorMessages(&wip_errors, &file);
+                var error_bundle = try wip_errors.toOwnedBundle("");
+                defer error_bundle.deinit(gpa);
+                error_bundle.renderToStdErr(color.renderOptions());
+
+                if (file.zir.loweringFailed()) {
+                    process.exit(1);
+                }
+            }
 
-    if (file.zir.hasCompileErrors()) {
-        var wip_errors: std.zig.ErrorBundle.Wip = undefined;
-        try wip_errors.init(gpa);
-        defer wip_errors.deinit();
-        try Compilation.addZirErrorMessages(&wip_errors, &file);
-        var error_bundle = try wip_errors.toOwnedBundle("");
-        defer error_bundle.deinit(gpa);
-        error_bundle.renderToStdErr(color.renderOptions());
+            if (!want_output_text) {
+                if (file.zir.hasCompileErrors()) {
+                    process.exit(1);
+                } else {
+                    return cleanExit();
+                }
+            }
+            if (!build_options.enable_debug_extensions) {
+                fatal("-t option only available in builds of zig with debug extensions", .{});
+            }
 
-        if (file.zir.loweringFailed()) {
-            process.exit(1);
-        }
-    }
+            {
+                const token_bytes = @sizeOf(Ast.TokenList) +
+                    file.tree.tokens.len * (@sizeOf(std.zig.Token.Tag) + @sizeOf(Ast.ByteOffset));
+                const tree_bytes = @sizeOf(Ast) + file.tree.nodes.len *
+                    (@sizeOf(Ast.Node.Tag) +
+                    @sizeOf(Ast.Node.Data) +
+                    @sizeOf(Ast.TokenIndex));
+                const instruction_bytes = file.zir.instructions.len *
+                    // Here we don't use @sizeOf(Zir.Inst.Data) because it would include
+                    // the debug safety tag but we want to measure release size.
+                    (@sizeOf(Zir.Inst.Tag) + 8);
+                const extra_bytes = file.zir.extra.len * @sizeOf(u32);
+                const total_bytes = @sizeOf(Zir) + instruction_bytes + extra_bytes +
+                    file.zir.string_bytes.len * @sizeOf(u8);
+                const stdout = io.getStdOut();
+                const fmtIntSizeBin = std.fmt.fmtIntSizeBin;
+                // zig fmt: off
+                try stdout.writer().print(
+                    \\# Source bytes:       {}
+                    \\# Tokens:             {} ({})
+                    \\# AST Nodes:          {} ({})
+                    \\# Total ZIR bytes:    {}
+                    \\# Instructions:       {d} ({})
+                    \\# String Table Bytes: {}
+                    \\# Extra Data Items:   {d} ({})
+                    \\
+                , .{
+                    fmtIntSizeBin(file.source.len),
+                    file.tree.tokens.len, fmtIntSizeBin(token_bytes),
+                    file.tree.nodes.len, fmtIntSizeBin(tree_bytes),
+                    fmtIntSizeBin(total_bytes),
+                    file.zir.instructions.len, fmtIntSizeBin(instruction_bytes),
+                    fmtIntSizeBin(file.zir.string_bytes.len),
+                    file.zir.extra.len, fmtIntSizeBin(extra_bytes),
+                });
+                // zig fmt: on
+            }
 
-    if (!want_output_text) {
-        if (file.zir.hasCompileErrors()) {
-            process.exit(1);
-        } else {
-            return cleanExit();
-        }
-    }
-    if (!build_options.enable_debug_extensions) {
-        fatal("-t option only available in builds of zig with debug extensions", .{});
-    }
+            try @import("print_zir.zig").renderAsTextToFile(gpa, &file, io.getStdOut());
 
-    {
-        const token_bytes = @sizeOf(Ast.TokenList) +
-            file.tree.tokens.len * (@sizeOf(std.zig.Token.Tag) + @sizeOf(Ast.ByteOffset));
-        const tree_bytes = @sizeOf(Ast) + file.tree.nodes.len *
-            (@sizeOf(Ast.Node.Tag) +
-            @sizeOf(Ast.Node.Data) +
-            @sizeOf(Ast.TokenIndex));
-        const instruction_bytes = file.zir.instructions.len *
-            // Here we don't use @sizeOf(Zir.Inst.Data) because it would include
-            // the debug safety tag but we want to measure release size.
-            (@sizeOf(Zir.Inst.Tag) + 8);
-        const extra_bytes = file.zir.extra.len * @sizeOf(u32);
-        const total_bytes = @sizeOf(Zir) + instruction_bytes + extra_bytes +
-            file.zir.string_bytes.len * @sizeOf(u8);
-        const stdout = io.getStdOut();
-        const fmtIntSizeBin = std.fmt.fmtIntSizeBin;
-        // zig fmt: off
-        try stdout.writer().print(
-            \\# Source bytes:       {}
-            \\# Tokens:             {} ({})
-            \\# AST Nodes:          {} ({})
-            \\# Total ZIR bytes:    {}
-            \\# Instructions:       {d} ({})
-            \\# String Table Bytes: {}
-            \\# Extra Data Items:   {d} ({})
-            \\
-        , .{
-            fmtIntSizeBin(file.source.len),
-            file.tree.tokens.len, fmtIntSizeBin(token_bytes),
-            file.tree.nodes.len, fmtIntSizeBin(tree_bytes),
-            fmtIntSizeBin(total_bytes),
-            file.zir.instructions.len, fmtIntSizeBin(instruction_bytes),
-            fmtIntSizeBin(file.zir.string_bytes.len),
-            file.zir.extra.len, fmtIntSizeBin(extra_bytes),
-        });
-        // zig fmt: on
-    }
+            if (file.zir.hasCompileErrors()) {
+                process.exit(1);
+            } else {
+                return cleanExit();
+            }
+        },
+        .zon => {
+            const zoir = try ZonGen.generate(gpa, file.tree);
+            defer zoir.deinit(gpa);
 
-    try @import("print_zir.zig").renderAsTextToFile(gpa, &file, io.getStdOut());
+            if (zoir.hasCompileErrors()) {
+                var wip_errors: std.zig.ErrorBundle.Wip = undefined;
+                try wip_errors.init(gpa);
+                defer wip_errors.deinit();
 
-    if (file.zir.hasCompileErrors()) {
-        process.exit(1);
-    } else {
-        return cleanExit();
+                {
+                    const src_path = try file.fullPath(gpa);
+                    defer gpa.free(src_path);
+                    try wip_errors.addZoirErrorMessages(zoir, file.tree, file.source, src_path);
+                }
+
+                var error_bundle = try wip_errors.toOwnedBundle("");
+                defer error_bundle.deinit(gpa);
+                error_bundle.renderToStdErr(color.renderOptions());
+
+                process.exit(1);
+            }
+
+            if (!want_output_text) {
+                return cleanExit();
+            }
+
+            if (!build_options.enable_debug_extensions) {
+                fatal("-t option only available in builds of zig with debug extensions", .{});
+            }
+
+            try @import("print_zoir.zig").renderToFile(zoir, arena, io.getStdOut());
+            return cleanExit();
+        },
     }
 }
 
diff --git a/src/print_zoir.zig b/src/print_zoir.zig
new file mode 100644
index 0000000000..b6cc8fe4d9
--- /dev/null
+++ b/src/print_zoir.zig
@@ -0,0 +1,122 @@
+pub fn renderToFile(zoir: Zoir, arena: Allocator, f: std.fs.File) (std.fs.File.WriteError || Allocator.Error)!void {
+    var bw = std.io.bufferedWriter(f.writer());
+    try renderToWriter(zoir, arena, bw.writer());
+    try bw.flush();
+}
+
+pub fn renderToWriter(zoir: Zoir, arena: Allocator, w: anytype) (@TypeOf(w).Error || Allocator.Error)!void {
+    assert(!zoir.hasCompileErrors());
+
+    const fmtIntSizeBin = std.fmt.fmtIntSizeBin;
+    const bytes_per_node = comptime n: {
+        var n: usize = 0;
+        for (@typeInfo(Zoir.Node.Repr).@"struct".fields) |f| {
+            n += @sizeOf(f.type);
+        }
+        break :n n;
+    };
+
+    const node_bytes = zoir.nodes.len * bytes_per_node;
+    const extra_bytes = zoir.extra.len * @sizeOf(u32);
+    const limb_bytes = zoir.limbs.len * @sizeOf(std.math.big.Limb);
+    const string_bytes = zoir.string_bytes.len;
+
+    // zig fmt: off
+    try w.print(
+        \\# Nodes:              {} ({})
+        \\# Extra Data Items:   {} ({})
+        \\# BigInt Limbs:       {} ({})
+        \\# String Table Bytes: {}
+        \\# Total ZON Bytes:    {}
+        \\
+    , .{
+        zoir.nodes.len, fmtIntSizeBin(node_bytes),
+        zoir.extra.len, fmtIntSizeBin(extra_bytes),
+        zoir.limbs.len, fmtIntSizeBin(limb_bytes),
+        fmtIntSizeBin(string_bytes),
+        fmtIntSizeBin(node_bytes + extra_bytes + limb_bytes + string_bytes),
+    });
+    // zig fmt: on
+    var pz: PrintZon = .{
+        .w = w.any(),
+        .arena = arena,
+        .zoir = zoir,
+        .indent = 0,
+    };
+
+    return @errorCast(pz.renderRoot());
+}
+
+const PrintZon = struct {
+    w: std.io.AnyWriter,
+    arena: Allocator,
+    zoir: Zoir,
+    indent: u32,
+
+    fn renderRoot(pz: *PrintZon) anyerror!void {
+        try pz.renderNode(.root);
+        try pz.w.writeByte('\n');
+    }
+
+    fn renderNode(pz: *PrintZon, node: Zoir.Node.Index) anyerror!void {
+        const zoir = pz.zoir;
+        try pz.w.print("%{d} = ", .{@intFromEnum(node)});
+        switch (node.get(zoir)) {
+            .true => try pz.w.writeAll("true"),
+            .false => try pz.w.writeAll("false"),
+            .null => try pz.w.writeAll("null"),
+            .pos_inf => try pz.w.writeAll("inf"),
+            .neg_inf => try pz.w.writeAll("-inf"),
+            .nan => try pz.w.writeAll("nan"),
+            .int_literal => |storage| switch (storage) {
+                .small => |x| try pz.w.print("int({d})", .{x}),
+                .big => |x| {
+                    const str = try x.toStringAlloc(pz.arena, 10, .lower);
+                    try pz.w.print("int(big {s})", .{str});
+                },
+            },
+            .float_literal => |x| try pz.w.print("float({d})", .{x}),
+            .char_literal => |x| try pz.w.print("char({d})", .{x}),
+            .enum_literal => |x| try pz.w.print("enum_literal({p})", .{std.zig.fmtId(x.get(zoir))}),
+            .string_literal => |x| try pz.w.print("str(\"{}\")", .{std.zig.fmtEscapes(x)}),
+            .empty_literal => try pz.w.writeAll("empty_literal(.{})"),
+            .array_literal => |vals| {
+                try pz.w.writeAll("array_literal({");
+                pz.indent += 1;
+                for (0..vals.len) |idx| {
+                    try pz.newline();
+                    try pz.renderNode(vals.at(@intCast(idx)));
+                    try pz.w.writeByte(',');
+                }
+                pz.indent -= 1;
+                try pz.newline();
+                try pz.w.writeAll("})");
+            },
+            .struct_literal => |s| {
+                try pz.w.writeAll("struct_literal({");
+                pz.indent += 1;
+                for (s.names, 0..s.vals.len) |name, idx| {
+                    try pz.newline();
+                    try pz.w.print("[{p}] ", .{std.zig.fmtId(name.get(zoir))});
+                    try pz.renderNode(s.vals.at(@intCast(idx)));
+                    try pz.w.writeByte(',');
+                }
+                pz.indent -= 1;
+                try pz.newline();
+                try pz.w.writeAll("})");
+            },
+        }
+    }
+
+    fn newline(pz: *PrintZon) !void {
+        try pz.w.writeByte('\n');
+        for (0..pz.indent) |_| {
+            try pz.w.writeByteNTimes(' ', 2);
+        }
+    }
+};
+
+const std = @import("std");
+const assert = std.debug.assert;
+const Allocator = std.mem.Allocator;
+const Zoir = std.zig.Zoir;