support --debug-dump-disasm

5 files changed, 198 insertions, 10 deletions

diff --git a/justfile b/justfile
index 96c3831..df88581 100644
--- a/justfile
+++ b/justfile
@@ -5,7 +5,7 @@ build: fmt
     @zig build
 
 run QUERY:
-    @zig build run -- {{QUERY}}
+    @zig build run -- --debug-dump-disasm "{{QUERY}}"
 
 test: fmt
     @zig build test
diff --git a/src/jq/codegen.zig b/src/jq/codegen.zig
index 23383bf..1fc889d 100644
--- a/src/jq/codegen.zig
+++ b/src/jq/codegen.zig
@@ -38,6 +38,22 @@ pub const Opcode = enum {
     store,
     append,
     each,
+
+    pub fn name(self: Opcode) []const u8 {
+        switch (self) {
+            inline else => |tag| {
+                const tag_name = @tagName(tag);
+                const upper_tag_name = comptime blk: {
+                    var buf: [tag_name.len]u8 = undefined;
+                    for (tag_name, 0..) |c, i| {
+                        buf[i] = std.ascii.toUpper(c);
+                    }
+                    break :blk buf;
+                };
+                return &upper_tag_name;
+            },
+        }
+    }
 };
 
 pub const Instr = union(Opcode) {
@@ -304,3 +320,26 @@ pub fn codegen(allocator: std.mem.Allocator, ast: *const Ast) ![]Instr {
     try gen.emit(.ret);
     return gen.toOwnedSlice();
 }
+
+pub fn disasm(allocator: std.mem.Allocator, instrs: []const Instr, constants: []const jv.Value, writer: *std.Io.Writer) !void {
+    for (instrs, 0..) |instr, i| {
+        try writer.print("{d:0>4} ", .{i});
+        try writer.writeAll(instr.op().name());
+        switch (instr) {
+            // program address
+            .jump, .jump_unless, .fork => |offset| try writer.print(" {d:0>4}", .{i + offset}),
+            // constant index
+            .@"const" => |idx| {
+                try writer.writeAll(" ");
+                const k = try jv.stringify(allocator, constants[@intFromEnum(idx)]);
+                defer allocator.free(k);
+                try writer.writeAll(k);
+            },
+            // variable index
+            .load, .store, .append => |idx| try writer.print(" ${d}", .{@intFromEnum(idx)}),
+            // no arguments
+            else => {},
+        }
+        try writer.writeAll("\n");
+    }
+}
diff --git a/src/jq/execute.zig b/src/jq/execute.zig
index d8081e6..cee7845 100644
--- a/src/jq/execute.zig
+++ b/src/jq/execute.zig
@@ -4,6 +4,7 @@ const tokenize = @import("./tokenize.zig").tokenize;
 const parse = @import("./parse.zig").parse;
 const Instr = @import("./codegen.zig").Instr;
 const codegen = @import("./codegen.zig").codegen;
+const disasm = @import("./codegen.zig").disasm;
 
 pub const ExecuteError = error{
     Unimplemented,
@@ -189,11 +190,6 @@ pub const Runtime = struct {
         const ast = try parse(self.allocator, compile_allocator.allocator(), tokens, &self.constants);
         const instrs = try codegen(self.allocator, ast);
         self.instrs = instrs;
-        // std.debug.print("BEGIN\n", .{});
-        // for (self.instrs) |instr| {
-        //     std.debug.print("{}\n", .{instr});
-        // }
-        // std.debug.print("END\n", .{});
     }
 
     pub fn compileFromSlice(self: *Self, query: []const u8) !void {
@@ -201,6 +197,10 @@ pub const Runtime = struct {
         return self.compileFromReader(&reader);
     }
 
+    pub fn dumpDisasm(self: *const Self, writer: *std.Io.Writer) !void {
+        try disasm(self.allocator, self.instrs, self.constants.items, writer);
+    }
+
     pub fn start(self: *Self, input: jv.Value) !void {
         try self.values.push(input.clone());
     }
diff --git a/src/main.zig b/src/main.zig
index 73d7886..51b2262 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -8,11 +8,20 @@ pub fn main() !void {
     const args = try std.process.argsAlloc(allocator);
     defer std.process.argsFree(allocator, args);
 
-    if (args.len < 2) {
-        try std.fs.File.stderr().writeAll("usage: zgjq <query>\n");
+    var debug_dump_disasm = false;
+    var query: ?[]const u8 = null;
+    for (args[1..]) |arg| {
+        if (std.mem.eql(u8, arg, "--debug-dump-disasm")) {
+            debug_dump_disasm = true;
+        } else if (query == null) {
+            query = arg;
+        }
+    }
+
+    if (query == null) {
+        try std.fs.File.stderr().writeAll("usage: zgjq [--debug-dump-disasm] <query>\n");
         std.process.exit(1);
     }
-    const query = args[1];
 
     const input = try std.fs.File.stdin().readToEndAlloc(allocator, std.math.maxInt(usize));
     defer allocator.free(input);
@@ -23,7 +32,15 @@ pub fn main() !void {
 
     var runtime = try jq.Runtime.init(allocator);
     defer runtime.deinit();
-    try runtime.compileFromSlice(query);
+    try runtime.compileFromSlice(query.?);
+
+    if (debug_dump_disasm) {
+        var buf: [4096]u8 = undefined;
+        var stderr = std.fs.File.stderr().writer(&buf);
+        try runtime.dumpDisasm(&stderr.interface);
+        try stderr.interface.flush();
+    }
+
     try runtime.start(json);
 
     const stdout = std.fs.File.stdout();
diff --git a/src/root.zig b/src/root.zig
index c81ae61..e7e0f6b 100644
--- a/src/root.zig
+++ b/src/root.zig
@@ -44,6 +44,21 @@ fn testRunMultiple(expected: []const []const u8, input: []const u8, query: []con
     try std.testing.expectEqual(null, try runtime.next());
 }
 
+fn testDisasm(expected: []const u8, query: []const u8) !void {
+    const allocator = std.testing.allocator;
+
+    var runtime = try jq.Runtime.init(allocator);
+    defer runtime.deinit();
+    try runtime.compileFromSlice(query);
+
+    var buf: [4096]u8 = undefined;
+    var writer = std.Io.Writer.fixed(&buf);
+    try runtime.dumpDisasm(&writer);
+    try writer.flush();
+
+    try std.testing.expectEqualStrings(expected, buf[0..writer.end]);
+}
+
 test "literals" {
     try testRun("\"hello\"", "null", "\"hello\"");
     try testRun("\"\"", "null", "\"\"");
@@ -397,3 +412,120 @@ test "each" {
     try testRunMultiple(&.{ "1", "2", "3" }, "[1,2,3]", ".[]");
     try testRunMultiple(&.{ "1", "2", "3" }, "[[1],[2],[3]]", ".[] | .[]");
 }
+
+test "disasm identity" {
+    try testDisasm(
+        \\0000 NOP
+        \\0001 RET
+        \\
+    , ".");
+}
+
+test "disasm field access" {
+    try testDisasm(
+        \\0000 NOP
+        \\0001 SUBEXP_BEGIN
+        \\0002 CONST "foo"
+        \\0003 SUBEXP_END
+        \\0004 INDEX
+        \\0005 RET
+        \\
+    , ".foo");
+}
+
+test "disasm pipe and comma" {
+    try testDisasm(
+        \\0000 NOP
+        \\0001 SUBEXP_BEGIN
+        \\0002 CONST "foo"
+        \\0003 SUBEXP_END
+        \\0004 INDEX
+        \\0005 FORK 0012
+        \\0006 NOP
+        \\0007 SUBEXP_BEGIN
+        \\0008 CONST "bar"
+        \\0009 SUBEXP_END
+        \\0010 INDEX
+        \\0011 JUMP 0017
+        \\0012 NOP
+        \\0013 SUBEXP_BEGIN
+        \\0014 CONST "baz"
+        \\0015 SUBEXP_END
+        \\0016 INDEX
+        \\0017 RET
+        \\
+    , ".foo | .bar, .baz");
+}
+
+test "disasm array construction" {
+    try testDisasm(
+        \\0000 DUP
+        \\0001 CONST []
+        \\0002 STORE $0
+        \\0003 NOP
+        \\0004 EACH
+        \\0005 APPEND $0
+        \\0006 BACKTRACK
+        \\0007 LOAD $0
+        \\0008 RET
+        \\
+    , "[.[]]");
+}
+
+test "disasm arithmetic" {
+    try testDisasm(
+        \\0000 SUBEXP_BEGIN
+        \\0001 CONST 1
+        \\0002 SUBEXP_END
+        \\0003 SUBEXP_BEGIN
+        \\0004 NOP
+        \\0005 SUBEXP_END
+        \\0006 ADD
+        \\0007 RET
+        \\
+    , ". + 1");
+}
+
+test "disasm optional index" {
+    try testDisasm(
+        \\0000 NOP
+        \\0001 SUBEXP_BEGIN
+        \\0002 CONST "foo"
+        \\0003 SUBEXP_END
+        \\0004 INDEX_OPT
+        \\0005 RET
+        \\
+    , ".foo?");
+}
+
+test "disasm comparison" {
+    try testDisasm(
+        \\0000 SUBEXP_BEGIN
+        \\0001 CONST 1
+        \\0002 SUBEXP_END
+        \\0003 SUBEXP_BEGIN
+        \\0004 NOP
+        \\0005 SUBEXP_END
+        \\0006 EQ
+        \\0007 RET
+        \\
+    , ". == 1");
+}
+
+test "disasm constants" {
+    try testDisasm(
+        \\0000 CONST null
+        \\0001 RET
+        \\
+    , "null");
+    try testDisasm(
+        \\0000 CONST true
+        \\0001 RET
+        \\
+    , "true");
+    try testDisasm(
+        \\0000 CONST false
+        \\0001 RET
+        \\
+    , "false");
+}