diff --git a/src/persistent_merkle_tree/Node.zig b/src/persistent_merkle_tree/Node.zig index 6a86188..a0d2c14 100644 --- a/src/persistent_merkle_tree/Node.zig +++ b/src/persistent_merkle_tree/Node.zig @@ -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(); @@ -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]; @@ -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 @@ -653,8 +673,12 @@ 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(); } @@ -662,6 +686,13 @@ pub const Id = enum(u32) { 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)])) { @@ -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 @@ -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; } @@ -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); @@ -734,16 +792,21 @@ 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); @@ -751,8 +814,6 @@ pub const Id = enum(u32) { 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 @@ -761,8 +822,12 @@ 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(); } @@ -770,6 +835,13 @@ pub const Id = enum(u32) { 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)])) { @@ -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(); } @@ -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 @@ -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; diff --git a/src/persistent_merkle_tree/gindex.zig b/src/persistent_merkle_tree/gindex.zig index 814cb18..6a4efc5 100644 --- a/src/persistent_merkle_tree/gindex.zig +++ b/src/persistent_merkle_tree/gindex.zig @@ -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); } diff --git a/src/persistent_merkle_tree/node_test.zig b/src/persistent_merkle_tree/node_test.zig index d0cf9e3..5aeb567 100644 --- a/src/persistent_merkle_tree/node_test.zig +++ b/src/persistent_merkle_tree/node_test.zig @@ -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; @@ -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);