Skip to content
This repository was archived by the owner on Nov 12, 2025. It is now read-only.

feat: bind to next node if possible in setNodes apis #70

Closed
wants to merge 6 commits into from
Closed

feat: bind to next node if possible in setNodes apis #70

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
0798a98
fix: avoid orphaned nodes during setNodes calls
twoeths Sep 24, 2025
f16b30e
fix: bring back Depth helpers test
twoeths Sep 24, 2025
9e70ab6
feat: bind to next node if possible
twoeths Sep 23, 2025
ab3614b
chore: add negative unit tests
twoeths Sep 24, 2025
8ab427f
feat: apply same algorithm to setNodes()
twoeths Sep 25, 2025
31da1de
fix: dedup gindex to path conversion
twoeths Sep 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
116 changes: 104 additions & 12 deletions src/persistent_merkle_tree/Node.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,18 @@ pub const Pool = struct {
return self.createUnsafe(self.nodes.items(.state));
}

/// Returns the number of nodes currently in use (not free)
pub fn getNodesInUse(self: *Pool) usize {
var count: usize = 0;
const states = self.nodes.items(.state);
for (states) |state| {
if (!state.isFree()) {
count += 1;
}
}
return count;
}

pub fn createLeaf(self: *Pool, hash: *const [32]u8, should_ref: bool) Allocator.Error!Id {
const node_id = try self.create();

Expand Down Expand Up @@ -605,8 +617,12 @@ pub const Id = enum(u32) {
const path_len = base_gindex.pathLen();

var path_parents_buf: [max_depth]Id = undefined;
// at each level, there is at most 1 unfinalized parent per traversal
var unfinalized_parents_buf: [max_depth]?Id = undefined;
var path_lefts_buf: [max_depth]Id = undefined;
var path_rights_buf: [max_depth]Id = undefined;
// right_move means it's part of the new tree, it happens when we traverse right
var right_move: [max_depth]bool = undefined;

const path_parents = path_parents_buf[0..path_len];
const path_lefts = path_lefts_buf[0..path_len];
Expand All @@ -626,25 +642,29 @@ pub const Id = enum(u32) {

// For each index specified, maintain/update path_lefts and path_rights from root (depth 0) all the way to path_len
// but only allocate and update path_parents from the next shared depth to path_len
for (0..indices.len) |i| {
var i: usize = 0;
while (i < indices.len) : (i += 1) {
// Calculate the gindex bits for the current index
const index = indices[i];
const gindex: Gindex = @enumFromInt(@as(Gindex.Uint, @intCast(@intFromEnum(base_gindex) | index)));
var path = gindex.toPath();

// warning: this next_index variable is not correct for the last index
const next_index = if (i + 1 < indices.len) indices[i + 1] else index;
const consume_next: bool = if (i + 1 < indices.len and index + 1 == next_index and path.leftN(path_len - 1)) true else false;
// Calculate the depth offset to navigate from current index to the next
const next_d_offset = if (i == indices.len - 1)
const next_d_offset = if (i + 1 == indices.len or (i + 2 == indices.len and consume_next))
// 0 because there is no next index, it also means node_id is now the new root
0
else
path_len - @as(Depth, @intCast(@bitSizeOf(usize) - @clz(index ^ indices[i + 1])));
path_len - @as(Depth, @intCast(@bitSizeOf(usize) - @clz(index ^ if (consume_next) indices[i + 2] else next_index)));

if (try pool.alloc(path_parents[next_d_offset..path_len])) {
states = pool.nodes.items(.state);
lefts = pool.nodes.items(.left);
rights = pool.nodes.items(.right);
}

var path = gindex.toPath();

// Navigate down (to the depth offset), attaching any new updates
// d_offset is the shared depth between the previous and current index so we can reuse path_lefts and path_rights up that point
// but update them to the path_parents to rebind starting from next_d_offset if needed
Expand All @@ -653,15 +673,26 @@ pub const Id = enum(u32) {
for (next_d_offset..d_offset) |bit_i| {
if (path.left()) {
path_lefts[bit_i] = path_parents[bit_i + 1];
right_move[bit_i] = false;
// move left, unfinalized
unfinalized_parents_buf[bit_i] = path_parents[bit_i];
} else {
path_rights[bit_i] = path_parents[bit_i + 1];
right_move[bit_i] = true;
}
path.next();
}
} else {
path.nextN(d_offset);
}

// right move at d_offset, make all unfinalized parents at lower levels as finalized
if (path.right()) {
for (d_offset + 1..path_len) |bit_i| {
unfinalized_parents_buf[bit_i] = null;
}
}

// Navigate down (from the depth offset) to the current index, populating parents
for (d_offset..path_len - 1) |bit_i| {
if (node_id.noChild(states[@intFromEnum(node_id)])) {
Expand All @@ -672,23 +703,31 @@ pub const Id = enum(u32) {
path_lefts[bit_i] = path_parents[bit_i + 1];
path_rights[bit_i] = rights[@intFromEnum(node_id)];
node_id = lefts[@intFromEnum(node_id)];
right_move[bit_i] = false;
unfinalized_parents_buf[bit_i] = path_parents[bit_i];
} else {
path_lefts[bit_i] = lefts[@intFromEnum(node_id)];
path_rights[bit_i] = path_parents[bit_i + 1];
node_id = rights[@intFromEnum(node_id)];
right_move[bit_i] = true;
}
path.next();
}

// final layer
if (node_id.noChild(states[@intFromEnum(node_id)])) {
return Error.InvalidNode;
}

if (path.left()) {
path_lefts[path_len - 1] = nodes[i];
path_rights[path_len - 1] = rights[@intFromEnum(node_id)];
path_rights[path_len - 1] = if (consume_next) nodes[i + 1] else rights[@intFromEnum(node_id)];
right_move[path_len - 1] = false;
unfinalized_parents_buf[path_len - 1] = path_parents[path_len - 1];
} else {
path_lefts[path_len - 1] = lefts[@intFromEnum(node_id)];
path_rights[path_len - 1] = nodes[i];
right_move[path_len - 1] = true;
}

// Rebind upwards depth diff times
Expand All @@ -697,7 +736,22 @@ pub const Id = enum(u32) {
path_lefts[next_d_offset..path_len],
path_rights[next_d_offset..path_len],
);

// unref prev parents if it's not part of the new tree
// can only unref after the rebind
for (next_d_offset..path_len) |bit_i| {
if (right_move[bit_i] and unfinalized_parents_buf[bit_i] != null) {
pool.unref(unfinalized_parents_buf[bit_i].?);
unfinalized_parents_buf[bit_i] = null;
}
}
node_id = path_parents[next_d_offset];

if (consume_next) {
// Move pointer one extra forward since node has consumed two nodes
i += 1;
}

d_offset = next_d_offset;
}

Expand All @@ -717,8 +771,12 @@ pub const Id = enum(u32) {
const path_len = base_gindex.pathLen();

var path_parents_buf: [max_depth]Id = undefined;
// at each level, there is at most 1 unfinalized parent per traversal
var unfinalized_parents_buf: [max_depth]?Id = undefined;
var path_lefts_buf: [max_depth]Id = undefined;
var path_rights_buf: [max_depth]Id = undefined;
// right_move means it's part of the new tree, it happens when we traverse right
var right_move: [max_depth]bool = undefined;

// TODO how to handle failing to resize here, especially after several allocs
errdefer pool.free(&path_parents_buf);
Expand All @@ -734,25 +792,28 @@ pub const Id = enum(u32) {

// For each index specified, maintain/update path_lefts and path_rights from root (depth 0) all the way to path_len
// but only allocate and update path_parents from the next shared depth to path_len
for (0..gindices.len) |i| {
var i: usize = 0;
while (i < gindices.len) : (i += 1) {
// Calculate the gindex bits for the current index
const gindex = gindices[i];
var path = gindex.toPath();

// warning: this next_index variable is not correct for the last index
const next_gindex = if (i + 1 < gindices.len) gindices[i + 1] else gindex;
const consume_next: bool = if (i + 1 < gindices.len and @intFromEnum(gindex) + 1 == @intFromEnum(next_gindex) and path.leftN(path_len - 1)) true else false;
// Calculate the depth offset to navigate from current index to the next
const next_d_offset = if (i == gindices.len - 1)
const next_d_offset = if (i + 1 == gindices.len or (i + 2 == gindices.len and consume_next))
// 0 because there is no next gindex, it also means node_id is now the new root
0
else
path_len - @as(Depth, @intCast(@bitSizeOf(usize) - @clz(@intFromEnum(gindex) ^ @intFromEnum(gindices[i + 1]))));
path_len - @as(Depth, @intCast(@bitSizeOf(usize) - @clz(@intFromEnum(gindex) ^ if (consume_next) @intFromEnum(gindices[i + 2]) else @intFromEnum(next_gindex))));

if (try pool.alloc(path_parents_buf[next_d_offset..path_len])) {
states = pool.nodes.items(.state);
lefts = pool.nodes.items(.left);
rights = pool.nodes.items(.right);
}

var path = gindex.toPath();

// Navigate down (to the depth offset), attaching any new updates
// d_offset is the shared depth between the previous and current index so we can reuse path_lefts and path_rights up that point
// but update them to the path_parents to rebind starting from next_d_offset if needed
Expand All @@ -761,15 +822,26 @@ pub const Id = enum(u32) {
for (next_d_offset..d_offset) |bit_i| {
if (path.left()) {
path_lefts_buf[bit_i] = path_parents_buf[bit_i + 1];
right_move[bit_i] = false;
// move left, unfinalized
unfinalized_parents_buf[bit_i] = path_parents_buf[bit_i];
} else {
path_rights_buf[bit_i] = path_parents_buf[bit_i + 1];
right_move[bit_i] = true;
}
path.next();
}
} else {
path.nextN(d_offset);
}

// right move at d_offset, make all unfinalized parents at lower levels as finalized
if (path.right()) {
for (d_offset + 1..path_len) |bit_i| {
unfinalized_parents_buf[bit_i] = null;
}
}

// Navigate down (from the depth offset) to the current index, populating parents
for (d_offset..path_len - 1) |bit_i| {
if (node_id.noChild(states[@intFromEnum(node_id)])) {
Expand All @@ -780,10 +852,13 @@ pub const Id = enum(u32) {
path_lefts_buf[bit_i] = path_parents_buf[bit_i + 1];
path_rights_buf[bit_i] = rights[@intFromEnum(node_id)];
node_id = lefts[@intFromEnum(node_id)];
right_move[bit_i] = false;
unfinalized_parents_buf[bit_i] = path_parents_buf[bit_i];
} else {
path_lefts_buf[bit_i] = lefts[@intFromEnum(node_id)];
path_rights_buf[bit_i] = path_parents_buf[bit_i + 1];
node_id = rights[@intFromEnum(node_id)];
right_move[bit_i] = true;
}
path.next();
}
Expand All @@ -793,10 +868,13 @@ pub const Id = enum(u32) {
}
if (path.left()) {
path_lefts_buf[path_len - 1] = nodes[i];
path_rights_buf[path_len - 1] = rights[@intFromEnum(node_id)];
path_rights_buf[path_len - 1] = if (consume_next) nodes[i + 1] else rights[@intFromEnum(node_id)];
right_move[path_len - 1] = false;
unfinalized_parents_buf[path_len - 1] = path_parents_buf[path_len - 1];
} else {
path_lefts_buf[path_len - 1] = lefts[@intFromEnum(node_id)];
path_rights_buf[path_len - 1] = nodes[i];
right_move[path_len - 1] = true;
}

// Rebind upwards depth diff times
Expand All @@ -805,8 +883,22 @@ pub const Id = enum(u32) {
path_lefts_buf[next_d_offset..path_len],
path_rights_buf[next_d_offset..path_len],
);
// unref prev parents if it's not part of the new tree
// can only unref after the rebind
for (next_d_offset..path_len) |bit_i| {
if (right_move[bit_i] and unfinalized_parents_buf[bit_i] != null) {
pool.unref(unfinalized_parents_buf[bit_i].?);
unfinalized_parents_buf[bit_i] = null;
}
}

node_id = path_parents_buf[next_d_offset];
d_offset = next_d_offset;

if (consume_next) {
// Move pointer one extra forward since node has consumed two nodes
i += 1;
}
}

return node_id;
Expand Down
8 changes: 8 additions & 0 deletions src/persistent_merkle_tree/gindex.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,10 +82,18 @@ pub const Gindex = enum(GindexUint) {
return @intFromEnum(path) & 1 == 0;
}

pub inline fn leftN(path: Path, n: Depth) bool {
return (@intFromEnum(path) >> n) & 1 == 0;
}

pub inline fn right(path: Path) bool {
return @intFromEnum(path) & 1 == 1;
}

pub inline fn rightN(path: Path, n: Depth) bool {
return (@intFromEnum(path) >> n) & 1 == 1;
}

pub inline fn next(path: *Path) void {
path.* = @enumFromInt(@intFromEnum(path.*) >> 1);
}
Expand Down
44 changes: 36 additions & 8 deletions src/persistent_merkle_tree/node_test.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,23 +239,41 @@ test "Depth helpers - round-trip setNodesAtDepth / getNodesAtDepth" {
const TestCase = struct {
depth: u6,
gindexes: []const usize,
new_nodes: ?u8,
};

fn createTestCase(d: u6, gindexes: anytype) TestCase {
fn createTestCase(d: u6, gindexes: anytype, new_nodes: ?u8) TestCase {
return .{
.depth = d,
.gindexes = &gindexes,
.new_nodes = new_nodes,
};
}

// refer to https://github.com/ChainSafe/ssz/blob/7f5580c2ea69f9307300ddb6010a8bc7ce2fc471/packages/persistent-merkle-tree/test/unit/tree.test.ts#L138
const test_cases = [_]TestCase{ createTestCase(1, [_]usize{2}), createTestCase(1, [_]usize{ 2, 3 }), createTestCase(2, [_]usize{4}), createTestCase(2, [_]usize{6}), createTestCase(2, [_]usize{ 4, 6 }), createTestCase(3, [_]usize{9}), createTestCase(3, [_]usize{12}), createTestCase(3, [_]usize{ 9, 10 }), createTestCase(3, [_]usize{ 13, 14 }), createTestCase(3, [_]usize{ 9, 10, 13, 14 }), createTestCase(3, [_]usize{ 8, 9, 10, 11, 12, 13, 14, 15 }), createTestCase(4, [_]usize{16}), createTestCase(4, [_]usize{ 16, 17 }), createTestCase(4, [_]usize{ 16, 20 }), createTestCase(4, [_]usize{ 16, 20, 30 }), createTestCase(4, [_]usize{ 16, 20, 30, 31 }), createTestCase(5, [_]usize{33}), createTestCase(5, [_]usize{ 33, 34 }), createTestCase(10, [_]usize{ 1024, 1061, 1098, 1135, 1172, 1209, 1246, 1283 }), createTestCase(
40,
[_]usize{ (2 << 39) + 1000, (2 << 39) + 1_000_000, (2 << 39) + 1_000_000_000 },
), createTestCase(
40,
[_]usize{ 1157505940782, 1349082402477, 1759777921993 },
) };
const test_cases = [_]TestCase{
createTestCase(1, [_]usize{2}, null), createTestCase(1, [_]usize{ 2, 3 }, null), createTestCase(2, [_]usize{4}, null), createTestCase(2, [_]usize{6}, null), createTestCase(2, [_]usize{ 4, 6 }, null), createTestCase(3, [_]usize{9}, null), createTestCase(3, [_]usize{12}, null), createTestCase(3, [_]usize{ 9, 10 }, null), createTestCase(3, [_]usize{ 13, 14 }, null), createTestCase(3, [_]usize{ 9, 10, 13, 14 }, null), createTestCase(3, [_]usize{ 8, 9, 10, 11, 12, 13, 14, 15 }, null), createTestCase(4, [_]usize{16}, null), createTestCase(4, [_]usize{ 16, 17 }, null), createTestCase(4, [_]usize{ 16, 20 }, null), createTestCase(4, [_]usize{ 16, 20, 30 }, null), createTestCase(4, [_]usize{ 16, 20, 30, 31 }, null), createTestCase(5, [_]usize{33}, null), createTestCase(5, [_]usize{ 33, 34 }, null), createTestCase(10, [_]usize{ 1024, 1061, 1098, 1135, 1172, 1209, 1246, 1283 }, null),
createTestCase(40, [_]usize{ (2 << 39) + 1000, (2 << 39) + 1_000_000, (2 << 39) + 1_000_000_000 }, null), createTestCase(40, [_]usize{ 1157505940782, 1349082402477, 1759777921993 }, null),
// 2 nodes next to each other at depth 4 but not sharing the same parent
createTestCase(4, [_]usize{ 17, 18, 31 }, null),
// new tests to also confirm the new nodes created to make sure there is no leaked/orphaned nodes during setNodes apis
// set all leaves at depth 4, need 15 new branch nodes
createTestCase(4, [_]usize{ 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31 }, 15),
// set first and last leafs, need 7 new branch nodes
createTestCase(4, [_]usize{ 16, 31 }, 7),
// set first, second last and last leafs, need 7 new branch nodes
createTestCase(4, [_]usize{ 16, 30, 31 }, 7),
// same to above plus first node in the right branch, need 9 new branch nodes
createTestCase(4, [_]usize{ 16, 24, 30, 31 }, 9),
// same to above, 24 and 25 should need only 1 parent, still need 9 new branch nodes
createTestCase(4, [_]usize{ 16, 24, 25, 30, 31 }, 9),
// first node plus the whole right branch, need 11 new branch nodes
createTestCase(4, [_]usize{ 16, 24, 25, 26, 27, 28, 29, 30, 31 }, 11),
// first node plus even nodes in the right branch, need 11 new branch nodes
createTestCase(4, [_]usize{ 16, 24, 26, 28, 30 }, 11),
// first node plus odd nodes in the right branch, need 11 new branch nodes
createTestCase(4, [_]usize{ 16, 25, 27, 29, 31 }, 11),
};

test "setNodesAtDepth, setNodes vs setNode multiple times" {
const allocator = std.testing.allocator;
Expand Down Expand Up @@ -284,8 +302,18 @@ test "setNodesAtDepth, setNodes vs setNode multiple times" {
root_ok = try root_ok.setNode(p, gindexes[i], leaf);
}

var old_nodes = pool.getNodesInUse();
root = try root.setNodesAtDepth(p, depth, indexes, leaves);
if (tc.new_nodes) |n| {
const new_nodes = pool.getNodesInUse() - old_nodes;
try std.testing.expectEqual(n, new_nodes);
}
old_nodes = pool.getNodesInUse();
root2 = try root.setNodes(p, gindexes, leaves);
if (tc.new_nodes) |n| {
const new_nodes = pool.getNodesInUse() - old_nodes;
try std.testing.expectEqual(n, new_nodes);
}

const hash_ok = root_ok.getRoot(p);

Expand Down