implement array index filter

5 files changed, 153 insertions, 31 deletions

diff --git a/src/jq/compile.zig b/src/jq/compile.zig
index d81f8ac..a28c10a 100644
--- a/src/jq/compile.zig
+++ b/src/jq/compile.zig
@@ -1,21 +1,35 @@
 const std = @import("std");
+const jv = @import("../jv.zig");
 const Ast = @import("./parse.zig").Ast;
 
 pub const Opcode = enum {
     nop,
-    identity,
+    array_index,
+    literal,
 };
 
-pub const Instr = struct {
-    op: Opcode,
+pub const Instr = union(Opcode) {
+    nop,
+    array_index,
+    literal: *jv.Value,
+
+    pub fn op(self: @This()) Opcode {
+        return self;
+    }
 };
 
 pub fn compile(allocator: std.mem.Allocator, compile_allocator: std.mem.Allocator, ast: *const Ast) ![]Instr {
-    _ = compile_allocator;
     var instrs = try std.array_list.Aligned(Instr, null).initCapacity(allocator, 16);
 
-    switch (ast.kind) {
-        .identity => try instrs.append(allocator, .{ .op = .identity }),
+    switch (ast.*) {
+        .identity => try instrs.append(allocator, .nop),
+        .array_index => |index| {
+            const index_instrs = try compile(allocator, compile_allocator, index);
+            defer allocator.free(index_instrs);
+            try instrs.appendSlice(allocator, index_instrs);
+            try instrs.append(allocator, .array_index);
+        },
+        .literal => |value| try instrs.append(allocator, .{ .literal = value }),
     }
 
     return instrs.toOwnedSlice(allocator);
diff --git a/src/jq/execute.zig b/src/jq/execute.zig
index 9857d46..dfad6c4 100644
--- a/src/jq/execute.zig
+++ b/src/jq/execute.zig
@@ -4,19 +4,44 @@ const Instr = @import("./compile.zig").Instr;
 
 pub const ExecuteError = error{
     Unimplemented,
+    InvalidType,
+    InternalError,
 };
 
 pub fn execute(allocator: std.mem.Allocator, instrs: []const Instr, input: jv.Value) !jv.Value {
-    _ = allocator;
+    var value_stack = try std.array_list.Aligned(jv.Value, null).initCapacity(allocator, 16);
+    defer value_stack.deinit(allocator);
+
+    try value_stack.append(allocator, input);
+
     const len = instrs.len;
     var pc: usize = 0;
     while (pc < len) {
         const cur = instrs[pc];
-        _ = switch (cur.op) {
-            .nop => void,
-            .identity => return input,
-        };
+        switch (cur) {
+            .nop => {},
+            .array_index => {
+                const v1 = value_stack.pop() orelse return error.InternalError;
+                const v1_integer = switch (v1) {
+                    .integer => |integer| integer,
+                    else => return error.InvalidType,
+                };
+                const v2 = value_stack.pop() orelse return error.InternalError;
+                const v2_array = switch (v2) {
+                    .array => |array| array,
+                    else => return error.InvalidType,
+                };
+                const index: usize = @intCast(v1_integer);
+                const result = if (index < v2_array.items.len) v2_array.items[index] else .null;
+                try value_stack.append(allocator, result);
+            },
+            .literal => |value| {
+                try value_stack.append(allocator, value.*);
+            },
+        }
         pc += 1;
     }
-    return ExecuteError.Unimplemented;
+
+    const result = value_stack.pop() orelse return error.InternalError;
+    return result;
 }
diff --git a/src/jq/parse.zig b/src/jq/parse.zig
index 0554dca..5269ef5 100644
--- a/src/jq/parse.zig
+++ b/src/jq/parse.zig
@@ -1,4 +1,5 @@
 const std = @import("std");
+const jv = @import("../jv.zig");
 const Token = @import("./tokenize.zig").Token;
 
 pub const ParseError = error{
@@ -8,26 +9,74 @@ pub const ParseError = error{
 
 pub const AstKind = enum {
     identity,
+    array_index,
+    literal,
 };
 
-pub const Ast = struct {
-    kind: AstKind,
+pub const Ast = union(AstKind) {
+    identity,
+    array_index: *Ast,
+    literal: *jv.Value,
+
+    pub fn kind(self: @This()) AstKind {
+        return self;
+    }
 };
 
 pub fn parse(allocator: std.mem.Allocator, tokens: []const Token) !*Ast {
-    if (tokens.len != 2) {
-        return ParseError.InvalidQuery;
+    if (tokens.len < 2) {
+        return error.InvalidQuery;
+    }
+
+    var i: usize = 0;
+    const t1 = tokens[i];
+    if (t1.kind() != .dot) {
+        return error.InvalidQuery;
+    }
+    i += 1;
+    const t2 = tokens[i];
+
+    if (t2.kind() == .end) {
+        const root = try allocator.create(Ast);
+        root.* = .identity;
+        return root;
+    }
+
+    if (t2.kind() != .bracket_left) {
+        return error.InvalidQuery;
+    }
+
+    i += 1;
+    if (tokens.len < 5) {
+        return error.UnexpectedEnd;
+    }
+    const t3 = tokens[i];
+    i += 1;
+    const t4 = tokens[i];
+    i += 1;
+    const t5 = tokens[i];
+
+    if (t3.kind() != .number) {
+        return error.InvalidQuery;
     }
-    const t1 = tokens[0];
-    const t2 = tokens[1];
-    if (t1.kind != .identity) {
-        return ParseError.InvalidQuery;
+    if (t4.kind() != .bracket_right) {
+        return error.InvalidQuery;
     }
-    if (t2.kind != .end) {
-        return ParseError.UnexpectedEnd;
+    if (t5.kind() != .end) {
+        return error.InvalidQuery;
     }
 
+    const index_value = try allocator.create(jv.Value);
+    index_value.* = .{
+        .integer = t3.number,
+    };
+    const index_node = try allocator.create(Ast);
+    index_node.* = .{
+        .literal = index_value,
+    };
     const root = try allocator.create(Ast);
-    root.kind = .identity;
+    root.* = .{
+        .array_index = index_node,
+    };
     return root;
 }
diff --git a/src/jq/tokenize.zig b/src/jq/tokenize.zig
index 0823ea1..92f96ca 100644
--- a/src/jq/tokenize.zig
+++ b/src/jq/tokenize.zig
@@ -2,32 +2,56 @@ const std = @import("std");
 
 pub const TokenizeError = error{
     UnexpectedEnd,
+    InvalidCharacter,
 };
 
 pub const TokenKind = enum {
     end,
-    identity,
+    dot,
+    bracket_left,
+    bracket_right,
+    number,
 };
 
-pub const Token = struct {
-    kind: TokenKind,
+pub const Token = union(TokenKind) {
+    end,
+    dot,
+    bracket_left,
+    bracket_right,
+    number: i64,
+
+    pub fn kind(self: @This()) TokenKind {
+        return self;
+    }
 };
 
 pub fn tokenize(allocator: std.mem.Allocator, query: []const u8) ![]Token {
     var tokens = try std.array_list.Aligned(Token, null).initCapacity(allocator, 16);
 
     const len = query.len;
+
+    if (len == 0) {
+        return error.UnexpectedEnd;
+    }
+
     var i: usize = 0;
     while (i < len) {
         const c = query[i];
-        if (c == '.') {
-            try tokens.append(allocator, .{ .kind = .identity });
-        } else {
-            return TokenizeError.UnexpectedEnd;
+        switch (c) {
+            '.' => try tokens.append(allocator, .dot),
+            '[' => try tokens.append(allocator, .bracket_left),
+            ']' => try tokens.append(allocator, .bracket_right),
+            else => {
+                if (std.ascii.isDigit(c)) {
+                    try tokens.append(allocator, .{ .number = (c - '0') });
+                } else {
+                    return error.InvalidCharacter;
+                }
+            },
         }
         i += 1;
     }
 
-    try tokens.append(allocator, .{ .kind = .end });
+    try tokens.append(allocator, .end);
     return tokens.toOwnedSlice(allocator);
 }
diff --git a/src/root.zig b/src/root.zig
index 61b61f5..d22a239 100644
--- a/src/root.zig
+++ b/src/root.zig
@@ -39,3 +39,13 @@ test "identity filter" {
     try testRun("[1,2,3]", allocator, "[1,2,3]", ".");
     try testRun("{\"a\":123}", allocator, "{\"a\":123}", ".");
 }
+
+test "array index filter" {
+    var debug_allocator = std.heap.DebugAllocator(.{}).init;
+    defer std.debug.assert(debug_allocator.deinit() == .ok);
+    const allocator = debug_allocator.allocator();
+
+    try testRun("null", allocator, "[]", ".[0]");
+    try testRun("1", allocator, "[1,2,3]", ".[0]");
+    try testRun("null", allocator, "[1,2,3]", ".[5]");
+}