Merkle tree tests

shielder/contract/Cargo.toml

@@ -12,6 +12,7 @@ ink = { version = "5.0.0-rc", default-features = false }
 
 [dev-dependencies]
 drink = "0.9.0"
+rand = { version = "0.8.5", default-features = false }
 
 [lib]
 path = "lib.rs"

shielder/contract/drink_tests/mod.rs

@@ -3,16 +3,17 @@ mod utils;
 use anyhow::Result;
 use drink::{session::Session, AccountId32};
 
-use crate::{
-    drink_tests::utils::{chain::*, ops::*, psp22::*, shielder::*},
-    test_utils::merkle::MerkleTree,
-};
+use crate::drink_tests::utils::{chain::*, ops::*, psp22::*, shielder::*};
+
+use rand::prelude::*;
 
 #[drink::contract_bundle_provider]
 pub enum BundleProvider {}
 
 #[drink::test]
 fn deploy_single_deposit_single_withdraw(mut session: Session) -> Result<()> {
+    let mut rng = StdRng::seed_from_u64(1);
+
     let alice = init_alice(&mut session)?;
     let bob = init_bob(&mut session)?;
 
@@ -21,14 +22,12 @@ fn deploy_single_deposit_single_withdraw(mut session: Session) -> Result<()> {
     let shielder_address = deploy_shielder(&mut session)?;
     let psp22_address = deploy_test_token(&mut session, 100)?;
 
-    let mut merkle_tree = MerkleTree::new();
-
     // CREATE ACCOUNT
     let user_shielded_data = create_shielder_account(
         &mut session,
         &shielder_address,
         &psp22_address,
-        &mut merkle_tree,
+        rng.gen::<u128>().into(),
     )?;
 
     // APPROVE TRANSFER
@@ -40,7 +39,7 @@ fn deploy_single_deposit_single_withdraw(mut session: Session) -> Result<()> {
         &shielder_address,
         deposit_op(&psp22_address, &alice, 10),
         user_shielded_data,
-        &mut merkle_tree,
+        rng.gen::<u128>().into(),
     )?;
 
     let alice_psp22_balance = get_psp22_balance(&mut session, &psp22_address, &alice)?;
@@ -58,7 +57,7 @@ fn deploy_single_deposit_single_withdraw(mut session: Session) -> Result<()> {
         &shielder_address,
         withdraw_op(&psp22_address, &bob, 1),
         user_shielded_data,
-        &mut merkle_tree,
+        rng.gen::<u128>().into(),
     )?;
 
     let bob_psp22_balance = get_psp22_balance(&mut session, &psp22_address, &bob)?;
@@ -72,6 +71,8 @@ fn deploy_single_deposit_single_withdraw(mut session: Session) -> Result<()> {
 
 #[drink::test]
 fn deploy_single_deposit_multiple_withdraw(mut session: Session) -> Result<()> {
+    let mut rng = StdRng::seed_from_u64(2);
+
     let alice = init_alice(&mut session)?;
     session = session.with_actor(alice.clone());
 
@@ -85,14 +86,12 @@ fn deploy_single_deposit_multiple_withdraw(mut session: Session) -> Result<()> {
     let shielder_address = deploy_shielder(&mut session)?;
     let psp22_address = deploy_test_token(&mut session, 100)?;
 
-    let mut merkle_tree = MerkleTree::new();
-
     // CREATE ACCOUNT
     let mut user_shielded_data = create_shielder_account(
         &mut session,
         &shielder_address,
         &psp22_address,
-        &mut merkle_tree,
+        rng.gen::<u128>().into(),
     )?;
 
     // APPROVE TRANSFER
@@ -108,7 +107,7 @@ fn deploy_single_deposit_multiple_withdraw(mut session: Session) -> Result<()> {
         &shielder_address,
         deposit_op(&psp22_address, &alice, 50),
         user_shielded_data,
-        &mut merkle_tree,
+        rng.gen::<u128>().into(),
     )?;
 
     for withdrawer_addr in withdrawers {
@@ -121,7 +120,7 @@ fn deploy_single_deposit_multiple_withdraw(mut session: Session) -> Result<()> {
             &shielder_address,
             withdraw_op(&psp22_address, &withdrawer_addr, 1),
             user_shielded_data,
-            &mut merkle_tree,
+            rng.gen::<u128>().into(),
         )?;
         let psp22_balance = get_psp22_balance(&mut session, &psp22_address, &withdrawer_addr)?;
         assert_eq!(psp22_balance, 1);
@@ -132,3 +131,78 @@ fn deploy_single_deposit_multiple_withdraw(mut session: Session) -> Result<()> {
 
     Ok(())
 }
+
+#[drink::test]
+fn deploy_multiple_deposit_multiple_withdraw(mut session: Session) -> Result<()> {
+    let mut rng = StdRng::seed_from_u64(3);
+
+    let alice = init_alice(&mut session)?;
+    session = session.with_actor(alice.clone());
+
+    let mut depositors = vec![];
+    for i in 3..11 {
+        let acc = AccountId32::new([i as u8; 32]);
+        init_acc_with_balance(&mut session, &acc)?;
+        depositors.push(acc.clone());
+    }
+
+    let mut withdrawers = vec![];
+    for i in 11..19 {
+        let acc = AccountId32::new([i as u8; 32]);
+        init_acc_with_balance(&mut session, &acc)?;
+        withdrawers.push(acc.clone());
+    }
+
+    let shielder_address = deploy_shielder(&mut session)?;
+    let psp22_address = deploy_test_token(&mut session, 800)?;
+    for depositor_addr in &depositors {
+        psp22_transfer(&mut session, &psp22_address, &depositor_addr, 100)?;
+    }
+
+    let mut user_shielded_data = vec![];
+    for (i, depositor_addr) in depositors.iter().enumerate() {
+        // SWITCH TO withdrawer
+        session = session.with_actor(depositor_addr.clone());
+
+        // CREATE ACCOUNT
+        user_shielded_data.push(create_shielder_account(
+            &mut session,
+            &shielder_address,
+            &psp22_address,
+            rng.gen::<u128>().into(),
+        )?);
+
+        // APPROVE TRANSFER
+        psp22_approve(&mut session, &psp22_address, &shielder_address, 50)?;
+
+        // DEPOSIT
+        user_shielded_data[i] = shielder_update(
+            &mut session,
+            &shielder_address,
+            deposit_op(&psp22_address, &depositor_addr, 50),
+            user_shielded_data[i],
+            rng.gen::<u128>().into(),
+        )?;
+    }
+
+    for (i, withdrawer_addr) in withdrawers.iter().enumerate() {
+        // SWITCH TO withdrawer
+        session = session.with_actor(withdrawer_addr.clone());
+
+        // WITHDRAW
+        user_shielded_data[i] = shielder_update(
+            &mut session,
+            &shielder_address,
+            withdraw_op(&psp22_address, &withdrawer_addr, 1),
+            user_shielded_data[i],
+            rng.gen::<u128>().into(),
+        )?;
+        let psp22_balance = get_psp22_balance(&mut session, &psp22_address, &withdrawer_addr)?;
+        assert_eq!(psp22_balance, 1);
+    }
+    let shielder_psp22_balance =
+        get_psp22_balance(&mut session, &psp22_address, &shielder_address)?;
+    assert_eq!(shielder_psp22_balance, 400 - 8);
+
+    Ok(())
+}

shielder/contract/drink_tests/utils/psp22.rs

@@ -73,3 +73,20 @@ pub fn psp22_approve(
     )??;
     Ok(())
 }
+
+pub fn psp22_transfer(
+    session: &mut Session<MinimalRuntime>,
+    token: &AccountId32,
+    to: &AccountId32,
+    amount: u128,
+) -> Result<()> {
+    let formatted_amount = format!("{}", amount);
+    let empty_arr: [u8; 0] = [];
+    session.call_with_address(
+        token.clone(),
+        "PSP22::transfer",
+        &[to.to_string(), formatted_amount, format!("{:?}", empty_arr)],
+        NO_ENDOWMENT,
+    )??;
+    Ok(())
+}

shielder/contract/drink_tests/utils/shielder.rs

@@ -15,10 +15,10 @@ use crate::{
         traits::Hashable,
         TOKENS_NUMBER,
     },
-    test_utils::merkle::MerkleTree,
     types::Scalar,
 };
 
+#[derive(Clone, Copy, Debug)]
 pub struct ShielderUserEnv {
     pub proof: ZkProof,
     pub nullifier: Scalar,
@@ -40,15 +40,14 @@ pub fn create_shielder_account(
     session: &mut Session<MinimalRuntime>,
     shielder_address: &AccountId32,
     token: &AccountId32,
-    merkle_tree: &mut MerkleTree,
+    nullifier: Scalar,
 ) -> Result<ShielderUserEnv> {
     let mut tokens: [Scalar; TOKENS_NUMBER] = [0_u128.into(); TOKENS_NUMBER];
     tokens[0] = Scalar::from_bytes(*((*token).as_ref()));
 
     let acc = Account::new(tokens);
 
     let id = 0_128.into();
-    let nullifier = 0_u128.into();
     let trapdoor = 0_u128.into();
     let op_priv = OpPriv {
         user: 0_u128.into(),
@@ -58,19 +57,17 @@ pub fn create_shielder_account(
 
     let h_note_new = Note::new(id, trapdoor, nullifier, acc.hash()).hash();
 
-    session.call_with_address(
+    let leaf_id_res: Result<u32, ()> = session.call_with_address(
         shielder_address.clone(),
         "add_note",
         &[format!("{:?}", h_note_new), format!("{:?}", proof)],
         NO_ENDOWMENT,
     )??;
 
-    merkle_tree.add_leaf(h_note_new).unwrap();
-
     Ok(ShielderUserEnv {
         proof,
         nullifier,
-        tree_leaf_id: 0,
+        tree_leaf_id: leaf_id_res.unwrap(),
     })
 }
 
@@ -79,13 +76,21 @@ pub fn shielder_update(
     shielder_address: &AccountId32,
     upd_op: UpdateOperation,
     user_shielded_data: ShielderUserEnv,
-    merkle_tree: &mut MerkleTree,
+    nullifier: Scalar,
 ) -> Result<ShielderUserEnv> {
-    let merkle_root = merkle_tree.root();
-    let merkle_proof = merkle_tree
-        .gen_proof(user_shielded_data.tree_leaf_id as usize)
-        .unwrap();
-    let nullifier_new = (u128::from(user_shielded_data.nullifier) + 1).into();
+    let merkle_root: Scalar = session.call_with_address(
+        shielder_address.clone(),
+        "notes_merkle_root",
+        NO_ARGS,
+        NO_ENDOWMENT,
+    )??;
+    let merkle_proof_res: Result<[Scalar; 10], ()> = session.call_with_address(
+        shielder_address.clone(),
+        "notes_merkle_path",
+        &[format!("{:?}", user_shielded_data.tree_leaf_id)],
+        NO_ENDOWMENT,
+    )??;
+    let merkle_proof = merkle_proof_res.unwrap();
     let trapdoor_new = 1_u128.into();
 
     let op_pub = upd_op.op_pub;
@@ -96,14 +101,13 @@ pub fn shielder_update(
         .update_account(
             operation,
             trapdoor_new,
-            nullifier_new,
+            nullifier,
             merkle_proof,
             user_shielded_data.tree_leaf_id,
         )
         .unwrap();
-    merkle_tree.add_leaf(note_hash).unwrap();
 
-    session.call_with_address(
+    let new_leaf_id_res: Result<u32, ()> = session.call_with_address(
         shielder_address.clone(),
         "update_note",
         &[
@@ -118,7 +122,7 @@ pub fn shielder_update(
 
     Ok(ShielderUserEnv {
         proof: new_proof,
-        nullifier: nullifier_new,
-        tree_leaf_id: user_shielded_data.tree_leaf_id + 1,
+        nullifier,
+        tree_leaf_id: new_leaf_id_res.unwrap(),
     })
 }

shielder/contract/lib.rs

@@ -8,8 +8,6 @@ mod drink_tests;
 mod errors;
 mod merkle;
 pub mod mocked_zk;
-#[cfg(test)]
-pub mod test_utils;
 mod traits;
 mod types;
 
@@ -19,7 +17,7 @@ pub mod contract {
 
     use crate::{
         errors::ShielderError,
-        merkle::MerkleTree,
+        merkle::{MerkleTree, DEPTH},
         mocked_zk::relations::ZkProof,
         traits::psp22::PSP22,
         types::{Scalar, Set},
@@ -71,10 +69,9 @@ pub mod contract {
             &mut self,
             h_note_new: Scalar,
             proof: ZkProof,
-        ) -> Result<(), ShielderError> {
+        ) -> Result<u32, ShielderError> {
             proof.verify_creation(h_note_new)?;
-            self.notes.add_leaf(h_note_new)?;
-            Ok(())
+            self.notes.add_leaf(h_note_new)
         }
 
         /// Updates existing note
@@ -87,13 +84,13 @@ pub mod contract {
             merkle_root: Scalar,
             nullifier_old: Scalar,
             proof: ZkProof,
-        ) -> Result<(), ShielderError> {
+        ) -> Result<u32, ShielderError> {
             self.notes.is_historical_root(merkle_root)?;
             self.nullify(nullifier_old)?;
             proof.verify_update(op_pub, h_note_new, merkle_root, nullifier_old)?;
-            self.notes.add_leaf(h_note_new)?;
+            let leaf_id = self.notes.add_leaf(h_note_new)?;
             self.process_operation(op_pub)?;
-            Ok(())
+            Ok(leaf_id)
         }
 
         fn process_operation(&mut self, op_pub: OpPub) -> Result<(), ShielderError> {
@@ -123,6 +120,18 @@ pub mod contract {
             Ok(())
         }
 
+        /// Returns merkle root of notes storage
+        #[ink(message)]
+        pub fn notes_merkle_root(&self) -> Scalar {
+            self.notes.root()
+        }
+
+        /// Returns merkle path
+        #[ink(message)]
+        pub fn notes_merkle_path(&self, note_id: u32) -> Result<[Scalar; DEPTH], ShielderError> {
+            self.notes.gen_proof(note_id)
+        }
+
         fn nullify(&mut self, nullifier: Scalar) -> Result<(), ShielderError> {
             self.nullifier_set
                 .insert(nullifier, &())