diff --git a/src/data_structures/README.md b/src/data_structures/README.md index bec6ae0a..33cde852 100644 --- a/src/data_structures/README.md +++ b/src/data_structures/README.md @@ -8,6 +8,7 @@ The Merkle tree algorithm is a cryptographic hashing algorithm used to create a The purpose of the Merkle tree algorithm is to provide a way to verify the integrity and authenticity of large amounts of data without needing to store the entire data set. The algorithm has applications in various areas of computer science, including cryptocurrency, file sharing, and database management. The Merkle tree algorithm is also used in creating digital signatures and verifying the authenticity of transactions. By providing a secure and efficient way to verify data integrity, the Merkle tree algorithm is an important tool in cryptography and information security. +A generic implementation is available to manage both pedersen (legacy) and poseidon hash methods. ## [Queue](./src/queue.cairo) diff --git a/src/data_structures/src/merkle_tree.cairo b/src/data_structures/src/merkle_tree.cairo index a4b34cd8..09d75f33 100644 --- a/src/data_structures/src/merkle_tree.cairo +++ b/src/data_structures/src/merkle_tree.cairo @@ -1,52 +1,121 @@ //! MerkleTree implementation. //! -//! # Example +//! # Examples +//! +//! ``` +//! // This version uses the pedersen hash method because the PedersenHasherImpl is in the scope. +//! use alexandria_data_structures::merkle_tree::{Hasher, MerkleTree, pedersen::PedersenHasherImpl, MerkleTreeTrait}; +//! +//! // Create a new merkle tree instance. +//! let mut merkle_tree: MerkleTree = MerkleTreeTrait::new(); +//! let mut proof = array![element_1, element_2]; +//! // Compute the merkle root. +//! let root = merkle_tree.compute_root(leaf, proof); +//! ``` +//! //! ``` -//! use alexandria::data_structures::merkle_tree::MerkleTreeTrait; +//! // This version uses the poseidon hash method because the PoseidonHasherImpl is in the scope. +//! use alexandria_data_structures::merkle_tree::{ Hasher, MerkleTree, poseidon::PoseidonHasherImpl, MerkleTreeTrait }; //! //! // Create a new merkle tree instance. -//! let mut merkle_tree = MerkleTreeTrait::new(); -//! let mut proof = array![]; -//! proof.append(element_1); -//! proof.append(element_2); +//! let mut merkle_tree: MerkleTree = MerkleTreeTrait::new(); +//! let mut proof = array![element_1, element_2]; //! // Compute the merkle root. //! let root = merkle_tree.compute_root(leaf, proof); +//! ``` // Core lib imports -use array::SpanTrait; -use hash::LegacyHash; -use traits::Into; +use array::{ArrayTrait, SpanTrait}; +use traits::{Into, Copy, Drop}; + +/// Hasher trait. + +trait HasherTrait { + fn new() -> T; + fn hash(ref self: T, data1: felt252, data2: felt252) -> felt252; +} + + +// Hasher representations. + +#[derive(Drop, Copy)] +struct Hasher {} + +/// Hasher impls. + +mod pedersen { + use pedersen::PedersenTrait; + use hash::HashStateTrait; + use super::{Hasher, HasherTrait}; + + impl PedersenHasherImpl of HasherTrait { + fn new() -> Hasher { + Hasher {} + } + fn hash(ref self: Hasher, data1: felt252, data2: felt252) -> felt252 { + let mut state = PedersenTrait::new(data1); + state = state.update(data2); + state.finalize() + } + } +} + +mod poseidon { + use poseidon::PoseidonTrait; + use hash::HashStateTrait; + use super::{Hasher, HasherTrait}; + + impl PoseidonHasherImpl of HasherTrait { + fn new() -> Hasher { + Hasher {} + } + fn hash(ref self: Hasher, data1: felt252, data2: felt252) -> felt252 { + let mut state = PoseidonTrait::new(); + state = state.update(data1); + state = state.update(data2); + state.finalize() + } + } +} /// MerkleTree representation. + #[derive(Drop)] -struct MerkleTree {} +struct MerkleTree { + hasher: T +} /// MerkleTree trait. -trait MerkleTreeTrait { +trait MerkleTreeTrait { /// Create a new merkle tree instance. - fn new() -> MerkleTree; + fn new() -> MerkleTree; /// Compute the merkle root of a given proof. - fn compute_root(ref self: MerkleTree, current_node: felt252, proof: Span) -> felt252; + fn compute_root( + ref self: MerkleTree, current_node: felt252, proof: Span + ) -> felt252; /// Verify a merkle proof. - fn verify(ref self: MerkleTree, root: felt252, leaf: felt252, proof: Span) -> bool; + fn verify(ref self: MerkleTree, root: felt252, leaf: felt252, proof: Span) -> bool; + /// Compute a merkle proof of given leaves and at a given index. + fn compute_proof(ref self: MerkleTree, leaves: Array, index: u32) -> Span; } -/// MerkleTree implementation. -impl MerkleTreeImpl of MerkleTreeTrait { +/// MerkleTree Legacy implementation. +impl MerkleTreeImpl< + T, impl THasher: HasherTrait, impl TCopy: Copy, impl TDrop: Drop +> of MerkleTreeTrait { /// Create a new merkle tree instance. - #[inline(always)] - fn new() -> MerkleTree { - MerkleTree {} + fn new() -> MerkleTree { + MerkleTree { hasher: HasherTrait::new() } } - /// Compute the merkle root of a given proof. + /// Compute the merkle root of a given proof using the generic T hasher. /// # Arguments /// * `current_node` - The current node of the proof. /// * `proof` - The proof. /// # Returns /// The merkle root. fn compute_root( - ref self: MerkleTree, mut current_node: felt252, mut proof: Span + ref self: MerkleTree, mut current_node: felt252, mut proof: Span ) -> felt252 { loop { match proof.pop_front() { @@ -54,11 +123,12 @@ impl MerkleTreeImpl of MerkleTreeTrait { // Compute the hash of the current node and the current element of the proof. // We need to check if the current node is smaller than the current element of the proof. // If it is, we need to swap the order of the hash. - if Into::::into(current_node) < (*proof_element).into() { - current_node = LegacyHash::hash(current_node, *proof_element); - } else { - current_node = LegacyHash::hash(*proof_element, current_node); - } + current_node = + if Into::::into(current_node) < (*proof_element).into() { + self.hasher.hash(current_node, *proof_element) + } else { + self.hasher.hash(*proof_element, current_node) + }; }, Option::None => { break current_node; @@ -67,7 +137,7 @@ impl MerkleTreeImpl of MerkleTreeTrait { } } - /// Verify a merkle proof. + /// Verify a merkle proof using the generic T hasher. /// # Arguments /// * `root` - The merkle root. /// * `leaf` - The leaf to verify. @@ -75,9 +145,107 @@ impl MerkleTreeImpl of MerkleTreeTrait { /// # Returns /// True if the proof is valid, false otherwise. fn verify( - ref self: MerkleTree, root: felt252, leaf: felt252, mut proof: Span + ref self: MerkleTree, root: felt252, mut leaf: felt252, mut proof: Span ) -> bool { - let computed_root = self.compute_root(leaf, proof); + let computed_root = loop { + match proof.pop_front() { + Option::Some(proof_element) => { + // Compute the hash of the current node and the current element of the proof. + // We need to check if the current node is smaller than the current element of the proof. + // If it is, we need to swap the order of the hash. + leaf = + if Into::::into(leaf) < (*proof_element).into() { + self.hasher.hash(leaf, *proof_element) + } else { + self.hasher.hash(*proof_element, leaf) + }; + }, + Option::None => { + break leaf; + }, + }; + }; computed_root == root } + + /// Compute a merkle proof of given leaves and at a given index using the generic T hasher. + /// # Arguments + /// * `leaves` - The sorted leaves. + /// * `index` - The index of the given. + /// # Returns + /// The merkle proof. + fn compute_proof( + ref self: MerkleTree, mut leaves: Array, index: u32 + ) -> Span { + let mut proof: Array = array![]; + compute_proof(leaves, self.hasher, index, ref proof); + proof.span() + } +} + +/// Helper function to compute a merkle proof of given leaves and at a given index. +/// # Arguments +/// * `nodes` - The sorted nodes. +/// * `index` - The index of the given. +/// * `hasher` - The hasher to use. +/// * `proof` - The proof array to fill. +fn compute_proof, impl TDrop: Drop>( + mut nodes: Array, mut hasher: T, index: u32, ref proof: Array +) { + // Break if we have reached the top of the tree + if nodes.len() == 1 { + return; + } + + // If odd number of nodes, add a null virtual leaf + if nodes.len() % 2 != 0 { + nodes.append(0); + } + + // Compute next level + let mut next_level: Array = get_next_level(nodes.span(), ref hasher); + + // Find neighbor node + let mut index_parent = 0; + let mut i = 0; + loop { + if i == index { + index_parent = i / 2; + if i % 2 == 0 { + proof.append(*nodes.at(i + 1)); + } else { + proof.append(*nodes.at(i - 1)); + } + break; + } + i += 1; + }; + + compute_proof(next_level, hasher, index_parent, ref proof) +} + +/// Helper function to compute the next layer of a merkle tree providing a layer of nodes. +/// # Arguments +/// * `nodes` - The sorted nodes. +/// * `hasher` - The hasher to use. +/// # Returns +/// The next layer of nodes. +fn get_next_level, impl TDrop: Drop>( + mut nodes: Span, ref hasher: T +) -> Array { + let mut next_level: Array = array![]; + loop { + if nodes.is_empty() { + break; + } + let left = *nodes.pop_front().expect('Index out of bounds'); + let right = *nodes.pop_front().expect('Index out of bounds'); + let node = if Into::::into(left) < right.into() { + hasher.hash(left, right) + } else { + hasher.hash(right, left) + }; + next_level.append(node); + }; + next_level } diff --git a/src/data_structures/src/tests/merkle_tree_test.cairo b/src/data_structures/src/tests/merkle_tree_test.cairo index 78e4fe81..28173576 100644 --- a/src/data_structures/src/tests/merkle_tree_test.cairo +++ b/src/data_structures/src/tests/merkle_tree_test.cairo @@ -1,74 +1,141 @@ -// Core lib imports -use array::ArrayTrait; -use hash::LegacyHash; // Internal imports -use alexandria_data_structures::merkle_tree::{MerkleTree, MerkleTreeTrait}; +use alexandria_data_structures::merkle_tree::{ + Hasher, MerkleTree, pedersen::PedersenHasherImpl, poseidon::PoseidonHasherImpl, MerkleTreeTrait, + MerkleTreeImpl +}; + + +mod regular_call_merkle_tree_pedersen { + // Internal imports + use alexandria_data_structures::merkle_tree::{ + Hasher, MerkleTree, pedersen::PedersenHasherImpl, MerkleTreeTrait, + }; + #[test] + #[available_gas(2000000)] + fn regular_call_merkle_tree_pedersen_test() { + // [Setup] Merkle tree. + let mut merkle_tree: MerkleTree = MerkleTreeTrait::new(); + let root = 0x15ac9e457789ef0c56e5d559809e7336a909c14ee2511503fa7af69be1ba639; + let leaf = 0x1; + let valid_proof = array![ + 0x2, 0x68ba2a188dd231112c1cb5aaa5d18be6d84f6c8683e5c3a6638dee83e727acc + ] + .span(); + let leaves = array![0x1, 0x2, 0x3]; + + // [Assert] Compute merkle root. + let computed_root = merkle_tree.compute_root(leaf, valid_proof); + assert(computed_root == root, 'compute valid root failed'); + + // [Assert] Compute merkle proof. + let mut input_leaves = leaves; + let index = 0; + let computed_proof = merkle_tree.compute_proof(input_leaves, index); + assert(computed_proof == valid_proof, 'compute valid proof failed'); + + // [Assert] Verify a valid proof. + let result = merkle_tree.verify(root, leaf, valid_proof); + assert(result, 'verify valid proof failed'); + + // [Assert] Verify an invalid proof. + let invalid_proof = array![ + 0x2 + 1, 0x68ba2a188dd231112c1cb5aaa5d18be6d84f6c8683e5c3a6638dee83e727acc + ] + .span(); + let result = merkle_tree.verify(root, leaf, invalid_proof); + assert(!result, 'verify invalid proof failed'); + + // [Assert] Verify a valid proof with an invalid leaf. + let invalid_leaf = 0x1 + 1; + let result = merkle_tree.verify(root, invalid_leaf, valid_proof); + assert(!result, 'wrong result'); + } +} #[test] #[available_gas(2000000)] -fn merkle_tree_test() { - let mut merkle_tree = MerkleTreeTrait::new(); - // Create a proof. - let proof = generate_proof_2_elements( - 275015828570532818958877094293872118179858708489648969448465143543997518327, - 3081470326846576744486900207655708080595997326743041181982939514729891127832 - ); - - let leaf = 1743721452664603547538108163491160873761573033120794192633007665066782417603; - let expected_merkle_root = - 455571898402516024591265345720711356365422160584912150000578530706912124657; - test_case_compute_root(ref merkle_tree, proof, leaf, expected_merkle_root); - - // Create a valid proof. - let mut valid_proof = generate_proof_2_elements( - 275015828570532818958877094293872118179858708489648969448465143543997518327, - 3081470326846576744486900207655708080595997326743041181982939514729891127832 - ); - // Verify the proof is valid. - test_case_verify(ref merkle_tree, expected_merkle_root, leaf, valid_proof, true); - - // Create an invalid proof. - let invalid_proof = generate_proof_2_elements( - 275015828570532818958877094293872118179858708489648969448465143543997518327 + 1, - 3081470326846576744486900207655708080595997326743041181982939514729891127832 - ); - // Verify the proof is invalid. - test_case_verify(ref merkle_tree, expected_merkle_root, leaf, invalid_proof, false); - - // Create a valid proof but we will pass a wrong leaf. - let valid_proof = generate_proof_2_elements( - 275015828570532818958877094293872118179858708489648969448465143543997518327, - 3081470326846576744486900207655708080595997326743041181982939514729891127832 - ); - // Verify the proof is invalid when passed the wrong leaf to verify. - test_case_verify( - ref merkle_tree, - expected_merkle_root, - 1743721452664603547538108163491160873761573033120794192633007665066782417603 + 1, - valid_proof, - false - ); -} +fn merkle_tree_pedersen_test() { + // [Setup] Merkle tree. + let mut merkle_tree: MerkleTree = MerkleTreeImpl::<_, PedersenHasherImpl>::new(); + let root = 0x15ac9e457789ef0c56e5d559809e7336a909c14ee2511503fa7af69be1ba639; + let leaf = 0x1; + let valid_proof = array![0x2, 0x68ba2a188dd231112c1cb5aaa5d18be6d84f6c8683e5c3a6638dee83e727acc] + .span(); + let leaves = array![0x1, 0x2, 0x3]; -fn test_case_compute_root( - ref merkle_tree: MerkleTree, proof: Array, leaf: felt252, expected_root: felt252 -) { - let mut merkle_tree = MerkleTreeTrait::new(); - let root = merkle_tree.compute_root(leaf, proof.span()); - assert(root == expected_root, 'wrong result'); -} + // [Assert] Compute merkle root. + let computed_root = MerkleTreeImpl::<_, + PedersenHasherImpl>::compute_root(ref merkle_tree, leaf, valid_proof); + assert(computed_root == root, 'compute valid root failed'); -fn test_case_verify( - ref merkle_tree: MerkleTree, - root: felt252, - leaf: felt252, - proof: Array, - expected_result: bool -) { - let result = merkle_tree.verify(root, leaf, proof.span()); - assert(result == expected_result, 'wrong result'); + // [Assert] Compute merkle proof. + let mut input_leaves = leaves; + let index = 0; + let computed_proof = MerkleTreeImpl::<_, + PedersenHasherImpl>::compute_proof(ref merkle_tree, input_leaves, index); + assert(computed_proof == valid_proof, 'compute valid proof failed'); + + // [Assert] Verify a valid proof. + let result = MerkleTreeImpl::<_, + PedersenHasherImpl>::verify(ref merkle_tree, root, leaf, valid_proof); + assert(result, 'verify valid proof failed'); + + // [Assert] Verify an invalid proof. + let invalid_proof = array![ + 0x2 + 1, 0x68ba2a188dd231112c1cb5aaa5d18be6d84f6c8683e5c3a6638dee83e727acc + ] + .span(); + let result = MerkleTreeImpl::<_, + PedersenHasherImpl>::verify(ref merkle_tree, root, leaf, invalid_proof); + assert(!result, 'verify invalid proof failed'); + + // [Assert] Verify a valid proof with an invalid leaf. + let invalid_leaf = 0x1 + 1; + let result = MerkleTreeImpl::<_, + PedersenHasherImpl>::verify(ref merkle_tree, root, invalid_leaf, valid_proof); + assert(!result, 'wrong result'); } -fn generate_proof_2_elements(element_1: felt252, element_2: felt252) -> Array { - array![element_1, element_2] +#[test] +#[available_gas(2000000)] +fn merkle_tree_poseidon_test() { + // [Setup] Merkle tree. + let mut merkle_tree: MerkleTree = MerkleTreeImpl::<_, PoseidonHasherImpl>::new(); + let root = 0x7abc09d19c8a03abd4333a23f7823975c7bdd325170f0d32612b8baa1457d47; + let leaf = 0x1; + let valid_proof = array![0x2, 0x47ef3ad11ad3f8fc055281f1721acd537563ec134036bc4bd4de2af151f0832] + .span(); + let leaves = array![0x1, 0x2, 0x3]; + + // [Assert] Compute merkle root. + let computed_root = MerkleTreeImpl::<_, + PoseidonHasherImpl>::compute_root(ref merkle_tree, leaf, valid_proof); + assert(computed_root == root, 'compute valid root failed'); + + // [Assert] Compute merkle proof. + let mut input_leaves = leaves; + let index = 0; + let computed_proof = MerkleTreeImpl::<_, + PoseidonHasherImpl>::compute_proof(ref merkle_tree, input_leaves, index); + assert(computed_proof == valid_proof, 'compute valid proof failed'); + + // [Assert] Verify a valid proof. + let result = MerkleTreeImpl::<_, + PoseidonHasherImpl>::verify(ref merkle_tree, root, leaf, valid_proof); + assert(result, 'verify valid proof failed'); + + // [Assert] Verify an invalid proof. + let invalid_proof = array![ + 0x2 + 1, 0x68ba2a188dd231112c1cb5aaa5d18be6d84f6c8683e5c3a6638dee83e727acc + ] + .span(); + let result = MerkleTreeImpl::<_, + PoseidonHasherImpl>::verify(ref merkle_tree, root, leaf, invalid_proof); + assert(!result, 'verify invalid proof failed'); + + // [Assert] Verify a valid proof with an invalid leaf. + let invalid_leaf = 0x1 + 1; + let result = MerkleTreeImpl::<_, + PoseidonHasherImpl>::verify(ref merkle_tree, root, invalid_leaf, valid_proof); + assert(!result, 'wrong result'); }