move mocks to separate modules and add unit tests

shielder/contract/errors.rs

@@ -1,12 +1,13 @@
 use psp22::PSP22Error;
 
-#[derive(scale::Encode, scale::Decode)]
+#[derive(PartialEq, Debug, scale::Encode, scale::Decode)]
 #[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
 pub enum ShielderError {
     PSP22(PSP22Error),
     NullifierIsInSet,
     MerkleTreeVerificationFail,
     MerkleTreeLimitExceeded,
+    MerkleTreeProofGenFail,
     ZkpVerificationFail,
 }
 

shielder/contract/lib.rs

@@ -15,7 +15,7 @@ mod contract {
     use crate::{
         errors::ShielderError,
         merkle::MerkleTree,
-        mocked_zk::{self, ZkProof},
+        mocked_zk::relations::ZkProof,
         types::{Scalar, Set},
     };
 
@@ -51,7 +51,12 @@ mod contract {
         }
 
         #[ink(message)]
-        pub fn add_note(&mut self, h_note_new: Scalar) -> Result<(), ShielderError> {
+        pub fn add_note(
+            &mut self,
+            h_note_new: Scalar,
+            proof: ZkProof,
+        ) -> Result<(), ShielderError> {
+            proof.verify_creation(h_note_new)?;
             self.notes.add_leaf(h_note_new)?;
             Ok(())
         }
@@ -67,7 +72,7 @@ mod contract {
         ) -> Result<(), ShielderError> {
             self.process_operation(op_pub)?;
             self.notes.is_historical_root(merkle_root)?;
-            mocked_zk::verify_update(proof, op_pub, h_note_new, merkle_root, nullifier_old)?;
+            proof.verify_update(op_pub, h_note_new, merkle_root, nullifier_old)?;
             self.nullify(nullifier_old)?;
             self.notes.add_leaf(h_note_new)?;
             Ok(())

shielder/contract/mocked_zk/account.rs

@@ -0,0 +1,57 @@
+use ink::env::hash::{CryptoHash, Sha2x256};
+
+use super::{ops::Operation, traits::Hashable, USDT_TOKEN};
+use crate::{contract::OpPub, types::Scalar};
+
+#[derive(Default, Clone, Copy, scale::Encode, scale::Decode)]
+#[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
+pub struct Account {
+    balance_aleph: Scalar,
+    balance_usdt: Scalar,
+}
+
+impl Hashable for Account {
+    fn hash(&self) -> Scalar {
+        let mut res = [0x0; 32];
+        Sha2x256::hash(
+            [self.balance_aleph.bytes, self.balance_usdt.bytes]
+                .concat()
+                .as_slice(),
+            &mut res,
+        );
+        Scalar { bytes: res }
+    }
+}
+
+impl Account {
+    pub fn new() -> Self {
+        Self {
+            balance_aleph: 0_u128.into(),
+            balance_usdt: 0_u128.into(),
+        }
+    }
+    pub fn update(&self, operation: Operation) -> Self {
+        match operation.op_pub {
+            OpPub::Deposit { amount, token, .. } => {
+                let mut balance_usdt = self.balance_usdt;
+                if token.as_ref() == USDT_TOKEN {
+                    balance_usdt = (u128::from(balance_usdt) + amount).into();
+                }
+                Self {
+                    balance_aleph: self.balance_aleph,
+                    balance_usdt,
+                }
+            }
+            OpPub::Withdraw { amount, token, .. } => {
+                let mut balance_usdt = self.balance_usdt;
+                if token.as_ref() == USDT_TOKEN {
+                    balance_usdt = (u128::from(balance_usdt) - amount).into();
+                }
+                Self {
+                    balance_aleph: self.balance_aleph,
+                    balance_usdt,
+                }
+            }
+        }
+    }
+}

shielder/contract/mocked_zk/mod.rs

@@ -0,0 +1,9 @@
+mod account;
+mod note;
+mod ops;
+pub mod relations;
+#[cfg(test)]
+mod tests;
+mod traits;
+
+const USDT_TOKEN: [u8; 32] = [0x2_u8; 32];

shielder/contract/mocked_zk/note.rs

@@ -0,0 +1,42 @@
+use ink::env::hash::{CryptoHash, Sha2x256};
+
+use super::traits::Hashable;
+use crate::types::Scalar;
+
+#[derive(Clone, Copy, scale::Encode, scale::Decode)]
+#[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
+pub struct Note {
+    id: Scalar,
+    trapdoor: Scalar,
+    nullifier: Scalar,
+    account_hash: Scalar,
+}
+
+impl Note {
+    pub fn new(id: Scalar, trapdoor: Scalar, nullifier: Scalar, account_hash: Scalar) -> Self {
+        Self {
+            id,
+            trapdoor,
+            nullifier,
+            account_hash,
+        }
+    }
+}
+
+impl Hashable for Note {
+    fn hash(&self) -> Scalar {
+        let mut res = [0x0; 32];
+        Sha2x256::hash(
+            [
+                self.id.bytes,
+                self.trapdoor.bytes,
+                self.nullifier.bytes,
+                self.account_hash.bytes,
+            ]
+            .concat()
+            .as_slice(),
+            &mut res,
+        );
+        Scalar { bytes: res }
+    }
+}

shielder/contract/mocked_zk/ops.rs

@@ -0,0 +1,34 @@
+use ink::primitives::AccountId;
+
+use crate::{contract::OpPub, errors::ShielderError};
+
+/// empty private operation
+#[derive(Clone, Copy, scale::Encode, scale::Decode)]
+#[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
+pub struct OpPriv {
+    pub user: AccountId,
+}
+
+#[derive(Clone, Copy, scale::Encode, scale::Decode)]
+#[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
+pub struct Operation {
+    pub op_pub: OpPub,
+}
+
+impl Operation {
+    pub fn combine(op_pub: OpPub, _op_priv: OpPriv) -> Result<Self, ShielderError> {
+        match op_pub {
+            OpPub::Deposit { user, .. } => {
+                if user != _op_priv.user {
+                    return Err(ShielderError::ZkpVerificationFail);
+                }
+            }
+            OpPub::Withdraw { user, .. } => {
+                if user != _op_priv.user {
+                    return Err(ShielderError::ZkpVerificationFail);
+                }
+            }
+        }
+        Ok(Operation { op_pub })
+    }
+}