From 47dd58ab60106a633b013c4df8ca672eac4506e1 Mon Sep 17 00:00:00 2001 From: twoeths <10568965+twoeths@users.noreply.github.com> Date: Tue, 7 Oct 2025 15:35:33 +0700 Subject: [PATCH] fix: bind to next node if possible --- src/persistent_merkle_tree/Node.zig | 44 +++++++++++++++++------- src/persistent_merkle_tree/gindex.zig | 8 +++++ src/persistent_merkle_tree/node_test.zig | 2 ++ 3 files changed, 42 insertions(+), 12 deletions(-) diff --git a/src/persistent_merkle_tree/Node.zig b/src/persistent_merkle_tree/Node.zig index f8ef75f..3d08f0d 100644 --- a/src/persistent_merkle_tree/Node.zig +++ b/src/persistent_merkle_tree/Node.zig @@ -648,25 +648,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 @@ -715,13 +719,15 @@ pub const Id = enum(u32) { } 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 { @@ -746,6 +752,12 @@ pub const Id = enum(u32) { } } 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; } @@ -792,16 +804,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); @@ -809,8 +826,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 @@ -865,7 +880,7 @@ 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 { @@ -891,6 +906,11 @@ pub const Id = enum(u32) { 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 779aa9a..289dc8c 100644 --- a/src/persistent_merkle_tree/node_test.zig +++ b/src/persistent_merkle_tree/node_test.zig @@ -272,6 +272,8 @@ const test_cases = [_]TestCase{ createTestCase(4, [_]usize{ 16, 20 }, null), createTestCase(4, [_]usize{ 16, 20, 30 }, null), createTestCase(4, [_]usize{ 16, 20, 30, 31 }, null), + // 2 nodes next to each other at depth 4 but not sharing the same parent + createTestCase(4, [_]usize{ 17, 18, 31 }, null), // depth 5 createTestCase(5, [_]usize{33}, null), createTestCase(5, [_]usize{ 33, 34 }, null),