address comments

relations/Cargo.toml

@@ -22,16 +22,3 @@ halo2-base = { git = "https://github.com/Cardinal-Cryptography/halo2-lib", branc
 #[dev-dependencies]
 # Fork with halo2curves = 0.6 to match the version of halo2-base
 poseidon = { git = "https://github.com/zemse/pse-poseidon" }
-
-#pprof = { version = "0.13", features = ["criterion", "flamegraph"] }
-#criterion = "0.5.1"
-#criterion-macro = "0.4"
-
-#test-case = "3.3.1"
-#test-log = "0.2.14"
-#env_logger = "0.11.1"
-#proptest = "1.4.0"
-
-#[[bench]]
-#name = "bench"
-#harness = false

relations/src/account.rs

@@ -3,11 +3,14 @@ use halo2_base::{gates::GateChip, utils::BigPrimeField, AssignedValue, Context};
 use crate::operation::{CircuitOperation, Operation};
 
 pub trait Account<F: BigPrimeField> {
+    type CircuitAccount: CircuitAccount<F>;
     type Op: Operation<F>;
 
     fn update(&self, op: Self::Op) -> Self;
 
     fn to_array(&self) -> Vec<F>;
+
+    fn load(&self, ctx: &mut Context<F>) -> Self::CircuitAccount;
 }
 
 pub trait CircuitAccount<F: BigPrimeField> {

relations/src/lib.rs

@@ -4,9 +4,6 @@ pub mod note;
 pub mod operation;
 pub mod relations;
 
-//#[cfg(test)]
-//pub mod tests;
-
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum Token {
     AZERO,

relations/src/merkle_proof.rs

@@ -59,32 +59,3 @@ impl<F: BigPrimeField, const TREE_HEIGHT: usize> CircuitMerkleProof<F, TREE_HEIG
         gate.assert_is_const(ctx, &eq, &F::ONE);
     }
 }
-
-#[cfg(test)]
-pub mod tests {
-    use poseidon::Poseidon;
-
-    use super::*;
-    use crate::poseidon_consts::{R_F, R_P};
-
-    impl<F: BigPrimeField, const TREE_HEIGHT: usize> MerkleProof<F, TREE_HEIGHT> {
-        pub fn verify(&self, root: F, leaf: F) -> bool {
-            let mut current_note = leaf;
-
-            for i in 0..TREE_HEIGHT {
-                let sibling = self.path[i];
-
-                let mut poseidon: Poseidon<F, T_WIDTH, RATE> = Poseidon::new(R_F, R_P);
-
-                if !self.path_shape[i] {
-                    poseidon.update(&[sibling, current_note]);
-                } else {
-                    poseidon.update(&[current_note, sibling]);
-                }
-                current_note = poseidon.squeeze();
-            }
-
-            current_note.eq(&root)
-        }
-    }
-}

relations/src/operation.rs

@@ -1,26 +1,14 @@
-use halo2_base::{utils::BigPrimeField, AssignedValue};
+use halo2_base::{utils::BigPrimeField, AssignedValue, Context};
 
-/// A trait representing an operation.
 pub trait Operation<F>
 where
     Self: Sized,
     F: BigPrimeField,
 {
-    /// The private data associated with the operation.
-    type OpPriv;
-    /// The public data associated with the operation.
-    type OpPub;
+    type CircuitOperation;
+    type OpPriv: Into<Vec<F>>;
+    type OpPub: Into<Vec<F>>;
 
-    /// Combines the private and public data to create an instance of the operation.
-    ///
-    /// # Arguments
-    ///
-    /// * `op_priv` - The private data of the operation.
-    /// * `op_pub` - The public data of the operation.
-    ///
-    /// # Returns
-    ///
-    /// An `Option` containing the combined operation, or `None` if the combination fails.
     fn combine(op_priv: Self::OpPriv, op_pub: Self::OpPub) -> Option<Self>;
 }
 
@@ -29,20 +17,8 @@ where
     Self: Sized,
     F: BigPrimeField,
 {
-    /// The private data associated with the operation.
-    type OpPriv;
-    /// The public data associated with the operation.
-    type OpPub: Into<Vec<AssignedValue<F>>> + Clone;
+    type OpPriv: From<Vec<AssignedValue<F>>>;
+    type OpPub: From<Vec<AssignedValue<F>>> + Into<Vec<AssignedValue<F>>> + Clone;
 
-    /// Combines the private and public data to create an instance of the operation.
-    ///
-    /// # Arguments
-    ///
-    /// * `op_priv` - The private data of the operation.
-    /// * `op_pub` - The public data of the operation.
-    ///
-    /// # Returns
-    ///
-    /// An `Option` containing the combined operation, or `None` if the combination fails.
     fn combine(op_priv: Self::OpPriv, op_pub: Self::OpPub) -> Option<Self>;
 }

relations/src/relations/update_account.rs

@@ -65,27 +65,31 @@ pub fn verify_account_circuit<F, A>(
 }
 
 #[allow(dead_code)]
-pub fn update_account_circuit<F, A>(
-    ctx: &mut Context<F>,
-    input: UpdateAccountInput<F, A>,
-    make_public: &mut Vec<AssignedValue<F>>,
-) where
+pub fn update_account_circuit<F, A>(ctx: &mut Context<F>, input: UpdateAccountInput<F, A>)
+where
     F: BigPrimeField,
     A: CircuitAccount<F>,
 {
-    let old_account_hash = input.old_account_hash;
-    let new_account_hash = input.new_account_hash;
-
-    make_public.extend([old_account_hash, new_account_hash]);
-
     let gate = GateChip::<F>::default();
     let mut poseidon =
         PoseidonHasher::<F, T_WIDTH, RATE>::new(OptimizedPoseidonSpec::new::<R_F, R_P, 0>());
     poseidon.initialize_consts(ctx, &gate);
 
     let old_account = input.old_account;
-    verify_account_circuit(ctx, &gate, &mut poseidon, &old_account, old_account_hash);
+    verify_account_circuit(
+        ctx,
+        &gate,
+        &mut poseidon,
+        &old_account,
+        input.old_account_hash,
+    );
 
     let new_account = old_account.update(input.operation, ctx, &gate);
-    verify_account_circuit(ctx, &gate, &mut poseidon, &new_account, new_account_hash);
+    verify_account_circuit(
+        ctx,
+        &gate,
+        &mut poseidon,
+        &new_account,
+        input.new_account_hash,
+    );
 }

relations/src/relations/update_note.rs

@@ -7,72 +7,81 @@ use halo2_base::{
 
 use super::update_account::{update_account_circuit, UpdateAccountInput};
 use crate::{
-    account::CircuitAccount,
-    merkle_proof::CircuitMerkleProof,
-    note::CircuitNote,
-    operation::CircuitOperation,
+    account::{Account, CircuitAccount},
+    merkle_proof::{CircuitMerkleProof, MerkleProof},
+    note::{CircuitNote, Note},
+    operation::{CircuitOperation, Operation},
     poseidon_consts::{RATE, R_F, R_P, T_WIDTH},
 };
 
+type OpFor<A, F> = <<A as Account<F>>::CircuitAccount as CircuitAccount<F>>::Op;
+
 pub struct UpdateNoteInput<F, A, const TREE_HEIGHT: usize>
 where
     F: BigPrimeField,
-    A: CircuitAccount<F>,
+    A: Account<F>,
 {
     //public inputs
-    pub op_pub: <A::Op as CircuitOperation<F>>::OpPub,
+    pub op_pub: <OpFor<A, F> as CircuitOperation<F>>::OpPub,
     pub new_note_hash: AssignedValue<F>,
     pub merkle_root: AssignedValue<F>,
-    pub old_nullifier: AssignedValue<F>,
+    //old_note.nullifier is also a public input
 
     //witnesses
     pub new_note: CircuitNote<F>,
     pub old_note: CircuitNote<F>,
-    pub new_trapdoor: AssignedValue<F>,
-    pub old_trapdoor: AssignedValue<F>,
-    pub new_nullifier: AssignedValue<F>,
     pub merkle_proof: CircuitMerkleProof<F, TREE_HEIGHT>,
-    pub op_priv: <A::Op as CircuitOperation<F>>::OpPriv,
-    pub note_id: AssignedValue<F>,
+    pub op_priv: <OpFor<A, F> as CircuitOperation<F>>::OpPriv,
 
-    pub old_account: A,
+    pub old_account: <A as Account<F>>::CircuitAccount,
 }
 
 //helper functions
 #[allow(clippy::too_many_arguments)]
 impl<F, A, const TREE_HEIGHT: usize> UpdateNoteInput<F, A, TREE_HEIGHT>
 where
     F: BigPrimeField,
-    A: CircuitAccount<F>,
+    A: Account<F>,
 {
     pub fn new(
-        op_pub: <A::Op as CircuitOperation<F>>::OpPub,
-        new_note_hash: AssignedValue<F>,
-        merkle_root: AssignedValue<F>,
-        old_nullifier: AssignedValue<F>,
-        new_note: CircuitNote<F>,
-        old_note: CircuitNote<F>,
-        new_trapdoor: AssignedValue<F>,
-        old_trapdoor: AssignedValue<F>,
-        new_nullifier: AssignedValue<F>,
-        merkle_proof: CircuitMerkleProof<F, TREE_HEIGHT>,
-        op_priv: <A::Op as CircuitOperation<F>>::OpPriv,
-        note_id: AssignedValue<F>,
+        ctx: &mut Context<F>,
+        op_pub: <A::Op as Operation<F>>::OpPub,
+        new_note_hash: F,
+        merkle_root: F,
+        new_note: Note<F>,
+        old_note: Note<F>,
+        merkle_proof: MerkleProof<F, TREE_HEIGHT>,
+        op_priv: <A::Op as Operation<F>>::OpPriv,
         old_account: A,
     ) -> Self {
+        let op_pub = op_pub
+            .into()
+            .iter()
+            .map(|x| ctx.load_witness(*x))
+            .collect::<Vec<AssignedValue<F>>>()
+            .into();
+        let new_note_hash = ctx.load_witness(new_note_hash);
+        let merkle_root = ctx.load_witness(merkle_root);
+        let new_note = new_note.load(ctx);
+        let old_note = old_note.load(ctx);
+        let merkle_proof = merkle_proof.load(ctx);
+        let op_priv = op_priv
+            .into()
+            .iter()
+            .map(|x| ctx.load_witness(*x))
+            .collect::<Vec<AssignedValue<F>>>()
+            .into();
+
+        let old_account = old_account.load(ctx);
+
         Self {
             op_pub,
             new_note_hash,
             merkle_root,
-            old_nullifier,
             new_note,
             old_note,
-            new_trapdoor,
-            old_trapdoor,
-            new_nullifier,
             merkle_proof,
             op_priv,
-            note_id,
             old_account,
         }
     }
@@ -99,45 +108,26 @@ pub fn update_note_circuit<F, A, const TREE_HEIGHT: usize>(
     make_public: &mut Vec<AssignedValue<F>>,
 ) where
     F: BigPrimeField,
-    A: CircuitAccount<F>,
+    A: Account<F>,
 {
+    let gate = GateChip::<F>::default();
+    let mut poseidon =
+        PoseidonHasher::<F, T_WIDTH, RATE>::new(OptimizedPoseidonSpec::new::<R_F, R_P, 0>());
+    poseidon.initialize_consts(ctx, &gate);
+
     let op_pub = input.op_pub;
 
     make_public.extend(op_pub.clone().into());
 
-    let outer_new_note_hash = input.new_note_hash;
+    let new_note_hash = input.new_note_hash;
     let merkle_root = input.merkle_root;
-    let old_nullifier = input.old_nullifier;
-
-    make_public.extend([outer_new_note_hash, merkle_root, old_nullifier]);
+    let old_nullifier = input.old_note.nullifier;
 
-    //TODO: Check how halo2 works with unused inputs
-    let _new_trapdoor = input.new_trapdoor;
-    let _new_nullifier = input.new_nullifier;
+    make_public.extend([new_note_hash, merkle_root, old_nullifier]);
 
-    let gate = GateChip::<F>::default();
-
-    let mut poseidon =
-        PoseidonHasher::<F, T_WIDTH, RATE>::new(OptimizedPoseidonSpec::new::<R_F, R_P, 0>());
-    poseidon.initialize_consts(ctx, &gate);
-
-    verify_note_circuit(
-        ctx,
-        &gate,
-        &mut poseidon,
-        &input.new_note,
-        outer_new_note_hash,
-    );
+    verify_note_circuit(ctx, &gate, &mut poseidon, &input.new_note, new_note_hash);
 
     let old_note_hash = poseidon.hash_fix_len_array(ctx, &gate, &input.old_note.to_array());
-    let old_account_hash = poseidon.hash_fix_len_array(ctx, &gate, &input.old_account.to_array());
-    let outer_old_note = CircuitNote {
-        note_id: input.note_id,
-        trapdoor: input.old_trapdoor,
-        nullifier: input.old_nullifier,
-        account_hash: old_account_hash,
-    };
-    verify_note_circuit(ctx, &gate, &mut poseidon, &outer_old_note, old_note_hash);
 
     let merkle_proof = input.merkle_proof;
 
@@ -154,5 +144,5 @@ pub fn update_note_circuit<F, A, const TREE_HEIGHT: usize>(
         input.old_account,
     );
 
-    update_account_circuit(ctx, update_account_input, make_public);
+    update_account_circuit(ctx, update_account_input);
 }