diff --git a/merklez/src/lib.nr b/merklez/src/lib.nr index 1e3015a..7e20148 100644 --- a/merklez/src/lib.nr +++ b/merklez/src/lib.nr @@ -38,4 +38,5 @@ mod utils; // Core Merkle operations // ----------------------------------------------------------------------------- pub mod merkle_root; +mod merkle_proof_mixed; pub mod merkle_proof; diff --git a/merklez/src/merkle_proof.nr b/merklez/src/merkle_proof.nr index e69de29..8c2b29a 100644 --- a/merklez/src/merkle_proof.nr +++ b/merklez/src/merkle_proof.nr @@ -0,0 +1,43 @@ +// ============================================================================= +// MERKLE PROOF - Generate a Merkle proof for a specific leaf +// ============================================================================= + +use crate::merkle_proof_mixed::merkle_proof_mixed; +use crate::types::{HashFn, Leaf, Proof}; + +// ============================================================================= +// MERKLE PROOF - Main function +// ============================================================================= + +/// Generates a Merkle proof for a specific leaf in the tree +/// +/// This function uses a bottom-up layer-by-layer approach that works for +/// any tree size (both power-of-2 and non-power-of-2 number of leaves). +/// +/// # Arguments +/// * `leaves` - Array of all leaf hashes in the tree +/// * `len` - Number of actual leaves (must be <= N) +/// * `leaf` - The leaf to generate a proof for +/// * `hash_fn` - Hash function to combine two hashes +/// +/// # Returns +/// A Proof struct containing the sibling nodes needed to verify the leaf +/// +/// # Type Parameters +/// * `N` - Maximum capacity for leaves array +/// * `P` - Maximum capacity for proof nodes (should be log2(N) or more) +/// +/// # Example +/// ``` +/// let leaves: [Leaf; 4] = [leaf_a, leaf_b, leaf_c, leaf_d]; +/// let proof = merkle_proof(leaves, 4, leaf_b, hash_fn); +/// // proof contains the sibling nodes needed to verify leaf_b +/// ``` +pub fn merkle_proof( + leaves: [Leaf; N], + len: u32, + leaf: Leaf, + hash_fn: HashFn, +) -> Proof

{ + merkle_proof_mixed(leaves, len, leaf, hash_fn) +} diff --git a/merklez/src/merkle_proof_mixed.nr b/merklez/src/merkle_proof_mixed.nr index e69de29..303eff4 100644 --- a/merklez/src/merkle_proof_mixed.nr +++ b/merklez/src/merkle_proof_mixed.nr @@ -0,0 +1,133 @@ +// ============================================================================= +// MERKLE PROOF MIXED - Proof generation for any tree (works with all sizes) +// ============================================================================= + +use crate::node::Node; +use crate::side::Side; +use crate::types::{HashFn, Leaf, Proof}; + +// ============================================================================= +// HELPER FUNCTIONS +// ============================================================================= + +/// Find the index of a leaf in an array of leaves +/// Returns the index if found, or N (out of bounds) if not found +fn find_leaf_index(leaves: [Leaf; N], len: u32, target: Leaf) -> u32 { + let mut index: u32 = N; // Default to out of bounds + let mut found = false; + for i in 0..N { + if (i < len) & (!found) { + let mut is_equal = true; + for j in 0..32 { + if leaves[i][j] != target[j] { + is_equal = false; + } + } + if is_equal { + index = i; + found = true; + } + } + } + index +} + +/// Compute the next layer up in the Merkle tree +/// Pairs leaves and hashes them together +/// Odd leaves are promoted to the next layer unchanged +fn up_layer(leaves: [Leaf; N], len: u32, hash_fn: HashFn) -> ([Leaf; N], u32) { + let mut new_layer: [Leaf; N] = [[0u8; 32]; N]; + let mut new_len: u32 = 0; + let mut i: u32 = 0; + + for _k in 0..N { + if i < len { + if i + 1 < len { + // Hash pair of leaves + new_layer[new_len] = hash_fn(leaves[i], leaves[i + 1]); + new_len += 1; + i += 2; + } else { + // Odd leaf, promote to next layer + new_layer[new_len] = leaves[i]; + new_len += 1; + i += 1; + } + } + } + + (new_layer, new_len) +} + +// ============================================================================= +// MERKLE PROOF MIXED - Main function (bottom-up approach) +// ============================================================================= + +/// Generates a Merkle proof using the bottom-up layer-by-layer approach +/// Works for any tree size (power-of-2 and non-power-of-2) +/// +/// # Arguments +/// * `leaves` - Array of all leaf hashes in the tree +/// * `len` - Number of actual leaves (must be <= N) +/// * `leaf` - The leaf to generate a proof for +/// * `hash_fn` - Hash function to combine two hashes +/// +/// # Returns +/// A Proof struct containing the sibling nodes needed to verify the leaf +pub fn merkle_proof_mixed( + leaves: [Leaf; N], + len: u32, + leaf: Leaf, + hash_fn: HashFn, +) -> Proof

{ + let mut proof_nodes: [Node; P] = [Node { data: [0u8; 32], side: Side::left() }; P]; + let mut proof_len: u32 = 0; + + // Handle edge cases: len == 0 or len == 1 means no proof needed + // We use a flag to skip the main logic in these cases + let should_process = len > 1; + + if should_process { + // Find the target leaf index + let initial_index = find_leaf_index(leaves, len, leaf); + assert(initial_index < len, "Leaf does not exist in the tree"); + + // Track current position in the tree + let mut leaf_index: u32 = initial_index; + let mut current_leaves: [Leaf; N] = leaves; + let mut current_len: u32 = len; + + // Build proof by traversing up the tree layer by layer + for _level in 0..P { + if current_len > 1 { + // Determine sibling position + if leaf_index % 2 == 0 { + // Current node is at even index, sibling is to the right + if leaf_index + 1 < current_len { + proof_nodes[proof_len] = Node { + data: current_leaves[leaf_index + 1], + side: Side::right(), + }; + proof_len += 1; + } + // If no sibling (odd-length layer), no node is added + } else { + // Current node is at odd index, sibling is to the left + proof_nodes[proof_len] = Node { + data: current_leaves[leaf_index - 1], + side: Side::left(), + }; + proof_len += 1; + } + + // Move up to the next layer + let (next_layer, next_len) = up_layer(current_leaves, current_len, hash_fn); + current_leaves = next_layer; + current_len = next_len; + leaf_index = leaf_index / 2; + } + } + } + + Proof { nodes: proof_nodes, len: proof_len } +} diff --git a/merklez/src/node.nr b/merklez/src/node.nr index 3dd5a64..77424c8 100644 --- a/merklez/src/node.nr +++ b/merklez/src/node.nr @@ -13,3 +13,15 @@ pub struct Node { // Whether this sibling is on the left or right pub side: Side, } + +impl Eq for Node { + fn eq(self, other: Self) -> bool { + let mut data_equal = true; + for i in 0..32 { + if self.data[i] != other.data[i] { + data_equal = false; + } + } + data_equal & (self.side == other.side) + } +} diff --git a/tests/src/merkle_proof/mod.nr b/tests/src/merkle_proof/mod.nr index e69de29..0d0a3a6 100644 --- a/tests/src/merkle_proof/mod.nr +++ b/tests/src/merkle_proof/mod.nr @@ -0,0 +1,301 @@ +use dep::keccak256::keccak256; +use dep::merklez::merkle_proof::merkle_proof; +use dep::merklez::merkle_root::merkle_root; +use dep::merklez::node::Node; +use dep::merklez::side::Side; +use dep::merklez::types::{Hash, Leaf, Proof, Root}; + +fn keccak_hash(left: Hash, right: Hash) -> Hash { + let mut input: [u8; 64] = [0u8; 64]; + for i in 0..32 { + input[i] = left[i]; + input[i + 32] = right[i]; + } + keccak256(input, 64u32) +} + +/// Helper function to verify a proof by computing the root +fn verify_proof(leaf: Leaf, proof: Proof

, hash_fn: fn(Hash, Hash) -> Hash) -> Root { + let mut current: Hash = leaf; + for i in 0..P { + if i < proof.len { + let node = proof.nodes[i]; + if node.side.is_right { + // Sibling is on the right + current = hash_fn(current, node.data); + } else { + // Sibling is on the left + current = hash_fn(node.data, current); + } + } + } + current +} + +// ============================================================================= +// TEST DATA - Standard leaves keccak256 of "a", "b", "c", "d", "e", "f" +// ============================================================================= + +fn get_leaf_a() -> Leaf { + [ + 58, 194, 37, 22, 141, 245, 66, 18, 162, 92, 28, 1, 253, 53, 190, 191, 234, 64, 143, 218, + 194, 227, 29, 221, 111, 128, 164, 187, 249, 165, 241, 203, + ] +} + +fn get_leaf_b() -> Leaf { + [ + 181, 85, 61, 227, 21, 224, 237, 245, 4, 217, 21, 10, 248, 45, 175, 165, 196, 102, 127, + 166, 24, 237, 10, 111, 25, 198, 155, 65, 22, 108, 85, 16, + ] +} + +fn get_leaf_c() -> Leaf { + [ + 11, 66, 182, 57, 60, 31, 83, 6, 15, 227, 221, 191, 205, 122, 173, 204, 168, 148, 70, 90, + 90, 67, 143, 105, 200, 125, 121, 11, 34, 153, 185, 178, + ] +} + +fn get_leaf_d() -> Leaf { + [ + 241, 145, 142, 133, 98, 35, 110, 177, 122, 220, 133, 2, 51, 47, 76, 156, 130, 188, 20, + 225, 155, 252, 10, 161, 10, 182, 116, 255, 117, 179, 210, 243, + ] +} + +fn get_leaf_e() -> Leaf { + [ + 168, 152, 44, 137, 216, 9, 135, 251, 154, 81, 14, 37, 152, 30, 233, 23, 2, 6, 190, 33, + 175, 60, 142, 14, 179, 18, 239, 29, 51, 130, 231, 97, + ] +} + +fn get_leaf_f() -> Leaf { + [ + 209, 232, 174, 183, 149, 0, 73, 110, 243, 220, 46, 87, 186, 116, 106, 131, 21, 208, 72, + 183, 166, 100, 162, 191, 148, 141, 180, 250, 145, 150, 4, 131, + ] +} + +// ============================================================================= +// TESTS - 6 leaves (even, non-power of 2) - Proof for leaf E (index 4) +// ============================================================================= + +#[test] +fn test_merkle_proof_leaves_even_make_proof() { + // 6 leaves: [a, b, c, d, e, f] + let leaves: [Leaf; 6] = [ + get_leaf_a(), + get_leaf_b(), + get_leaf_c(), + get_leaf_d(), + get_leaf_e(), + get_leaf_f(), + ]; + + // Target: leaf E (index 4) + let target = get_leaf_e(); + + // Expected proof from Rust implementation + let expected_proof: [Node; 2] = [ + Node { + data: [ + 209, 232, 174, 183, 149, 0, 73, 110, 243, 220, 46, 87, 186, 116, 106, 131, 21, + 208, 72, 183, 166, 100, 162, 191, 148, 141, 180, 250, 145, 150, 4, 131, + ], + side: Side::right(), + }, + Node { + data: [ + 104, 32, 63, 144, 233, 208, 125, 197, 133, 146, 89, 215, 83, 110, 135, 166, + 186, 157, 52, 95, 37, 82, 181, 185, 222, 41, 153, 221, 206, 156, 225, 191, + ], + side: Side::left(), + }, + ]; + + let proof: Proof<4> = merkle_proof(leaves, 6, target, keccak_hash); + + assert(proof.len == 2); + assert(proof.nodes[0] == expected_proof[0]); + assert(proof.nodes[1] == expected_proof[1]); + + // Also verify the proof produces the correct root + let expected_root = merkle_root(leaves, 6, keccak_hash); + let computed_root = verify_proof(target, proof, keccak_hash); + assert(computed_root == expected_root); +} + +// ============================================================================= +// TESTS - 5 leaves (odd, non-power of 2) - Proof for leaf D (index 3) +// ============================================================================= + +#[test] +fn test_merkle_proof_leaves_odd_make_proof() { + // 5 leaves: [a, b, c, d, e] + let leaves: [Leaf; 5] = [ + get_leaf_a(), + get_leaf_b(), + get_leaf_c(), + get_leaf_d(), + get_leaf_e(), + ]; + + // Target: leaf D (index 3) + let target = get_leaf_d(); + + // Expected proof from Rust implementation + let expected_proof: [Node; 3] = [ + Node { + data: [ + 11, 66, 182, 57, 60, 31, 83, 6, 15, 227, 221, 191, 205, 122, 173, 204, 168, + 148, 70, 90, 90, 67, 143, 105, 200, 125, 121, 11, 34, 153, 185, 178, + ], + side: Side::left(), + }, + Node { + data: [ + 128, 91, 33, 216, 70, 177, 137, 239, 174, 176, 55, 125, 107, 176, 210, 1, 179, + 135, 42, 54, 62, 96, 124, 37, 8, 143, 2, 91, 12, 106, 225, 248, + ], + side: Side::left(), + }, + Node { + data: [ + 168, 152, 44, 137, 216, 9, 135, 251, 154, 81, 14, 37, 152, 30, 233, 23, 2, 6, + 190, 33, 175, 60, 142, 14, 179, 18, 239, 29, 51, 130, 231, 97, + ], + side: Side::right(), + }, + ]; + + let proof: Proof<4> = merkle_proof(leaves, 5, target, keccak_hash); + + assert(proof.len == 3); + assert(proof.nodes[0] == expected_proof[0]); + assert(proof.nodes[1] == expected_proof[1]); + assert(proof.nodes[2] == expected_proof[2]); + + // Also verify the proof produces the correct root + let expected_root = merkle_root(leaves, 5, keccak_hash); + let computed_root = verify_proof(target, proof, keccak_hash); + assert(computed_root == expected_root); +} + +// ============================================================================= +// TESTS - 4 leaves (power of 2) - Proof for leaf A (index 0) +// ============================================================================= + +#[test] +fn test_merkle_proof_leaves_base_2_make_proof() { + // 4 leaves: [a, b, c, d] + let leaves: [Leaf; 4] = [get_leaf_a(), get_leaf_b(), get_leaf_c(), get_leaf_d()]; + + // Target: leaf A (index 0) + let target = get_leaf_a(); + + // Expected proof from Rust implementation + let expected_proof: [Node; 2] = [ + Node { + data: [ + 181, 85, 61, 227, 21, 224, 237, 245, 4, 217, 21, 10, 248, 45, 175, 165, 196, + 102, 127, 166, 24, 237, 10, 111, 25, 198, 155, 65, 22, 108, 85, 16, + ], + side: Side::right(), + }, + Node { + data: [ + 210, 83, 165, 45, 76, 176, 13, 226, 137, 94, 133, 242, 82, 158, 41, 118, 230, + 170, 170, 92, 24, 16, 107, 104, 171, 102, 129, 62, 20, 65, 86, 105, + ], + side: Side::right(), + }, + ]; + + let proof: Proof<4> = merkle_proof(leaves, 4, target, keccak_hash); + + assert(proof.len == 2); + assert(proof.nodes[0] == expected_proof[0]); + assert(proof.nodes[1] == expected_proof[1]); + + // Also verify the proof produces the correct root + let expected_root = merkle_root(leaves, 4, keccak_hash); + let computed_root = verify_proof(target, proof, keccak_hash); + assert(computed_root == expected_root); +} + +// ============================================================================= +// TESTS - Additional coverage +// ============================================================================= + +#[test] +fn test_merkle_proof_2_leaves_first() { + let leaves: [Leaf; 2] = [get_leaf_a(), get_leaf_b()]; + let target = get_leaf_a(); + let expected_root = merkle_root(leaves, 2, keccak_hash); + + let proof: Proof<2> = merkle_proof(leaves, 2, target, keccak_hash); + + assert(proof.len == 1); + // First leaf (index 0) - sibling is on the right + assert(proof.nodes[0].side == Side::right()); + assert(proof.nodes[0].data == get_leaf_b()); + + let computed_root = verify_proof(target, proof, keccak_hash); + assert(computed_root == expected_root); +} + +#[test] +fn test_merkle_proof_2_leaves_second() { + let leaves: [Leaf; 2] = [get_leaf_a(), get_leaf_b()]; + let target = get_leaf_b(); + let expected_root = merkle_root(leaves, 2, keccak_hash); + + let proof: Proof<2> = merkle_proof(leaves, 2, target, keccak_hash); + + assert(proof.len == 1); + // Second leaf (index 1) - sibling is on the left + assert(proof.nodes[0].side == Side::left()); + assert(proof.nodes[0].data == get_leaf_a()); + + let computed_root = verify_proof(target, proof, keccak_hash); + assert(computed_root == expected_root); +} + +#[test] +fn test_merkle_proof_single_leaf() { + let leaves: [Leaf; 1] = [get_leaf_a()]; + let target = get_leaf_a(); + + let proof: Proof<1> = merkle_proof(leaves, 1, target, keccak_hash); + + // Single leaf tree has no proof nodes + assert(proof.len == 0); +} + +#[test] +fn test_merkle_proof_3_leaves() { + let leaves: [Leaf; 3] = [get_leaf_a(), get_leaf_b(), get_leaf_c()]; + let target = get_leaf_c(); + let expected_root = merkle_root(leaves, 3, keccak_hash); + + let proof: Proof<3> = merkle_proof(leaves, 3, target, keccak_hash); + + let computed_root = verify_proof(target, proof, keccak_hash); + assert(computed_root == expected_root); +} + +#[test] +fn test_merkle_proof_4_leaves_all_positions() { + let leaves: [Leaf; 4] = [get_leaf_a(), get_leaf_b(), get_leaf_c(), get_leaf_d()]; + let expected_root = merkle_root(leaves, 4, keccak_hash); + + // Test all 4 positions + for i in 0..4 { + let target = leaves[i]; + let proof: Proof<4> = merkle_proof(leaves, 4, target, keccak_hash); + let computed_root = verify_proof(target, proof, keccak_hash); + assert(computed_root == expected_root); + } +}