1 files changed, 44 insertions, 22 deletions

diff --git a/src/jq/execute.zig b/src/jq/execute.zig
index 8982458..5108a64 100644
--- a/src/jq/execute.zig
+++ b/src/jq/execute.zig
@@ -8,34 +8,34 @@ pub const ExecuteError = error{
     InternalError,
 };
 
+const SaveableStack = @import("./saveable_stack.zig").SaveableStack;
+
 const ValueStack = struct {
     const Self = @This();
-    const Stack = std.ArrayList(jv.Value);
+    const Stack = SaveableStack(jv.Value);
 
     stack: Stack,
-    allocator: std.mem.Allocator,
 
     pub fn init(allocator: std.mem.Allocator) !Self {
         return .{
-            .stack = try Stack.initCapacity(allocator, 16),
-            .allocator = allocator,
+            .stack = try Stack.init(allocator),
         };
     }
 
     pub fn deinit(self: *Self) void {
-        self.stack.deinit(self.allocator);
+        self.stack.deinit();
     }
 
     pub fn push(self: *Self, value: jv.Value) !void {
-        try self.stack.append(self.allocator, value);
+        try self.stack.push(value);
     }
 
-    pub fn pop(self: *Self) ExecuteError!jv.Value {
-        return self.stack.pop() orelse return error.InternalError;
+    pub fn pop(self: *Self) jv.Value {
+        return self.stack.pop();
     }
 
     pub fn popInteger(self: *Self) ExecuteError!i64 {
-        const value = try self.pop();
+        const value = self.pop();
         return switch (value) {
             .integer => |i| i,
             else => error.InvalidType,
@@ -43,7 +43,7 @@ const ValueStack = struct {
     }
 
     pub fn popNumber(self: *Self) ExecuteError!f64 {
-        const value = try self.pop();
+        const value = self.pop();
         return switch (value) {
             .integer => |i| @floatFromInt(i),
             .float => |f| f,
@@ -52,7 +52,7 @@ const ValueStack = struct {
     }
 
     pub fn popString(self: *Self) ExecuteError![]const u8 {
-        const value = try self.pop();
+        const value = self.pop();
         return switch (value) {
             .string => |s| s,
             else => error.InvalidType,
@@ -60,7 +60,7 @@ const ValueStack = struct {
     }
 
     pub fn popArray(self: *Self) ExecuteError!jv.Array {
-        const value = try self.pop();
+        const value = self.pop();
         return switch (value) {
             .array => |a| a,
             else => error.InvalidType,
@@ -68,7 +68,7 @@ const ValueStack = struct {
     }
 
     pub fn popObject(self: *Self) ExecuteError!jv.Object {
-        const value = try self.pop();
+        const value = self.pop();
         return switch (value) {
             .object => |o| o,
             else => error.InvalidType,
@@ -76,15 +76,29 @@ const ValueStack = struct {
     }
 
     pub fn dup(self: *Self) !void {
-        const top = self.stack.items[self.stack.items.len - 1];
+        const top = self.stack.top().*;
         try self.push(top);
     }
 
-    pub fn swap(self: *Self) void {
-        const len = self.stack.items.len;
-        const tmp = self.stack.items[len - 1];
-        self.stack.items[len - 1] = self.stack.items[len - 2];
-        self.stack.items[len - 2] = tmp;
+    pub fn swap(self: *Self) !void {
+        std.debug.assert(self.ensureSize(2));
+
+        const a = self.pop();
+        const b = self.pop();
+        try self.push(a);
+        try self.push(b);
+    }
+
+    pub fn save(self: *Self) !void {
+        try self.stack.save();
+    }
+
+    pub fn restore(self: *Self) void {
+        self.stack.restore();
+    }
+
+    pub fn ensureSize(self: *Self, n: usize) bool {
+        return self.stack.ensureSize(n);
     }
 };
 
@@ -101,27 +115,35 @@ pub fn execute(allocator: std.mem.Allocator, instrs: []const Instr, input: jv.Va
         switch (cur) {
             .nop => {},
             .subexp_begin => try value_stack.dup(),
-            .subexp_end => value_stack.swap(),
+            .subexp_end => try value_stack.swap(),
             .array_index => {
+                std.debug.assert(value_stack.ensureSize(2));
+
                 const array = try value_stack.popArray();
                 const index: usize = @intCast(try value_stack.popInteger());
                 const result = if (index < array.items.len) array.items[index] else .null;
                 try value_stack.push(result);
             },
             .add => {
-                _ = try value_stack.pop();
+                std.debug.assert(value_stack.ensureSize(3));
+
+                _ = value_stack.pop();
                 const lhs = try value_stack.popInteger();
                 const rhs = try value_stack.popInteger();
                 const result = lhs + rhs;
                 try value_stack.push(.{ .integer = result });
             },
             .object_key => |key| {
+                std.debug.assert(value_stack.ensureSize(1));
+
                 const obj = try value_stack.popObject();
                 const result = obj.get(key) orelse .null;
                 try value_stack.push(result);
             },
             .literal => |value| {
-                _ = try value_stack.pop();
+                std.debug.assert(value_stack.ensureSize(1));
+
+                _ = value_stack.pop();
                 try value_stack.push(value.*);
             },
         }