feat(merkle-tree): Contract with tests (#228)

* feat(merkle-tree): Contract with tests

* feat(merkle-tree): Corrections according to PR reviews

* feat(merkle-tree): Contract with tests

* fix: 2024_07 edition

* fix: Replace Map simulating Array with Vec - streamline md file explanations

* fix: scarb fmt

---------

Co-authored-by: julio4 <30329843+julio4@users.noreply.github.com>

Scarb.lock

@@ -110,6 +110,10 @@ version = "0.1.0"
 name = "mappings"
 version = "0.1.0"
 
+[[package]]
+name = "merkle_tree"
+version = "0.1.0"
+
 [[package]]
 name = "nft_dutch_auction"
 version = "0.1.0"

listings/applications/merkle_tree/.gitignore

@@ -0,0 +1 @@
+target

listings/applications/merkle_tree/Scarb.toml

@@ -0,0 +1,15 @@
+[package]
+name = "merkle_tree"
+version.workspace = true
+edition = "2024_07"
+
+[dependencies]
+starknet.workspace = true
+
+[dev-dependencies]
+cairo_test.workspace = true
+
+[scripts]
+test.workspace = true
+
+[[target.starknet-contract]]

listings/applications/merkle_tree/src/contract.cairo

@@ -0,0 +1,117 @@
+#[generate_trait]
+pub impl ByteArrayHashTraitImpl of ByteArrayHashTrait {
+ fn hash(self: @ByteArray) -> felt252 {
+ let mut serialized_byte_arr: Array<felt252> = ArrayTrait::new();
+ self.serialize(ref serialized_byte_arr);
+
+ core::poseidon::poseidon_hash_span(serialized_byte_arr.span())
+ }
+}
+
+#[starknet::interface]
+pub trait IMerkleTree<TContractState> {
+ fn build_tree(ref self: TContractState, data: Array<ByteArray>) -> Array<felt252>;
+ fn get_root(self: @TContractState) -> felt252;
+ // function to verify if leaf node exists in the merkle tree
+ fn verify(
+ self: @TContractState, proof: Array<felt252>, root: felt252, leaf: felt252, index: usize
+ ) -> bool;
+}
+
+mod errors {
+ pub const NOT_POW_2: felt252 = 'Data length is not a power of 2';
+ pub const NOT_PRESENT: felt252 = 'No element in merkle tree';
+}
+
+#[starknet::contract]
+pub mod MerkleTree {
+ use core::poseidon::PoseidonTrait;
+ use core::hash::{HashStateTrait, HashStateExTrait};
+ use starknet::storage::{
+ StoragePointerWriteAccess, StoragePointerReadAccess, Vec, MutableVecTrait, VecTrait
+ };
+ use super::ByteArrayHashTrait;
+
+ #[storage]
+ struct Storage {
+ pub hashes: Vec<felt252>
+ }
+
+ #[derive(Drop, Serde, Copy)]
+ struct Vec2 {
+ x: u32,
+ y: u32
+ }
+
+ #[abi(embed_v0)]
+ impl IMerkleTreeImpl of super::IMerkleTree<ContractState> {
+ fn build_tree(ref self: ContractState, mut data: Array<ByteArray>) -> Array<felt252> {
+ let data_len = data.len();
+ assert(data_len > 0 && (data_len & (data_len - 1)) == 0, super::errors::NOT_POW_2);
+
+ let mut _hashes: Array<felt252> = ArrayTrait::new();
+
+ // first, hash every leaf
+ for value in data {
+ _hashes.append(value.hash());
+ };
+
+ // then, hash all levels above leaves
+ let mut current_nodes_lvl_len = data_len;
+ let mut hashes_offset = 0;
+
+ while current_nodes_lvl_len > 0 {
+ let mut i = 0;
+ while i < current_nodes_lvl_len - 1 {
+ let left_elem = *_hashes.at(hashes_offset + i);
+ let right_elem = *_hashes.at(hashes_offset + i + 1);
+
+ let hash = PoseidonTrait::new().update_with((left_elem, right_elem)).finalize();
+ _hashes.append(hash);
+
+ i += 2;
+ };
+
+ hashes_offset += current_nodes_lvl_len;
+ current_nodes_lvl_len /= 2;
+ };
+
+ // write to the contract state (useful for the get_root function)
+ for hash in _hashes.span() {
+ self.hashes.append().write(*hash);
+ };
+
+ _hashes
+ }
+
+ fn get_root(self: @ContractState) -> felt252 {
+ let merkle_tree_length = self.hashes.len();
+ assert(merkle_tree_length > 0, super::errors::NOT_PRESENT);
+
+ self.hashes.at(merkle_tree_length - 1).read()
+ }
+
+ fn verify(
+ self: @ContractState,
+ mut proof: Array<felt252>,
+ root: felt252,
+ leaf: felt252,
+ mut index: usize
+ ) -> bool {
+ let mut current_hash = leaf;
+
+ while let Option::Some(value) = proof.pop_front() {
+ current_hash =
+ if index % 2 == 0 {
+ PoseidonTrait::new().update_with((current_hash, value)).finalize()
+ } else {
+ PoseidonTrait::new().update_with((value, current_hash)).finalize()
+ };
+
+ index /= 2;
+ };
+
+ current_hash == root
+ }
+ }
+}

listings/applications/merkle_tree/src/lib.cairo

@@ -0,0 +1,4 @@
+mod contract;
+
+#[cfg(test)]
+mod tests;

listings/applications/merkle_tree/src/tests.cairo

@@ -0,0 +1,182 @@
+use merkle_tree::contract::IMerkleTreeDispatcherTrait;
+use merkle_tree::contract::{IMerkleTreeDispatcher, MerkleTree, ByteArrayHashTrait};
+use starknet::syscalls::deploy_syscall;
+use starknet::{ContractAddress, SyscallResultTrait};
+use starknet::testing::set_contract_address;
+use core::poseidon::PoseidonTrait;
+use core::hash::{HashStateTrait, HashStateExTrait};
+use starknet::storage::{VecTrait, StoragePointerReadAccess};
+
+fn deploy_util(class_hash: felt252, calldata: Array<felt252>) -> ContractAddress {
+ let (address, _) = deploy_syscall(class_hash.try_into().unwrap(), 0, calldata.span(), false)
+ .unwrap_syscall();
+ address
+}
+
+fn setup() -> IMerkleTreeDispatcher {
+ let contract_address = deploy_util(MerkleTree::TEST_CLASS_HASH, array![]);
+
+ IMerkleTreeDispatcher { contract_address }
+}
+
+#[test]
+fn should_deploy() {
+ let deploy = setup();
+
+ let state = @MerkleTree::contract_state_for_testing();
+ // "link" a new MerkleTree struct to the deployed MerkleTree contract
+ // in order to access its internal state fields for assertions
+ set_contract_address(deploy.contract_address);
+
+ assert_eq!(state.hashes.len(), 0);
+}
+
+#[test]
+fn build_tree_succeeds() {
+ /// Set up
+ let deploy = setup();
+
+ let data_1 = "alice -> bob: 2";
+ let data_2 = "bob -> john: 5";
+ let data_3 = "alice -> john: 1";
+ let data_4 = "john -> alex: 8";
+ let arguments = array![data_1.clone(), data_2.clone(), data_3.clone(), data_4.clone()];
+
+ /// When
+ let actual_hashes = deploy.build_tree(arguments);
+
+ /// Then
+ let mut expected_hashes: Array<felt252> = array![];
+
+ // leaves' hashes
+ expected_hashes.append(data_1.hash());
+ expected_hashes.append(data_2.hash());
+ expected_hashes.append(data_3.hash());
+ expected_hashes.append(data_4.hash());
+
+ // hashes for level above leaves
+ let hash_0 = PoseidonTrait::new()
+ .update_with((*expected_hashes.at(0), *expected_hashes.at(1)))
+ .finalize();
+ let hash_1 = PoseidonTrait::new()
+ .update_with((*expected_hashes.at(2), *expected_hashes.at(3)))
+ .finalize();
+ expected_hashes.append(hash_0);
+ expected_hashes.append(hash_1);
+
+ // root hash
+ let root_hash = PoseidonTrait::new().update_with((hash_0, hash_1)).finalize();
+ expected_hashes.append(root_hash);
+
+ // verify returned result
+ assert_eq!(actual_hashes, expected_hashes);
+
+ // verify get_root
+ assert_eq!(deploy.get_root(), root_hash);
+
+ // verify contract storage state
+
+ let state = @MerkleTree::contract_state_for_testing();
+ // "link" a new MerkleTree struct to the deployed MerkleTree contract
+ // in order to access its internal state fields for assertions
+ set_contract_address(deploy.contract_address);
+
+ assert_eq!(state.hashes.len(), expected_hashes.len().into());
+
+ for i in 0
+ ..expected_hashes
+ .len() {
+ assert_eq!(state.hashes.at(i.into()).read(), *expected_hashes.at(i));
+ }
+}
+
+#[test]
+#[should_panic(expected: ('Data length is not a power of 2', 'ENTRYPOINT_FAILED'))]
+fn build_tree_fails() {
+ /// Set up
+ let deploy = setup();
+
+ let data_1 = "alice -> bob: 2";
+ let data_2 = "bob -> john: 5";
+ let data_3 = "alice -> john: 1";
+ // number of arguments not a power of 2
+ let arguments = array![data_1, data_2, data_3];
+
+ /// When
+ deploy.build_tree(arguments);
+}
+
+#[test]
+fn verify_leaf_succeeds() {
+ /// Set up
+ let deploy = setup();
+
+ let data_1 = "bob -> alice: 1";
+ let data_2 = "alex -> john: 3";
+ let data_3 = "alice -> alex: 8";
+ let data_4 = "alex -> bob: 8";
+ let arguments = array![data_1.clone(), data_2.clone(), data_3.clone(), data_4.clone()];
+
+ let hashes = deploy.build_tree(arguments);
+
+ // ----> hashes tree :
+ //
+ // hashes[6]
+ // / \
+ // hashes[4] hashes[5]
+ // / \ / \
+ // hashes[0] hashes[1] hashes[2] hashes[3]
+
+ let res = deploy
+ .verify(
+ array![*hashes.at(3), *hashes.at(4)], // proof
+ *hashes.at(6), // root
+ data_3.hash(), // leaf
+ 2 // leaf index
+ );
+
+ assert(res, 'Leaf should be in merkle tree');
+}
+
+#[test]
+#[available_gas(20000000)]
+fn verify_leaf_fails() {
+ /// Set up
+ let deploy = setup();
+
+ let data_1 = "bob -> alice: 1";
+ let data_2 = "alex -> john: 3";
+ let data_3 = "alice -> alex: 8";
+ let data_4 = "alex -> bob: 8";
+ let arguments = array![data_1.clone(), data_2.clone(), data_3.clone(), data_4.clone()];
+
+ let hashes = deploy.build_tree(arguments);
+
+ // ----- hashes tree -----
+ // hashes[6]
+ // / \
+ // hashes[4] hashes[5]
+ // / \ / \
+ // hashes[0] hashes[1] hashes[2] hashes[3]
+
+ let wrong_leaf: ByteArray = "alice -> alex: 9";
+ let res = deploy
+ .verify(
+ array![*hashes.at(3), *hashes.at(4)], // proof
+ *hashes.at(6), // root
+ wrong_leaf.hash(), // leaf
+ 2 // leaf index
+ );
+ assert(!res, '1- Leaf should NOT be in tree');
+
+ let wrong_proof = array![*hashes.at(4), *hashes.at(3)];
+ let res = deploy
+ .verify(
+ wrong_proof, // proof
+ *hashes.at(6), // root
+ data_3.hash(), // leaf
+ 2 // leaf index
+ );
+ assert(!res, '2- Leaf should NOT be in tree');
+}
+

src/SUMMARY.md

@@ -61,6 +61,7 @@ Summary
 - [Constant Product AMM](./applications/constant-product-amm.md)
 - [TimeLock](./applications/timelock.md)
 - [Staking](./applications/staking.md)
+- [Merkle Tree](./applications/merkle_tree.md)
 - [Simple Storage with Starknet-js](./applications/simple_storage_starknetjs.md)
 - [Crowdfunding Campaign](./applications/crowdfunding.md)
 - [AdvancedFactory: Crowdfunding](./applications/advanced_factory.md)