use reference counting to manage JSON value lifetime

1 files changed, 83 insertions, 39 deletions

diff --git a/src/jq/execute.zig b/src/jq/execute.zig
index 5749b0b..3096be2 100644
--- a/src/jq/execute.zig
+++ b/src/jq/execute.zig
@@ -34,52 +34,58 @@ const ValueStack = struct {
     }
 
     pub fn pop(self: *Self) jv.Value {
+        // Values beyond the savepoint boundary belong to a previous segment
+        // that may be restored later. We must clone them because restore()
+        // expects those values to be still available.
+        if (self.stack.isBeyondSavepointBoundary()) {
+            return self.stack.pop().clone();
+        }
         return self.stack.pop();
     }
 
     pub fn popInteger(self: *Self) ExecuteError!i64 {
         const value = self.pop();
-        return switch (value.kind()) {
-            .integer => value.integer(),
+        return switch (value) {
+            .integer => |i| i,
             else => error.InvalidType,
         };
     }
 
     pub fn popNumber(self: *Self) ExecuteError!f64 {
         const value = self.pop();
-        return switch (value.kind()) {
-            .integer => @floatFromInt(value.integer()),
-            .float => value.float(),
+        return switch (value) {
+            .integer => |i| @floatFromInt(i),
+            .float => |f| f,
             else => error.InvalidType,
         };
     }
 
     pub fn popString(self: *Self) ExecuteError![]const u8 {
         const value = self.pop();
-        return switch (value.kind()) {
-            .string => value.string(),
+        return switch (value) {
+            .string => |s| s,
             else => error.InvalidType,
         };
     }
 
     pub fn popArray(self: *Self) ExecuteError!jv.Array {
         const value = self.pop();
-        return switch (value.kind()) {
-            .array => value.array(),
+        return switch (value) {
+            .array => |a| a,
             else => error.InvalidType,
         };
     }
 
     pub fn popObject(self: *Self) ExecuteError!jv.Object {
         const value = self.pop();
-        return switch (value.kind()) {
-            .object => value.object(),
+        return switch (value) {
+            .object => |o| o,
             else => error.InvalidType,
         };
     }
 
     pub fn dup(self: *Self) !void {
-        const top = self.stack.top().*;
+        const top = self.stack.top().*.clone();
         try self.push(top);
     }
 
@@ -96,13 +102,35 @@ const ValueStack = struct {
         try self.stack.save();
     }
 
-    pub fn restore(self: *Self) void {
+    pub fn restore(self: *Self, allocator: std.mem.Allocator) void {
+        self.discardAllValuesAboveSavepoint(allocator);
         self.stack.restore();
     }
 
     pub fn ensureSize(self: *Self, n: usize) bool {
         return self.stack.ensureSize(n);
     }
+
+    // Discard all values pushed above the current savepoint.
+    fn discardAllValuesAboveSavepoint(self: *Self, allocator: std.mem.Allocator) void {
+        if (self.stack.savepoints.items.len == 0) return;
+        const sp = self.stack.savepoints.items[self.stack.savepoints.items.len - 1];
+
+        var seg_idx = self.stack.active_segment_index;
+        while (seg_idx > sp.segment_index) : (seg_idx -= 1) {
+            const seg = &self.stack.segments.items[seg_idx];
+            for (seg.data.items) |item| {
+                item.deinit(allocator);
+            }
+        }
+
+        const seg = &self.stack.segments.items[sp.segment_index];
+        if (seg.data.items.len > sp.offset) {
+            for (seg.data.items[sp.offset..]) |item| {
+                item.deinit(allocator);
+            }
+        }
+    }
 };
 
 pub const Runtime = struct {
@@ -122,7 +150,7 @@ pub const Runtime = struct {
         try constants.append(allocator, jv.Value.null);
         try constants.append(allocator, jv.Value.false);
         try constants.append(allocator, jv.Value.true);
-        try constants.append(allocator, jv.Value.initArray(jv.Array.init(allocator)));
+        try constants.append(allocator, jv.Value.initArray(try jv.Array.init(allocator)));
 
         return .{
             .allocator = allocator,
@@ -136,12 +164,15 @@ pub const Runtime = struct {
     }
 
     pub fn deinit(self: *Self) void {
-        for (self.variables.items) |*value| {
+        for (self.variables.items) |value| {
             value.deinit(self.allocator);
         }
         self.variables.deinit(self.allocator);
-        for (self.constants.items) |*value| {
-            value.deinit(self.allocator);
+        for (self.constants.items) |value| {
+            switch (value) {
+                .string => |s| self.allocator.free(s),
+                else => value.deinit(self.allocator),
+            }
         }
         self.constants.deinit(self.allocator);
         self.allocator.free(self.instrs);
@@ -171,7 +202,7 @@ pub const Runtime = struct {
     }
 
     pub fn start(self: *Self, input: jv.Value) !void {
-        try self.values.push(input);
+        try self.values.push(input.clone());
     }
 
     pub fn next(self: *Self) !?jv.Value {
@@ -217,7 +248,7 @@ pub const Runtime = struct {
                 .pop => {
                     std.debug.assert(self.values.ensureSize(1));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                 },
                 .subexp_begin => try self.values.dup(),
                 .subexp_end => try self.values.swap(),
@@ -226,7 +257,9 @@ pub const Runtime = struct {
 
                     const base = self.values.pop();
                     const key = self.values.pop();
-                    const result = try jv.ops.index(base, key);
+                    const result = (try jv.ops.index(base, key)).clone();
+                    base.deinit(self.allocator);
+                    key.deinit(self.allocator);
                     try self.values.push(result);
                 },
                 .index_opt => {
@@ -234,13 +267,16 @@ pub const Runtime = struct {
 
                     const base = self.values.pop();
                     const key = self.values.pop();
-                    const result = jv.ops.index(base, key) catch jv.Value.null;
+                    const idx_result: jv.Value = jv.ops.index(base, key) catch .null;
+                    const result = idx_result.clone();
+                    base.deinit(self.allocator);
+                    key.deinit(self.allocator);
                     try self.values.push(result);
                 },
                 .add => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = try self.values.popInteger();
                     const rhs = try self.values.popInteger();
                     const result = lhs + rhs;
@@ -249,7 +285,7 @@ pub const Runtime = struct {
                 .sub => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = try self.values.popInteger();
                     const rhs = try self.values.popInteger();
                     const result = lhs - rhs;
@@ -258,7 +294,7 @@ pub const Runtime = struct {
                 .mul => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = try self.values.popInteger();
                     const rhs = try self.values.popInteger();
                     const result = lhs * rhs;
@@ -267,7 +303,7 @@ pub const Runtime = struct {
                 .div => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = try self.values.popInteger();
                     const rhs = try self.values.popInteger();
                     const result = @divTrunc(lhs, rhs);
@@ -276,7 +312,7 @@ pub const Runtime = struct {
                 .mod => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = try self.values.popInteger();
                     const rhs = try self.values.popInteger();
                     const result = @mod(lhs, rhs);
@@ -285,7 +321,7 @@ pub const Runtime = struct {
                 .eq => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = self.values.pop();
                     const rhs = self.values.pop();
                     const result = try jv.ops.compare(lhs, rhs, .eq);
@@ -294,7 +330,7 @@ pub const Runtime = struct {
                 .ne => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = self.values.pop();
                     const rhs = self.values.pop();
                     const result = try jv.ops.compare(lhs, rhs, .ne);
@@ -303,7 +339,7 @@ pub const Runtime = struct {
                 .lt => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = self.values.pop();
                     const rhs = self.values.pop();
                     const result = try jv.ops.compare(lhs, rhs, .lt);
@@ -312,7 +348,7 @@ pub const Runtime = struct {
                 .gt => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = self.values.pop();
                     const rhs = self.values.pop();
                     const result = try jv.ops.compare(lhs, rhs, .gt);
@@ -321,7 +357,7 @@ pub const Runtime = struct {
                 .le => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = self.values.pop();
                     const rhs = self.values.pop();
                     const result = try jv.ops.compare(lhs, rhs, .le);
@@ -330,7 +366,7 @@ pub const Runtime = struct {
                 .ge => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = self.values.pop();
                     const rhs = self.values.pop();
                     const result = try jv.ops.compare(lhs, rhs, .ge);
@@ -339,19 +375,25 @@ pub const Runtime = struct {
                 .alt => {
                     std.debug.assert(self.values.ensureSize(3));
 
-                    _ = self.values.pop();
+                    self.values.pop().deinit(self.allocator);
                     const lhs = self.values.pop();
                     const rhs = self.values.pop();
-                    try self.values.push(if (jv.ops.isFalsy(lhs)) rhs else lhs);
+                    if (jv.ops.isFalsy(lhs)) {
+                        lhs.deinit(self.allocator);
+                        try self.values.push(rhs);
+                    } else {
+                        rhs.deinit(self.allocator);
+                        try self.values.push(lhs);
+                    }
                 },
                 .@"const" => |idx| {
                     std.debug.assert(self.values.ensureSize(1));
 
-                    _ = self.values.pop();
-                    try self.values.push(self.constants.items[@intFromEnum(idx)]);
+                    self.values.pop().deinit(self.allocator);
+                    try self.values.push(self.constants.items[@intFromEnum(idx)].clone());
                 },
                 .load => |idx| {
-                    try self.values.push(self.variables.items[@intFromEnum(idx)]);
+                    try self.values.push(self.variables.items[@intFromEnum(idx)].clone());
                 },
                 .store => |idx| {
                     std.debug.assert(self.values.ensureSize(1));
@@ -360,12 +402,14 @@ pub const Runtime = struct {
                     while (self.variables.items.len <= @intFromEnum(idx)) {
                         try self.variables.append(self.allocator, jv.Value.null);
                     }
+                    self.variables.items[@intFromEnum(idx)].deinit(self.allocator);
                     self.variables.items[@intFromEnum(idx)] = self.values.pop();
                 },
                 .append => |idx| {
                     std.debug.assert(self.values.ensureSize(1));
 
-                    try self.variables.items[@intFromEnum(idx)].arrayAppend(self.values.pop());
+                    const var_ptr = &self.variables.items[@intFromEnum(idx)];
+                    try var_ptr.arrayAppend(self.allocator, self.values.pop());
                 },
             }
         }
@@ -381,7 +425,7 @@ pub const Runtime = struct {
     fn restore_stack(self: *Self) bool {
         if (self.forks.pop()) |target_pc| {
             self.pc = target_pc;
-            self.values.restore();
+            self.values.restore(self.allocator);
             return true;
         }
         return false;