add tweakable merkle tree and path

build.zig

@@ -24,11 +24,17 @@ pub fn build(b: *Builder) void {
         .target = target,
     });
     const run_main_tests = b.addRunArtifact(main_tests);
+
     const tests_tests = b.addTest(.{
         .root_source_file = .{ .cwd_relative = "src/tests.zig" },
         .optimize = optimize,
         .target = target,
     });
+    const tweak_hash = b.dependency("tweak_hash", .{
+        .target = target,
+        .optimize = optimize,
+    }).module("tweak_hash");
+    tests_tests.root_module.addImport("tweak_hash", tweak_hash);
     tests_tests.root_module.addImport("ssz.zig", mod);
     const run_tests_tests = b.addRunArtifact(tests_tests);
 

build.zig.zon

@@ -2,5 +2,11 @@
     .name = .ssz,
     .fingerprint = 0x1d34bd0ceb1dfc2d,
     .version = "0.0.4",
+    .dependencies = .{
+        .tweak_hash = .{
+            .url = "git+https://github.com/bhaskar1001101/tweak-hash#63611fd65d9db28a7ea342ff2fd61b1e201727fd",
+            .hash = "tweak_hash-0.0.0-b5xZ7JkRAABSp1CPFJpayRnzIt2tnLjzjTPfveqLiUXu",
+        },
+    },
     .paths = .{""},
 }

src/lib.zig

@@ -818,3 +818,88 @@ fn bytesToBits(comptime N: usize, src: [N]u8) [N * 8]bool {
     }
     return bitvector;
 }
+
+pub const MerkleTree = [][]u8;
+pub const MerklePath = [][]u8;
+
+// returns a full merkle tree with a tweakable hash function from provided leaf hashes
+pub fn merkleizeWithTweak(comptime TweakHash: type, hash: TweakHash, allocator: std.mem.Allocator, parameter: []u8, chunks: [][]u8) !MerkleTree {
+    const num_leaves = chunks.len;
+    std.debug.assert(std.math.isPowerOfTwo(num_leaves));
+
+    const node_count = (2 * num_leaves) - 1;
+    var nodes = try allocator.alloc([]u8, node_count);
+
+    for (chunks, 0..) |c, i| {
+        const leaf_pos = node_count - num_leaves + i;
+        nodes[leaf_pos] = try allocator.dupe(u8, c);
+    }
+
+    var level: u8 = 1;
+    var level_size: usize = num_leaves / 2;
+    var level_offset: usize = node_count - num_leaves - level_size;
+
+    while (level_size > 0) {
+        for (0..level_size) |i| {
+            const left_child = nodes[level_offset + level_size + i * 2];
+            const right_child = nodes[level_offset + level_size + i * 2 + 1];
+
+            var combined = [_][]u8{ left_child, right_child };
+
+            nodes[level_offset + i] = try allocator.alloc(u8, hash.hash_size);
+            const tweak = hash.tree_tweak(level, @as(u32, @intCast(i)));
+            hash.hash(parameter, tweak, &combined, nodes[level_offset + i]);
+        }
+
+        level += 1;
+        level_size /= 2;
+        level_offset -= level_size;
+    }
+
+    return nodes;
+}
+
+// returns a merkle path from a leaf to the root
+pub fn buildPath(allocator: std.mem.Allocator, tree: MerkleTree, leaf_index: usize) !MerklePath {
+    const height = std.math.log2_int(usize, tree.len);
+
+    var path = try allocator.alloc([]u8, height);
+    var current_index = leaf_index;
+    const num_leaves = @as(u32, 1) << @intCast(height);
+    const total_nodes = (2 * num_leaves) - 1;
+    var node_index = total_nodes - num_leaves + current_index;
+    for (0..height) |level| {
+        const is_left = current_index % 2 == 0;
+        const sibling_offset: isize = if (is_left) 1 else -1;
+
+        const sibling_node_index: usize = @intCast(@as(isize, @intCast(node_index)) + sibling_offset);
+        path[level] = try allocator.dupe(u8, tree[sibling_node_index]);
+
+        current_index /= 2;
+        node_index = (node_index - 1) / 2;
+    }
+
+    return path;
+}
+
+// verify a merkle path
+pub fn verifyPath(comptime TweakHash: type, parameter: []u8, hash: TweakHash, leaf_index: usize, root: []u8, leaf_hash: []u8, path: MerklePath) bool {
+    var current_index = leaf_index;
+    for (0..path.len) |level| {
+        const is_left = current_index % 2 == 0;
+        const sibling = path[level];
+
+        const combined = if (is_left) [_][]u8{ leaf_hash, sibling } else [_][]u8{ sibling, leaf_hash };
+        const tweak = hash.tree_tweak(@as(u8, @intCast(level + 1)), @as(u32, @intCast(current_index / 2)));
+
+        hash.hash(parameter, tweak, &combined, leaf_hash);
+
+        current_index /= 2;
+    }
+
+    return std.mem.eql(u8, leaf_hash, root);
+}
+
+pub fn treeRoot(tree: MerkleTree) []u8 {
+    return tree[0];
+}

src/tests.zig

@@ -857,3 +857,84 @@ test "structs with nested fixed/variable size u8 array" {
     try expect(var_signed_block.message.body.slot == deserialized_var_block.message.body.slot);
     try expect(std.mem.eql(u8, var_signed_block.message.body.data[0..], deserialized_var_block.message.body.data[0..]));
 }
+
+test "MerkleTree build, path, and verify" {
+    const testing = std.testing;
+    const allocator = testing.allocator;
+
+    const num_leaves: usize = 1024;
+    // Each leaf will consist of `leaf_len` chunks of `hash.hash_size` bytes.
+    const leaf_len = 3;
+
+    const parameter = try allocator.alloc(u8, 16);
+    defer allocator.free(parameter);
+    std.crypto.random.bytes(parameter);
+
+    const ShaTweakHash = @import("tweak_hash").ShaTweakHash;
+    const hash = ShaTweakHash.init(16, 24);
+
+    // Generate random leaves. Each leaf is a slice of leaf_len chunks.
+    var leaves = try allocator.alloc([][]u8, num_leaves);
+    defer {
+        for (leaves) |leaf_chunks| {
+            for (leaf_chunks) |chunk| {
+                allocator.free(chunk);
+            }
+            allocator.free(leaf_chunks);
+        }
+        allocator.free(leaves);
+    }
+
+    // Fill every leaf with `leaf_len` random chunks.
+    for (0..num_leaves) |i| {
+        leaves[i] = try allocator.alloc([]u8, leaf_len);
+        for (0..leaf_len) |j| {
+            leaves[i][j] = try allocator.alloc(u8, hash.hash_size);
+            std.crypto.random.bytes(leaves[i][j]);
+        }
+    }
+
+    // Hash the leaves with the level 0 tweak.
+    var leaf_hashes = try allocator.alloc([]u8, num_leaves);
+    defer {
+        for (leaf_hashes) |h| allocator.free(h);
+        allocator.free(leaf_hashes);
+    }
+    for (leaves, 0..) |leaf_chunks, i| {
+        leaf_hashes[i] = try allocator.alloc(u8, hash.hash_size);
+        const tweak = hash.tree_tweak(0, @as(u32, @intCast(i)));
+        hash.hash(parameter, tweak, leaf_chunks, leaf_hashes[i]);
+    }
+
+    const tree = try libssz.merkleizeWithTweak(ShaTweakHash, hash, allocator, parameter, leaf_hashes);
+    defer {
+        for (tree) |node| {
+            allocator.free(node);
+        }
+        allocator.free(tree);
+    }
+
+    const root = libssz.treeRoot(tree);
+
+    // build path for all leaves and verify them
+    for (0..num_leaves) |idx| {
+        const path = try libssz.buildPath(allocator, tree, idx);
+        defer {
+            for (path) |node| {
+                allocator.free(node);
+            }
+            allocator.free(path);
+        }
+
+        const leaf_tweak = hash.tree_tweak(0, @as(u32, @intCast(idx)));
+        const leaf_hash_expected = try allocator.alloc(u8, hash.hash_size);
+        defer allocator.free(leaf_hash_expected);
+        hash.hash(parameter, leaf_tweak, leaves[idx], leaf_hash_expected);
+
+        const leaf_hash_to_verify = try allocator.dupe(u8, leaf_hash_expected);
+        defer allocator.free(leaf_hash_to_verify);
+
+        const ok = libssz.verifyPath(ShaTweakHash, parameter, hash, idx, root, leaf_hash_to_verify, path);
+        try testing.expect(ok);
+    }
+}