implement arbitrary query in index access

2 files changed, 30 insertions, 25 deletions

diff --git a/src/jq/parse.zig b/src/jq/parse.zig
index 13a035b..58c3722 100644
--- a/src/jq/parse.zig
+++ b/src/jq/parse.zig
@@ -107,25 +107,26 @@ const TokenStream = struct {
 
 const Parser = struct {
     const Self = @This();
+    const Error = ParseError || std.mem.Allocator.Error;
 
     allocator: std.mem.Allocator,
     parse_allocator: std.mem.Allocator,
     tokens: *TokenStream,
     constants: *std.ArrayList(jv.Value),
 
-    fn parseProgram(self: *Self) !*Ast {
+    fn parseProgram(self: *Self) Error!*Ast {
         return self.parseBody();
     }
 
-    fn parseBody(self: *Self) !*Ast {
+    fn parseBody(self: *Self) Error!*Ast {
         return self.parseQuery();
     }
 
-    fn parseQuery(self: *Self) !*Ast {
+    fn parseQuery(self: *Self) Error!*Ast {
         return self.parseQuery2();
     }
 
-    fn parseQuery2(self: *Self) !*Ast {
+    fn parseQuery2(self: *Self) Error!*Ast {
         var lhs = try self.parseQuery3();
         while (self.tokens.consumeIf(.pipe)) {
             const rhs = try self.parseQuery3();
@@ -140,7 +141,7 @@ const Parser = struct {
         return lhs;
     }
 
-    fn parseQuery3(self: *Self) !*Ast {
+    fn parseQuery3(self: *Self) Error!*Ast {
         var lhs = try self.parseExpr();
         while (self.tokens.consumeIf(.comma)) {
             const rhs = try self.parseExpr();
@@ -154,7 +155,7 @@ const Parser = struct {
         return lhs;
     }
 
-    fn parseExpr(self: *Self) !*Ast {
+    fn parseExpr(self: *Self) Error!*Ast {
         var lhs = try self.parseExpr2();
         while (self.tokens.consumeIf(.slash_slash)) {
             const rhs = try self.parseExpr2();
@@ -169,7 +170,7 @@ const Parser = struct {
         return lhs;
     }
 
-    fn parseExpr2(self: *Self) !*Ast {
+    fn parseExpr2(self: *Self) Error!*Ast {
         const lhs = try self.parseExpr3();
         const token = self.tokens.peek() catch return lhs;
         const op: BinaryOp = switch (token.kind()) {
@@ -194,7 +195,7 @@ const Parser = struct {
         return ast;
     }
 
-    fn parseExpr3(self: *Self) !*Ast {
+    fn parseExpr3(self: *Self) Error!*Ast {
         const lhs = try self.parseExpr4();
         if (!self.tokens.consumeIf(.keyword_or)) {
             return lhs;
@@ -209,7 +210,7 @@ const Parser = struct {
         return ast;
     }
 
-    fn parseExpr4(self: *Self) !*Ast {
+    fn parseExpr4(self: *Self) Error!*Ast {
         const lhs = try self.parseExpr5();
         if (!self.tokens.consumeIf(.keyword_and)) {
             return lhs;
@@ -224,7 +225,7 @@ const Parser = struct {
         return ast;
     }
 
-    fn parseExpr5(self: *Self) !*Ast {
+    fn parseExpr5(self: *Self) Error!*Ast {
         const lhs = try self.parseExpr6();
         const token = self.tokens.peek() catch return lhs;
         const op: BinaryOp = switch (token.kind()) {
@@ -247,7 +248,7 @@ const Parser = struct {
         return ast;
     }
 
-    fn parseExpr6(self: *Self) !*Ast {
+    fn parseExpr6(self: *Self) Error!*Ast {
         var lhs = try self.parseExpr7();
         while (true) {
             const token = self.tokens.peek() catch return lhs;
@@ -268,7 +269,7 @@ const Parser = struct {
         }
     }
 
-    fn parseExpr7(self: *Self) !*Ast {
+    fn parseExpr7(self: *Self) Error!*Ast {
         var lhs = try self.parseTerm();
         while (true) {
             const token = self.tokens.peek() catch return lhs;
@@ -290,7 +291,7 @@ const Parser = struct {
         }
     }
 
-    fn parseTerm(self: *Self) !*Ast {
+    fn parseTerm(self: *Self) Error!*Ast {
         var result = try self.parsePrimary();
         while (true) {
             const token = self.tokens.peek() catch return result;
@@ -303,7 +304,7 @@ const Parser = struct {
         return result;
     }
 
-    fn parsePrimary(self: *Self) !*Ast {
+    fn parsePrimary(self: *Self) Error!*Ast {
         const first_token = try self.tokens.peek();
         switch (first_token) {
             .keyword_null => {
@@ -392,20 +393,15 @@ const Parser = struct {
         }
     }
 
-    fn parseSuffix(self: *Self, base: *Ast) !*Ast {
+    fn parseSuffix(self: *Self, base: *Ast) Error!*Ast {
         _ = try self.tokens.expect(.bracket_left);
-        const index_token = try self.tokens.expect(.number);
+        const index_expr = try self.parseExpr();
         _ = try self.tokens.expect(.bracket_right);
 
         const is_optional = self.tokens.consumeIf(.question);
 
-        try self.constants.append(self.allocator, .{ .integer = @intFromFloat(index_token.number) });
-        const idx: ConstIndex = @enumFromInt(self.constants.items.len - 1);
-        const index_node = try self.parse_allocator.create(Ast);
-        index_node.* = .{ .literal = idx };
-
         const ast = try self.parse_allocator.create(Ast);
-        ast.* = .{ .index = .{ .base = base, .index = index_node, .is_optional = is_optional } };
+        ast.* = .{ .index = .{ .base = base, .index = index_expr, .is_optional = is_optional } };
         return ast;
     }
 };
diff --git a/src/root.zig b/src/root.zig
index b1f8614..6cdc276 100644
--- a/src/root.zig
+++ b/src/root.zig
@@ -65,7 +65,7 @@ test "identity filter" {
     try testRun("{\"a\":123}", "{\"a\":123}", ".");
 }
 
-test "array index filter" {
+test "index access" {
     try testRun("null", "[]", ".[0]");
     try testRun("1", "[1,2,3]", ".[0]");
     try testRun("null", "[1,2,3]", ".[5]");
@@ -77,14 +77,23 @@ test "array index filter" {
         \\ 61,62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,
         \\ 81,82,83,84,85,86,87,88,89,90,91,92,93,94,95,96,97,98,99,100]
     , ".[100]");
-}
 
-test "object key filter" {
     try testRun("123", "{\"a\":123}", ".a");
     try testRun("null", "{\"a\":123}", ".b");
     try testRun("\"hello\"", "{\"foo\":\"hello\"}", ".foo");
     try testRun("[1,2,3]", "{\"arr\":[1,2,3]}", ".arr");
     try testRun("{\"bar\":true}", "{\"foo\":{\"bar\":true}}", ".foo");
+
+    try testRun("123", "{\"a\":123}", ".[\"a\"]");
+    try testRun("null", "{\"a\":123}", ".[\"b\"]");
+    try testRun("\"hello\"", "{\"foo\":\"hello\"}", ".[\"foo\"]");
+
+    try testRun("42", "{\"foo bar\":42}", ".[\"foo bar\"]");
+    try testRun("\"value\"", "{\"key with spaces\":\"value\"}", ".[\"key with spaces\"]");
+
+    try testRun("\"world\"", "{\"key\":\"hello\",\"hello\":\"world\"}", ".[.key]");
+    try testRun("3", "[1,2,3,4,5]", ".[1 + 1]");
+    try testRun("5", "[1,2,3,4,5]", ".[2 * 2]");
 }
 
 test "arithmetic operations" {