implement saveable stack for value stack

2 files changed, 446 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.*);
             },
         }
diff --git a/src/jq/saveable_stack.zig b/src/jq/saveable_stack.zig
new file mode 100644
index 0000000..a3085d6
--- /dev/null
+++ b/src/jq/saveable_stack.zig
@@ -0,0 +1,402 @@
+const std = @import("std");
+
+const SEGMENT_INITIAL_CAPACITY: usize = 16;
+
+const StackPosition = struct {
+    segment_index: usize,
+    offset: usize,
+};
+
+fn Segment(comptime T: type) type {
+    return struct {
+        const Self = @This();
+
+        data: std.ArrayList(T),
+        previous_position: StackPosition,
+
+        fn init(allocator: std.mem.Allocator, previous_position: StackPosition) !Self {
+            const data = try std.ArrayList(T).initCapacity(allocator, SEGMENT_INITIAL_CAPACITY);
+            return .{
+                .data = data,
+                .previous_position = previous_position,
+            };
+        }
+
+        fn deinit(self: *Self, allocator: std.mem.Allocator) void {
+            self.data.deinit(allocator);
+        }
+
+        fn len(self: *const Self) usize {
+            return self.data.items.len;
+        }
+
+        fn top(self: *Self) *T {
+            return &self.data.items[self.data.items.len - 1];
+        }
+
+        fn itemAt(self: *Self, index: usize) *T {
+            return &self.data.items[index];
+        }
+
+        fn pop(self: *Self) T {
+            return self.data.pop().?;
+        }
+
+        fn push(self: *Self, allocator: std.mem.Allocator, value: T) !void {
+            try self.data.append(allocator, value);
+        }
+
+        fn clear(self: *Self) void {
+            self.data.clearRetainingCapacity();
+        }
+
+        fn shrink(self: *Self, new_len: usize) void {
+            self.data.shrinkRetainingCapacity(new_len);
+        }
+    };
+}
+
+pub fn SaveableStack(comptime T: type) type {
+    return struct {
+        const Self = @This();
+
+        allocator: std.mem.Allocator,
+        segments: std.ArrayList(Segment(T)),
+        active_segment_index: usize,
+        savepoints: std.ArrayList(StackPosition),
+
+        pub fn init(allocator: std.mem.Allocator) !Self {
+            var stack = Self{
+                .segments = .{},
+                .active_segment_index = 0,
+                .savepoints = .{},
+                .allocator = allocator,
+            };
+
+            const first_segment = try Segment(T).init(allocator, .{ .segment_index = 0, .offset = 0 });
+            try stack.segments.append(allocator, first_segment);
+
+            return stack;
+        }
+
+        pub fn deinit(self: *Self) void {
+            for (self.segments.items) |*seg| {
+                seg.deinit(self.allocator);
+            }
+            self.segments.deinit(self.allocator);
+            self.savepoints.deinit(self.allocator);
+        }
+
+        pub fn push(self: *Self, value: T) !void {
+            try self.activeSegment().push(self.allocator, value);
+        }
+
+        pub fn isEmpty(self: *Self) bool {
+            const seg = self.activeSegment();
+
+            if (seg.len() > 0) {
+                return false;
+            }
+
+            return seg.previous_position.offset == 0;
+        }
+
+        pub fn pop(self: *Self) T {
+            std.debug.assert(!self.isEmpty());
+
+            const seg = self.activeSegment();
+
+            if (seg.len() > 0) {
+                return seg.pop();
+            }
+
+            seg.previous_position.offset -= 1;
+            const result = self.segmentAt(seg.previous_position.segment_index).itemAt(seg.previous_position.offset).*;
+
+            if (seg.previous_position.offset == 0 and seg.previous_position.segment_index > 0) {
+                seg.previous_position = self.segmentAt(seg.previous_position.segment_index).previous_position;
+            }
+
+            return result;
+        }
+
+        pub fn top(self: *Self) *T {
+            std.debug.assert(!self.isEmpty());
+
+            const seg = self.activeSegment();
+
+            if (seg.len() > 0) {
+                return seg.top();
+            }
+
+            const pos = seg.previous_position;
+            return self.segmentAt(pos.segment_index).itemAt(pos.offset - 1);
+        }
+
+        pub fn ensureSize(self: *Self, n: usize) bool {
+            return self.size() >= n;
+        }
+
+        pub fn size(self: *Self) usize {
+            var total: usize = 0;
+            var pos = StackPosition{
+                .segment_index = self.active_segment_index,
+                .offset = self.activeSegment().len(),
+            };
+
+            while (true) {
+                total += pos.offset;
+                if (pos.segment_index == 0) break;
+                pos = self.segmentAt(pos.segment_index).previous_position;
+            }
+
+            return total;
+        }
+
+        pub fn save(self: *Self) !void {
+            const current_pos = StackPosition{
+                .segment_index = self.active_segment_index,
+                .offset = self.activeSegment().len(),
+            };
+
+            try self.savepoints.append(self.allocator, current_pos);
+
+            const effective_pos = self.activeSegment().previous_position;
+
+            const new_seg_idx = self.active_segment_index + 1;
+
+            if (new_seg_idx < self.segments.items.len) {
+                const seg = &self.segments.items[new_seg_idx];
+                seg.clear();
+                seg.previous_position = if (current_pos.offset > 0) current_pos else effective_pos;
+            } else {
+                const pos = if (current_pos.offset > 0) current_pos else effective_pos;
+                const new_segment = try Segment(T).init(self.allocator, pos);
+                try self.segments.append(self.allocator, new_segment);
+            }
+
+            self.active_segment_index = new_seg_idx;
+        }
+
+        pub fn restore(self: *Self) void {
+            std.debug.assert(self.savepoints.items.len > 0);
+
+            const sp = self.savepoints.pop().?;
+
+            while (self.active_segment_index > sp.segment_index) {
+                self.activeSegment().clear();
+                self.active_segment_index -= 1;
+            }
+
+            self.activeSegment().shrink(sp.offset);
+        }
+
+        fn activeSegment(self: *Self) *Segment(T) {
+            std.debug.assert(self.segments.items.len > 0);
+
+            return &self.segments.items[self.active_segment_index];
+        }
+
+        fn segmentAt(self: *Self, index: usize) *Segment(T) {
+            return &self.segments.items[index];
+        }
+    };
+}
+
+test "basic push and pop" {
+    var stack = try SaveableStack(i32).init(std.testing.allocator);
+    defer stack.deinit();
+
+    try stack.push(10);
+    try stack.push(20);
+    try stack.push(30);
+
+    try std.testing.expectEqual(30, stack.top().*);
+
+    try std.testing.expectEqual(30, stack.pop());
+    try std.testing.expectEqual(20, stack.pop());
+    try std.testing.expectEqual(10, stack.pop());
+}
+
+test "save and restore" {
+    var stack = try SaveableStack(i32).init(std.testing.allocator);
+    defer stack.deinit();
+
+    try stack.push(1);
+    try stack.push(2);
+
+    try stack.save();
+
+    try stack.push(3);
+    try stack.push(4);
+
+    try std.testing.expectEqual(4, stack.top().*);
+
+    stack.restore();
+
+    try std.testing.expectEqual(2, stack.top().*);
+    try std.testing.expectEqual(2, stack.pop());
+    try std.testing.expectEqual(1, stack.pop());
+}
+
+test "nested save and restore" {
+    var stack = try SaveableStack(i32).init(std.testing.allocator);
+    defer stack.deinit();
+
+    try stack.push(1);
+    try stack.save();
+
+    try stack.push(2);
+    try stack.save();
+
+    try stack.push(3);
+
+    try std.testing.expectEqual(3, stack.top().*);
+
+    stack.restore();
+    try std.testing.expectEqual(2, stack.top().*);
+
+    stack.restore();
+    try std.testing.expectEqual(1, stack.top().*);
+}
+
+test "pop across segments" {
+    var stack = try SaveableStack(i32).init(std.testing.allocator);
+    defer stack.deinit();
+
+    try stack.push(1);
+    try stack.push(2);
+
+    try stack.save();
+
+    try std.testing.expectEqual(2, stack.pop());
+    try std.testing.expectEqual(1, stack.pop());
+}
+
+test "save restore with segment reuse" {
+    var stack = try SaveableStack(i32).init(std.testing.allocator);
+    defer stack.deinit();
+
+    try stack.push(100);
+    try stack.save();
+    try stack.push(200);
+    stack.restore();
+
+    try stack.save();
+    try stack.push(300);
+
+    try std.testing.expectEqual(300, stack.top().*);
+    try std.testing.expectEqual(300, stack.pop());
+    try std.testing.expectEqual(100, stack.pop());
+}
+
+test "save, pop from previous, push, restore" {
+    var stack = try SaveableStack(i32).init(std.testing.allocator);
+    defer stack.deinit();
+
+    try stack.push(1);
+    try stack.push(2);
+
+    try stack.save();
+
+    try std.testing.expectEqual(2, stack.pop());
+    try std.testing.expectEqual(1, stack.pop());
+
+    try stack.push(3);
+
+    try std.testing.expectEqual(3, stack.top().*);
+
+    stack.restore();
+
+    try std.testing.expectEqual(2, stack.top().*);
+    try std.testing.expectEqual(2, stack.pop());
+    try std.testing.expectEqual(1, stack.pop());
+}
+
+test "nested save with pop across multiple segments" {
+    var stack = try SaveableStack(i32).init(std.testing.allocator);
+    defer stack.deinit();
+
+    try stack.push(1);
+
+    try stack.save();
+    try stack.push(2);
+
+    try stack.save();
+    try stack.push(3);
+
+    try std.testing.expectEqual(3, stack.pop());
+    try std.testing.expectEqual(2, stack.pop());
+    try std.testing.expectEqual(1, stack.pop());
+
+    stack.restore();
+
+    try std.testing.expectEqual(2, stack.top().*);
+    try std.testing.expectEqual(2, stack.pop());
+    try std.testing.expectEqual(1, stack.pop());
+}
+
+test "save on empty segment after pop" {
+    var stack = try SaveableStack(i32).init(std.testing.allocator);
+    defer stack.deinit();
+
+    try stack.push(1);
+    try stack.push(2);
+
+    try stack.save();
+
+    try std.testing.expectEqual(2, stack.pop());
+    try std.testing.expectEqual(1, stack.pop());
+
+    try stack.save();
+
+    try stack.push(3);
+    try std.testing.expectEqual(3, stack.top().*);
+
+    stack.restore();
+
+    try std.testing.expect(stack.isEmpty());
+
+    stack.restore();
+
+    try std.testing.expectEqual(2, stack.top().*);
+    try std.testing.expectEqual(2, stack.pop());
+    try std.testing.expectEqual(1, stack.pop());
+}
+
+test "isEmpty" {
+    var stack = try SaveableStack(i32).init(std.testing.allocator);
+    defer stack.deinit();
+
+    try std.testing.expect(stack.isEmpty());
+
+    try stack.push(1);
+    try std.testing.expect(!stack.isEmpty());
+
+    try stack.push(2);
+    try std.testing.expect(!stack.isEmpty());
+
+    _ = stack.pop();
+    try std.testing.expect(!stack.isEmpty());
+
+    _ = stack.pop();
+    try std.testing.expect(stack.isEmpty());
+}
+
+test "isEmpty with save and restore" {
+    var stack = try SaveableStack(i32).init(std.testing.allocator);
+    defer stack.deinit();
+
+    try std.testing.expect(stack.isEmpty());
+
+    try stack.push(1);
+    try stack.save();
+
+    try std.testing.expect(!stack.isEmpty());
+
+    _ = stack.pop();
+    try std.testing.expect(stack.isEmpty());
+
+    stack.restore();
+    try std.testing.expect(!stack.isEmpty());
+}