From 6467883cf432a115fa33119ccc60361e2a3b76e8 Mon Sep 17 00:00:00 2001
From: Travis Staloch <1562827+travisstaloch@users.noreply.github.com>
Date: Sun, 8 Oct 2023 23:19:23 -0700
Subject: [PATCH] dom: fix Array.at() + add tests

this bug was discovered in simdjzon-rpc. previously the new tests
failed.
---
 src/dom.zig   | 23 ++++++++++++-----------
 src/tests.zig | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 11 deletions(-)

diff --git a/src/dom.zig b/src/dom.zig
index 98e94cd..c49bc22 100644
--- a/src/dom.zig
+++ b/src/dom.zig
@@ -1504,22 +1504,22 @@ pub const Parser = struct {
         return parser.stage2();
     }
 
-    pub inline fn element(parser: Parser) Element {
-        return .{
-            .tape = .{ .doc = &parser.doc, .idx = 1 },
-        };
+    pub fn element(parser: *const Parser) Element {
+        return .{ .tape = .{ .doc = &parser.doc, .idx = 1 } };
     }
 };
 
 pub const Array = struct {
     tape: TapeRef,
-    pub fn at(a: Array, idx: usize) ?Element {
-        var it = TapeRefIterator.init(a);
-        const target_idx = idx + it.tape.idx + idx;
-        while (true) {
-            if (it.tape.idx == target_idx)
-                return Element{ .tape = .{ .doc = it.tape.doc, .idx = it.tape.idx } };
-            _ = it.next() orelse break;
+    pub fn at(_a: Array, idx: usize) ?Element {
+        var i: usize = 0;
+        const end = _a.tape.after_element() - 1;
+        var a = Array{ .tape = .{ .idx = _a.tape.idx + 1, .doc = _a.tape.doc } };
+        while (a.tape.idx < end) {
+            if (i == idx)
+                return .{ .tape = a.tape };
+            a.tape.idx = a.tape.after_element();
+            i += 1;
         }
         return null;
     }
@@ -1688,6 +1688,7 @@ const TapeRef = struct {
 
     pub fn get_string_length(tr: TapeRef) u32 {
         const string_buf_index = tr.value();
+        // std.debug.print("get_string_length string_buf_index={}\n", .{string_buf_index});
         return mem.readIntLittle(u32, (tr.doc.string_buf.ptr + string_buf_index)[0..@sizeOf(u32)]);
     }
 
diff --git a/src/tests.zig b/src/tests.zig
index 36b6f63..8004b14 100644
--- a/src/tests.zig
+++ b/src/tests.zig
@@ -192,6 +192,52 @@ test "search tape" {
     try testing.expectEqual(@as(i64, 116), try array1.get_int64());
 }
 
+test "array of objects" {
+    var parser = try dom.Parser.initFixedBuffer(allr,
+        \\[{"a": 1}, {"b": 2}]
+    , .{});
+    defer parser.deinit();
+    try parser.parse();
+    const ele = parser.element();
+    try testing.expect(ele.is(.ARRAY));
+    const array = try ele.get_array();
+    const array0 = array.at(0) orelse unreachable;
+    try testing.expect(array0.is(.OBJECT));
+    const array0_afield = array0.at_key("a") orelse unreachable;
+    try testing.expect(array0_afield.is(.INT64));
+    try testing.expectEqual(@as(i64, 1), try array0_afield.get_int64());
+}
+
+test "array of objects 2" {
+    var parser = try dom.Parser.initFixedBuffer(allr,
+        \\[{"foo": null, "params": [1,2,4],  "id": 1},
+        \\ {"foo": null, "params": [1,2,10], "id": 2}]
+    , .{});
+    defer parser.deinit();
+    try parser.parse();
+    const ele = parser.element();
+    try testing.expect(ele.is(.ARRAY));
+    const array = try ele.get_array();
+
+    const array0 = array.at(0) orelse unreachable;
+    try testing.expect(array0.is(.OBJECT));
+    const params0 = array0.at_key("params") orelse unreachable;
+    try testing.expect(params0.is(.ARRAY));
+    const params0_3 = params0.at(2) orelse unreachable;
+    try testing.expectEqual(@as(i64, 4), try params0_3.get_int64());
+    const id0 = array0.at_key("id") orelse unreachable;
+    try testing.expectEqual(@as(i64, 1), try id0.get_int64());
+
+    const array1 = array.at(1) orelse unreachable;
+    try testing.expect(array1.is(.OBJECT));
+    const params1 = array1.at_key("params") orelse unreachable;
+    try testing.expect(params1.is(.ARRAY));
+    const params1_3 = params1.at(2) orelse unreachable;
+    try testing.expectEqual(@as(i64, 10), try params1_3.get_int64());
+    const id1 = array1.at_key("id") orelse unreachable;
+    try testing.expectEqual(@as(i64, 2), try id1.get_int64());
+}
+
 test "json pointer" {
     const input =
         \\{"a": {"b": [1,2,3], "c": 3.1415, "d": true, "e": "e-string", "f": null}}