implement array iteration

5 files changed, 53 insertions, 2 deletions

diff --git a/docs/jq_grammar.md b/docs/jq_grammar.md
index f9850cf..697a062 100644
--- a/docs/jq_grammar.md
+++ b/docs/jq_grammar.md
@@ -88,6 +88,7 @@ term:
     primary { suffix }*
 
 suffix:
+    '[' ']'
     '[' query ']' '?'?
     '[' query ':' query ']' '?'?
     '[' query ':' ']' '?'?
diff --git a/src/jq/codegen.zig b/src/jq/codegen.zig
index 62535c1..23383bf 100644
--- a/src/jq/codegen.zig
+++ b/src/jq/codegen.zig
@@ -37,6 +37,7 @@ pub const Opcode = enum {
     load,
     store,
     append,
+    each,
 };
 
 pub const Instr = union(Opcode) {
@@ -72,6 +73,7 @@ pub const Instr = union(Opcode) {
     load: VariableIndex,
     store: VariableIndex,
     append: VariableIndex,
+    each,
 
     pub fn op(self: Self) Opcode {
         return self;
@@ -264,6 +266,12 @@ const Codegen = struct {
                 try self.emit(.backtrack);
                 try self.emit(.{ .load = v });
             },
+            .each => |each| {
+                // <base>
+                // EACH
+                try self.generate(each.base);
+                try self.emit(.each);
+            },
         }
     }
 
diff --git a/src/jq/execute.zig b/src/jq/execute.zig
index 9fa410f..d8081e6 100644
--- a/src/jq/execute.zig
+++ b/src/jq/execute.zig
@@ -208,11 +208,11 @@ pub const Runtime = struct {
     pub fn next(self: *Self) !?jv.Value {
         std.debug.assert(self.instrs.len > 0);
 
-        _ = self.restore_stack();
+        var is_backtracking = self.restore_stack();
 
         while (self.pc < self.instrs.len) : (self.pc += 1) {
             const cur = self.instrs[self.pc];
-            // std.debug.print("{}\n", .{cur});
+            // std.debug.print("{}  ({})\n", .{ cur, self.values.stack.size() });
             switch (cur) {
                 .nop => {},
                 .ret => {
@@ -435,6 +435,32 @@ pub const Runtime = struct {
                     const var_ptr = &self.variables.items[@intFromEnum(idx)];
                     try var_ptr.arrayAppend(self.allocator, self.values.pop());
                 },
+                .each => {
+                    std.debug.assert(self.values.ensureSize(1));
+
+                    const key = if (is_backtracking)
+                        jv.Value.initInteger(self.values.pop().integer + 1)
+                    else
+                        jv.Value.initInteger(0);
+                    const base = self.values.pop();
+                    if (base.array.len() <= key.integer) {
+                        base.deinit(self.allocator);
+                        if (self.restore_stack()) {
+                            self.pc -= 1;
+                            is_backtracking = true;
+                            continue;
+                        } else {
+                            return null;
+                        }
+                    }
+                    const idx_result: jv.Value = try jv.ops.index(base, key);
+                    const result = idx_result.clone();
+                    try self.values.push(base);
+                    try self.values.push(key);
+                    try self.save_stack(self.pc);
+                    try self.values.push(result);
+                    is_backtracking = false;
+                },
             }
         }
 
@@ -442,12 +468,14 @@ pub const Runtime = struct {
     }
 
     fn save_stack(self: *Self, target_pc: usize) !void {
+        // std.debug.print("STACK SAVED: {}\n", .{target_pc});
         try self.forks.append(self.allocator, target_pc);
         try self.values.save();
     }
 
     fn restore_stack(self: *Self) bool {
         if (self.forks.pop()) |target_pc| {
+            // std.debug.print("STACK RESTORED: {}\n", .{target_pc});
             self.pc = target_pc;
             self.values.restore(self.allocator);
             return true;
diff --git a/src/jq/parse.zig b/src/jq/parse.zig
index ad64824..17138ae 100644
--- a/src/jq/parse.zig
+++ b/src/jq/parse.zig
@@ -20,6 +20,7 @@ pub const AstKind = enum {
     pipe,
     comma,
     construct_array,
+    each,
 };
 
 pub const BinaryOp = enum {
@@ -56,6 +57,7 @@ pub const Ast = union(AstKind) {
     pipe: struct { lhs: *Ast, rhs: *Ast },
     comma: struct { lhs: *Ast, rhs: *Ast },
     construct_array: struct { items: *Ast },
+    each: struct { base: *Ast },
 
     pub fn kind(self: @This()) AstKind {
         return self;
@@ -400,6 +402,13 @@ const Parser = struct {
     fn parseSuffix(self: *Self, base: *Ast) Error!*Ast {
         _ = try self.tokens.expect(.bracket_left);
 
+        // Handle [] form.
+        if (self.tokens.consumeIf(.bracket_right)) {
+            const ast = try self.compile_allocator.create(Ast);
+            ast.* = .{ .each = .{ .base = base } };
+            return ast;
+        }
+
         // Handle [:to] form.
         if (self.tokens.consumeIf(.colon)) {
             const to_expr = try self.parseQuery();
diff --git a/src/root.zig b/src/root.zig
index ac9c21c..c81ae61 100644
--- a/src/root.zig
+++ b/src/root.zig
@@ -392,3 +392,8 @@ test "array constructor" {
         \\]
     , "{\"a\":1,\"b\":2}", "[.a, [.a, [.a, .b]]]");
 }
+
+test "each" {
+    try testRunMultiple(&.{ "1", "2", "3" }, "[1,2,3]", ".[]");
+    try testRunMultiple(&.{ "1", "2", "3" }, "[[1],[2],[3]]", ".[] | .[]");
+}