fix(persistent_merkle_tree): address batchGetRoot review comments

Author: lodekeeper-z

src/persistent_merkle_tree/Node.zig

@@ -419,6 +419,8 @@ pub const Id = enum(u32) {
     /// Collects dirty branches via BFS, then hashes them bottom-up in batches
     /// using SIMD-accelerated `hashBatch`. Produces identical results to `getRoot`
     /// but is ~2x faster on large trees with many dirty nodes.
+    ///
+    /// `depth` must be the full height of `node_id`'s tree.
     pub fn batchGetRoot(node_id: Id, pool: *Pool, depth: Depth, allocator: Allocator) Error!*const [32]u8 {
         const states = pool.nodes.items(.state);
         const hashes = pool.nodes.items(.hash);
@@ -427,21 +429,15 @@ pub const Id = enum(u32) {
         if (!states[@intFromEnum(node_id)].isBranchLazy()) {
             return &hashes[@intFromEnum(node_id)];
         }
-
-        // Fast path: depth 0 means the node is a leaf, just return its hash.
-        if (depth == 0) {
-            return &hashes[@intFromEnum(node_id)];
-        }
+        if (depth == 0) return Error.InvalidNode;
 
         const lefts = pool.nodes.items(.left);
         const rights = pool.nodes.items(.right);
 
         // Collect dirty branch nodes per level via BFS.
         // levels[i] holds dirty node IDs at depth i (0 = root level).
         var levels: [max_depth]std.ArrayListUnmanaged(Id) = undefined;
-        for (0..depth) |i| {
-            levels[i] = .empty;
-        }
+        @memset(levels[0..depth], .empty);
         defer for (0..depth) |i| {
             levels[i].deinit(allocator);
         };
@@ -466,16 +462,15 @@ pub const Id = enum(u32) {
         // Bottom-up batch hash: process from deepest level to root.
         var pairs = std.ArrayListUnmanaged([32]u8).empty;
         defer pairs.deinit(allocator);
+
         var outs = std.ArrayListUnmanaged([32]u8).empty;
         defer outs.deinit(allocator);
 
-        var level_i: usize = depth - 1;
-        while (true) : (level_i -= 1) {
+        var level_i: usize = depth;
+        while (level_i > 0) {
+            level_i -= 1;
             const dirty_nodes = levels[level_i].items;
-            if (dirty_nodes.len == 0) {
-                if (level_i == 0) break;
-                continue;
-            }
+            if (dirty_nodes.len == 0) continue;
 
             // Build input pairs: [left_hash, right_hash] for each dirty node.
             pairs.clearRetainingCapacity();
@@ -487,17 +482,14 @@ pub const Id = enum(u32) {
 
             // Batch hash
             outs.clearRetainingCapacity();
-            try outs.ensureTotalCapacity(allocator, dirty_nodes.len);
-            outs.items.len = dirty_nodes.len;
+            try outs.resize(allocator, dirty_nodes.len);
             hashBatch(outs.items, pairs.items) catch unreachable;
 
             // Write results back
             for (dirty_nodes, 0..) |id, i| {
                 hashes[@intFromEnum(id)] = outs.items[i];
                 states[@intFromEnum(id)].setBranchComputed();
             }
-
-            if (level_i == 0) break;
         }
 
         return &hashes[@intFromEnum(node_id)];

src/persistent_merkle_tree/node_test.zig

@@ -640,26 +640,38 @@ test "batchGetRoot matches getRoot for dirty branches at various depths" {
         const depth: Depth = @intCast(depth_usize);
         const max_length = @as(usize, 1) << @intCast(depth);
 
-        // Build a full tree, then modify some leaves to create dirty branches.
         var leaves = try allocator.alloc(Node.Id, max_length);
         defer allocator.free(leaves);
         for (0..max_length) |i| {
             leaves[i] = try pool.createLeafFromUint(@intCast(i + 1));
+            try pool.ref(leaves[i]);
         }
+        defer for (leaves) |leaf| pool.unref(leaf);
 
-        const root = try Node.fillWithContents(p, leaves, depth);
-        defer pool.unref(root);
+        // Build two identical dirty trees. `getRoot` mutates in-place, so the
+        // batch path must run on a separate tree to exercise dirty branches.
+        const leaves_for_get_root = try allocator.dupe(Node.Id, leaves);
+        defer allocator.free(leaves_for_get_root);
+
+        const leaves_for_batch = try allocator.dupe(Node.Id, leaves);
+        defer allocator.free(leaves_for_batch);
+
+        const root_get = try Node.fillWithContents(p, leaves_for_get_root, depth);
+        defer pool.unref(root_get);
+
+        const root_batch = try Node.fillWithContents(p, leaves_for_batch, depth);
+        defer pool.unref(root_batch);
+
+        const new_leaf_get = try pool.createLeafFromUint(0xBEEF);
+        const dirty_get = try root_get.setNodeAtDepth(p, depth, 0, new_leaf_get);
+        defer pool.unref(dirty_get);
 
-        // Modify a few leaves to make branches dirty.
-        const new_leaf = try pool.createLeafFromUint(0xBEEF);
-        const dirty_root = try root.setNodeAtDepth(p, depth, 0, new_leaf);
-        defer pool.unref(dirty_root);
+        const new_leaf_batch = try pool.createLeafFromUint(0xBEEF);
+        const dirty_batch = try root_batch.setNodeAtDepth(p, depth, 0, new_leaf_batch);
+        defer pool.unref(dirty_batch);
 
-        // Clone the tree for getRoot (since getRoot mutates state in-place).
-        // Both should produce the same hash.
-        const expected = dirty_root.getRoot(p);
-        // getRoot already computed the hashes, but batchGetRoot should return the same result.
-        const actual = try dirty_root.batchGetRoot(p, depth, allocator);
+        const expected = dirty_get.getRoot(p);
+        const actual = try dirty_batch.batchGetRoot(p, depth, allocator);
         try std.testing.expectEqualSlices(u8, expected, actual);
     }
 }
@@ -673,31 +685,47 @@ test "batchGetRoot matches getRoot with multiple dirty leaves" {
     const depth: Depth = 4;
     const max_length = @as(usize, 1) << depth;
 
-    // Build a full tree with all leaves set.
     var leaves = try allocator.alloc(Node.Id, max_length);
     defer allocator.free(leaves);
     for (0..max_length) |i| {
         leaves[i] = try pool.createLeafFromUint(@intCast(i + 1));
+        try pool.ref(leaves[i]);
     }
+    defer for (leaves) |leaf| pool.unref(leaf);
+
+    // Build two identical dirty trees. `getRoot` mutates in-place, so the
+    // batch path must run on a separate tree to exercise dirty branches.
+    const leaves_for_get_root = try allocator.dupe(Node.Id, leaves);
+    defer allocator.free(leaves_for_get_root);
+
+    const leaves_for_batch = try allocator.dupe(Node.Id, leaves);
+    defer allocator.free(leaves_for_batch);
+
+    const root_get = try Node.fillWithContents(p, leaves_for_get_root, depth);
+    defer pool.unref(root_get);
+    var modified_get = root_get;
 
-    const root = try Node.fillWithContents(p, leaves, depth);
+    const root_batch = try Node.fillWithContents(p, leaves_for_batch, depth);
+    defer pool.unref(root_batch);
+    var modified_batch = root_batch;
 
-    // Modify multiple scattered leaves.
-    var modified = root;
     const modify_indices = [_]usize{ 0, 3, 7, 12, 15 };
     for (modify_indices) |idx| {
-        const new_leaf = try pool.createLeafFromUint(@intCast(0xF000 + idx));
-        const old = modified;
-        modified = try modified.setNodeAtDepth(p, depth, idx, new_leaf);
-        if (old != root) pool.unref(old);
+        const new_leaf_get = try pool.createLeafFromUint(@intCast(0xF000 + idx));
+        const old_get = modified_get;
+        modified_get = try modified_get.setNodeAtDepth(p, depth, idx, new_leaf_get);
+        if (old_get != root_get) pool.unref(old_get);
+
+        const new_leaf_batch = try pool.createLeafFromUint(@intCast(0xF000 + idx));
+        const old_batch = modified_batch;
+        modified_batch = try modified_batch.setNodeAtDepth(p, depth, idx, new_leaf_batch);
+        if (old_batch != root_batch) pool.unref(old_batch);
     }
-    defer pool.unref(modified);
-    pool.unref(root);
+    defer pool.unref(modified_get);
+    defer pool.unref(modified_batch);
 
-    // Use getRoot on a copy-by-reference to get expected hash, then verify batchGetRoot matches.
-    // Since getRoot computes in-place, calling it first then batchGetRoot should still agree.
-    const expected = modified.getRoot(p);
-    const actual = try modified.batchGetRoot(p, depth, allocator);
+    const expected = modified_get.getRoot(p);
+    const actual = try modified_batch.batchGetRoot(p, depth, allocator);
     try std.testing.expectEqualSlices(u8, expected, actual);
 }