feat(merkle-tree): Contract with tests (NethermindEth#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 <[email protected]>

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)