allow query to return multiple values

4 files changed, 103 insertions, 51 deletions

diff --git a/src/jq.zig b/src/jq.zig
index cf0606b..cf29e78 100644
--- a/src/jq.zig
+++ b/src/jq.zig
@@ -13,4 +13,4 @@ pub const Instr = @import("./jq/compile.zig").Instr;
 pub const compile = @import("./jq/compile.zig").compile;
 
 pub const ExecuteError = @import("./jq/execute.zig").ExecuteError;
-pub const execute = @import("./jq/execute.zig").execute;
+pub const Runtime = @import("./jq/execute.zig").Runtime;
diff --git a/src/jq/compile.zig b/src/jq/compile.zig
index f4044c9..5719f9c 100644
--- a/src/jq/compile.zig
+++ b/src/jq/compile.zig
@@ -4,6 +4,7 @@ const Ast = @import("./parse.zig").Ast;
 
 pub const Opcode = enum {
     nop,
+    ret,
     subexp_begin,
     subexp_end,
     array_index,
@@ -14,6 +15,7 @@ pub const Opcode = enum {
 
 pub const Instr = union(Opcode) {
     nop,
+    ret,
     subexp_begin,
     subexp_end,
     array_index,
@@ -26,13 +28,13 @@ pub const Instr = union(Opcode) {
     }
 };
 
-pub fn compile(allocator: std.mem.Allocator, compile_allocator: std.mem.Allocator, ast: *const Ast) ![]Instr {
+fn compileExpr(allocator: std.mem.Allocator, compile_allocator: std.mem.Allocator, ast: *const Ast) ![]Instr {
     var instrs = try std.ArrayList(Instr).initCapacity(allocator, 16);
 
     switch (ast.*) {
         .identity => try instrs.append(allocator, .nop),
         .array_index => |index| {
-            const index_instrs = try compile(allocator, compile_allocator, index);
+            const index_instrs = try compileExpr(allocator, compile_allocator, index);
             defer allocator.free(index_instrs);
             try instrs.append(allocator, .subexp_begin);
             try instrs.appendSlice(allocator, index_instrs);
@@ -42,9 +44,9 @@ pub fn compile(allocator: std.mem.Allocator, compile_allocator: std.mem.Allocato
         .object_key => |key| try instrs.append(allocator, .{ .object_key = key }),
         .literal => |value| try instrs.append(allocator, .{ .literal = value }),
         .binary_expr => |binary_expr| {
-            const rhs_instrs = try compile(allocator, compile_allocator, binary_expr.rhs);
+            const rhs_instrs = try compileExpr(allocator, compile_allocator, binary_expr.rhs);
             defer allocator.free(rhs_instrs);
-            const lhs_instrs = try compile(allocator, compile_allocator, binary_expr.lhs);
+            const lhs_instrs = try compileExpr(allocator, compile_allocator, binary_expr.lhs);
             defer allocator.free(lhs_instrs);
             try instrs.append(allocator, .subexp_begin);
             try instrs.appendSlice(allocator, rhs_instrs);
@@ -55,9 +57,9 @@ pub fn compile(allocator: std.mem.Allocator, compile_allocator: std.mem.Allocato
             try instrs.append(allocator, .add);
         },
         .pipe => |pipe_expr| {
-            const lhs_instrs = try compile(allocator, compile_allocator, pipe_expr.lhs);
+            const lhs_instrs = try compileExpr(allocator, compile_allocator, pipe_expr.lhs);
             defer allocator.free(lhs_instrs);
-            const rhs_instrs = try compile(allocator, compile_allocator, pipe_expr.rhs);
+            const rhs_instrs = try compileExpr(allocator, compile_allocator, pipe_expr.rhs);
             defer allocator.free(rhs_instrs);
             try instrs.appendSlice(allocator, lhs_instrs);
             try instrs.appendSlice(allocator, rhs_instrs);
@@ -66,3 +68,12 @@ pub fn compile(allocator: std.mem.Allocator, compile_allocator: std.mem.Allocato
 
     return instrs.toOwnedSlice(allocator);
 }
+
+pub fn compile(allocator: std.mem.Allocator, compile_allocator: std.mem.Allocator, ast: *const Ast) ![]Instr {
+    var instrs = try std.ArrayList(Instr).initCapacity(allocator, 16);
+    const expr_instrs = try compileExpr(allocator, compile_allocator, ast);
+    defer allocator.free(expr_instrs);
+    try instrs.appendSlice(allocator, expr_instrs);
+    try instrs.append(allocator, .ret);
+    return instrs.toOwnedSlice(allocator);
+}
diff --git a/src/jq/execute.zig b/src/jq/execute.zig
index 5108a64..3916d5b 100644
--- a/src/jq/execute.zig
+++ b/src/jq/execute.zig
@@ -102,53 +102,73 @@ const ValueStack = struct {
     }
 };
 
-pub fn execute(allocator: std.mem.Allocator, instrs: []const Instr, input: jv.Value) !jv.Value {
-    var value_stack = try ValueStack.init(allocator);
-    defer value_stack.deinit();
+pub const Runtime = struct {
+    const Self = @This();
+
+    values: ValueStack,
+    instrs: []const Instr,
+    input: jv.Value,
+    pc: usize,
+
+    pub fn init(allocator: std.mem.Allocator, instrs: []const Instr, input: jv.Value) !Self {
+        var self = Self{
+            .values = try ValueStack.init(allocator),
+            .instrs = instrs,
+            .input = input,
+            .pc = 0,
+        };
+        try self.values.push(input);
+        return self;
+    }
 
-    try value_stack.push(input);
+    pub fn deinit(self: *Self) void {
+        self.values.deinit();
+    }
 
-    const len = instrs.len;
-    var pc: usize = 0;
-    while (pc < len) {
-        const cur = instrs[pc];
-        switch (cur) {
-            .nop => {},
-            .subexp_begin => try value_stack.dup(),
-            .subexp_end => try value_stack.swap(),
-            .array_index => {
-                std.debug.assert(value_stack.ensureSize(2));
+    pub fn next(self: *Self) !?jv.Value {
+        while (self.pc < self.instrs.len) : (self.pc += 1) {
+            const cur = self.instrs[self.pc];
+            switch (cur) {
+                .nop => {},
+                .ret => {
+                    self.pc += 1;
+                    return self.values.pop();
+                },
+                .subexp_begin => try self.values.dup(),
+                .subexp_end => try self.values.swap(),
+                .array_index => {
+                    std.debug.assert(self.values.ensureSize(2));
 
-                const array = try value_stack.popArray();
-                const index: usize = @intCast(try value_stack.popInteger());
-                const result = if (index < array.items.len) array.items[index] else .null;
-                try value_stack.push(result);
-            },
-            .add => {
-                std.debug.assert(value_stack.ensureSize(3));
+                    const array = try self.values.popArray();
+                    const index: usize = @intCast(try self.values.popInteger());
+                    const result = if (index < array.items.len) array.items[index] else .null;
+                    try self.values.push(result);
+                },
+                .add => {
+                    std.debug.assert(self.values.ensureSize(3));
 
-                _ = value_stack.pop();
-                const lhs = try value_stack.popInteger();
-                const rhs = try value_stack.popInteger();
-                const result = lhs + rhs;
-                try value_stack.push(.{ .integer = result });
-            },
-            .object_key => |key| {
-                std.debug.assert(value_stack.ensureSize(1));
+                    _ = self.values.pop();
+                    const lhs = try self.values.popInteger();
+                    const rhs = try self.values.popInteger();
+                    const result = lhs + rhs;
+                    try self.values.push(.{ .integer = result });
+                },
+                .object_key => |key| {
+                    std.debug.assert(self.values.ensureSize(1));
 
-                const obj = try value_stack.popObject();
-                const result = obj.get(key) orelse .null;
-                try value_stack.push(result);
-            },
-            .literal => |value| {
-                std.debug.assert(value_stack.ensureSize(1));
+                    const obj = try self.values.popObject();
+                    const result = obj.get(key) orelse .null;
+                    try self.values.push(result);
+                },
+                .literal => |value| {
+                    std.debug.assert(self.values.ensureSize(1));
 
-                _ = value_stack.pop();
-                try value_stack.push(value.*);
-            },
+                    _ = self.values.pop();
+                    try self.values.push(value.*);
+                },
+            }
         }
-        pc += 1;
-    }
 
-    return value_stack.pop();
-}
+        return null;
+    }
+};
diff --git a/src/root.zig b/src/root.zig
index b89b54f..9c22deb 100644
--- a/src/root.zig
+++ b/src/root.zig
@@ -14,15 +14,36 @@ pub fn run(allocator: std.mem.Allocator, input: []const u8, query: []const u8) !
     const parsed = try jv.parse(allocator, input);
     defer parsed.deinit();
     const json = parsed.value;
-    const result = try jq.execute(allocator, instrs, json);
+
+    var runtime = try jq.Runtime.init(allocator, instrs, json);
+    defer runtime.deinit();
+    const result = try runtime.next() orelse return error.NoResult;
     const output = try jv.stringify(allocator, result);
     return output;
 }
 
 fn testRun(expected: []const u8, allocator: std.mem.Allocator, input: []const u8, query: []const u8) !void {
-    const result = try run(allocator, input, query);
+    var compile_allocator = std.heap.ArenaAllocator.init(allocator);
+    defer compile_allocator.deinit();
+    var reader = std.Io.Reader.fixed(query);
+    const tokens = try jq.tokenize(compile_allocator.allocator(), &reader);
+    const ast = try jq.parse(compile_allocator.allocator(), tokens);
+    const instrs = try jq.compile(allocator, compile_allocator.allocator(), ast);
+    defer allocator.free(instrs);
+
+    const parsed = try jv.parse(allocator, input);
+    defer parsed.deinit();
+    const json = parsed.value;
+
+    var runtime = try jq.Runtime.init(allocator, instrs, json);
+    defer runtime.deinit();
+
+    const result_value = try runtime.next() orelse return error.NoResult;
+    const result = try jv.stringify(allocator, result_value);
     defer allocator.free(result);
     try std.testing.expectEqualStrings(expected, result);
+
+    try std.testing.expectEqual(null, try runtime.next());
 }
 
 test "identity filter" {