From 35478ec42692827f2dcfb5c39d13a1bd3582094b Mon Sep 17 00:00:00 2001 From: Enrico Bottazzi <85900164+enricobottazzi@users.noreply.github.com> Date: Mon, 6 Nov 2023 16:29:41 +0100 Subject: [PATCH] feat: added method implementation to `Tree` trait and moved outside of utils --- zk_prover/examples/gen_inclusion_verifier.rs | 2 +- .../examples/nova_incremental_verifier.rs | 2 +- zk_prover/src/circuits/tests.rs | 4 +- zk_prover/src/merkle_sum_tree/mst.rs | 19 +--- zk_prover/src/merkle_sum_tree/tree.rs | 87 +++++++++++++++++-- .../src/merkle_sum_tree/utils/create_proof.rs | 44 ---------- .../src/merkle_sum_tree/utils/index_of.rs | 17 ---- zk_prover/src/merkle_sum_tree/utils/mod.rs | 6 -- .../utils/proof_verification.rs | 32 ------- 9 files changed, 86 insertions(+), 127 deletions(-) delete mode 100644 zk_prover/src/merkle_sum_tree/utils/create_proof.rs delete mode 100644 zk_prover/src/merkle_sum_tree/utils/index_of.rs delete mode 100644 zk_prover/src/merkle_sum_tree/utils/proof_verification.rs diff --git a/zk_prover/examples/gen_inclusion_verifier.rs b/zk_prover/examples/gen_inclusion_verifier.rs index 815ef7ad..312aace1 100644 --- a/zk_prover/examples/gen_inclusion_verifier.rs +++ b/zk_prover/examples/gen_inclusion_verifier.rs @@ -56,7 +56,7 @@ fn main() { let user_entry = merkle_sum_tree.get_entry(user_index); // Generate the circuit with the actual inputs - let mut circuit = + let circuit = MstInclusionCircuit::::init(merkle_proof, user_entry.clone()); let instances = circuit.instances(); diff --git a/zk_prover/examples/nova_incremental_verifier.rs b/zk_prover/examples/nova_incremental_verifier.rs index 7ab36e43..74ff0a2d 100644 --- a/zk_prover/examples/nova_incremental_verifier.rs +++ b/zk_prover/examples/nova_incremental_verifier.rs @@ -235,7 +235,7 @@ fn main() { run_test(circuit_filepath.clone(), witness_gen_filepath); } -use num_traits::{FromPrimitive, Num}; +use num_traits::Num; use poseidon_rs::{Fr, Poseidon}; // Note that we cannot reuse the MerkleSumTree implementation from zk_prover because it is not compatible with circom's Poseidon Hasher diff --git a/zk_prover/src/circuits/tests.rs b/zk_prover/src/circuits/tests.rs index 18a8db31..f4429d6d 100644 --- a/zk_prover/src/circuits/tests.rs +++ b/zk_prover/src/circuits/tests.rs @@ -257,7 +257,7 @@ mod test { let user_entry = merkle_sum_tree.get_entry(user_index); // Only now we can instantiate the circuit with the actual inputs - let mut circuit = MstInclusionCircuit::::init( + let circuit = MstInclusionCircuit::::init( merkle_proof, user_entry.clone(), ); @@ -741,7 +741,7 @@ mod test { let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap(); let user_entry = merkle_sum_tree.get_entry(user_index); - let mut circuit = MstInclusionCircuit::::init( + let circuit = MstInclusionCircuit::::init( merkle_proof, user_entry.clone(), ); diff --git a/zk_prover/src/merkle_sum_tree/mst.rs b/zk_prover/src/merkle_sum_tree/mst.rs index 42ee4003..b84fe221 100644 --- a/zk_prover/src/merkle_sum_tree/mst.rs +++ b/zk_prover/src/merkle_sum_tree/mst.rs @@ -1,7 +1,7 @@ use crate::merkle_sum_tree::utils::{ - build_merkle_tree_from_leaves, compute_leaves, create_proof, index_of, parse_csv_to_entries, + build_merkle_tree_from_leaves, compute_leaves, parse_csv_to_entries, }; -use crate::merkle_sum_tree::{Entry, MerkleProof, Node, Tree}; +use crate::merkle_sum_tree::{Entry, Node, Tree}; use num_bigint::BigUint; /// Merkle Sum Tree Data Structure. @@ -43,14 +43,6 @@ impl Tree fn nodes(&self) -> &[Vec>] { &self.nodes } - - /// Generates a MerkleProof for the user with the given index. No mini tree index is required for a MerkleSumTree. - fn generate_proof( - &self, - user_index: usize, - ) -> Result, &'static str> { - create_proof(user_index, self.depth, &self.nodes, &self.root) - } } impl MerkleSumTree { @@ -170,13 +162,6 @@ impl MerkleSumTree Option - where - [usize; N_ASSETS + 1]: Sized, - { - index_of(username, balances, &self.nodes) - } /// Returns the index of the leaf with the matching username pub fn index_of_username(&self, username: &str) -> Result> diff --git a/zk_prover/src/merkle_sum_tree/tree.rs b/zk_prover/src/merkle_sum_tree/tree.rs index da229cf3..079ba429 100644 --- a/zk_prover/src/merkle_sum_tree/tree.rs +++ b/zk_prover/src/merkle_sum_tree/tree.rs @@ -1,5 +1,6 @@ -use crate::merkle_sum_tree::utils::verify_proof; -use crate::merkle_sum_tree::{MerkleProof, Node}; +use crate::merkle_sum_tree::{Entry, MerkleProof, Node}; +use halo2_proofs::halo2curves::bn256::Fr as Fp; +use num_bigint::BigUint; /// A trait representing the basic operations for a Merkle-Sum-like Tree. pub trait Tree { @@ -16,10 +17,46 @@ pub trait Tree { fn nodes(&self) -> &[Vec>]; /// Generates a MerkleProof for the user with the given index. - fn generate_proof( - &self, - user_index: usize, - ) -> Result, &'static str>; + fn generate_proof(&self, index: usize) -> Result, &'static str> { + let nodes = self.nodes(); + let depth = *self.depth(); + let root = self.root(); + + if index >= nodes[0].len() { + return Err("The leaf does not exist in this tree"); + } + + let mut sibling_hashes = vec![Fp::from(0); depth]; + let mut sibling_sums = vec![[Fp::from(0); N_ASSETS]; depth]; + let mut path_indices = vec![Fp::from(0); depth]; + let mut current_index = index; + + let leaf = &nodes[0][index]; + + for level in 0..depth { + let position = current_index % 2; + let level_start_index = current_index - position; + let level_end_index = level_start_index + 2; + + path_indices[level] = Fp::from(position as u64); + + for i in level_start_index..level_end_index { + if i != current_index { + sibling_hashes[level] = nodes[level][i].hash; + sibling_sums[level] = nodes[level][i].balances; + } + } + current_index /= 2; + } + + Ok(MerkleProof { + leaf: leaf.clone(), + root_hash: root.hash, + sibling_hashes, + sibling_sums, + path_indices, + }) + } /// Verifies a MerkleProof. fn verify_proof(&self, proof: &MerkleProof) -> bool @@ -27,6 +64,42 @@ pub trait Tree { [usize; N_ASSETS + 1]: Sized, [usize; 2 * (1 + N_ASSETS)]: Sized, { - verify_proof(proof) + let mut node = proof.leaf.clone(); + + let mut balances = proof.leaf.balances; + + for i in 0..proof.sibling_hashes.len() { + let sibling_node = Node { + hash: proof.sibling_hashes[i], + balances: proof.sibling_sums[i], + }; + + if proof.path_indices[i] == 0.into() { + node = Node::middle(&node, &sibling_node); + } else { + node = Node::middle(&sibling_node, &node); + } + + for (balance, sibling_balance) in balances.iter_mut().zip(sibling_node.balances.iter()) + { + *balance += sibling_balance; + } + } + + proof.root_hash == node.hash && balances == node.balances + } + + /// Returns the index of the user with the given username and balances in the tree + fn index_of(&self, username: &str, balances: [BigUint; N_ASSETS]) -> Option + where + [usize; N_ASSETS + 1]: Sized, + { + let entry: Entry = Entry::new(username.to_string(), balances).unwrap(); + let leaf = entry.compute_leaf(); + let leaf_hash = leaf.hash; + + self.nodes()[0] + .iter() + .position(|node| node.hash == leaf_hash) } } diff --git a/zk_prover/src/merkle_sum_tree/utils/create_proof.rs b/zk_prover/src/merkle_sum_tree/utils/create_proof.rs deleted file mode 100644 index 9d97c0b2..00000000 --- a/zk_prover/src/merkle_sum_tree/utils/create_proof.rs +++ /dev/null @@ -1,44 +0,0 @@ -use crate::merkle_sum_tree::{MerkleProof, Node}; -use halo2_proofs::halo2curves::bn256::Fr as Fp; - -pub fn create_proof( - index: usize, - depth: usize, - nodes: &[Vec>], - root: &Node, -) -> Result, &'static str> { - if index >= nodes[0].len() { - return Err("The leaf does not exist in this tree"); - } - - let mut sibling_hashes = vec![Fp::from(0); depth]; - let mut sibling_sums = vec![[Fp::from(0); N_ASSETS]; depth]; - let mut path_indices = vec![Fp::from(0); depth]; - let mut current_index = index; - - let leaf = &nodes[0][index]; - - for level in 0..depth { - let position = current_index % 2; - let level_start_index = current_index - position; - let level_end_index = level_start_index + 2; - - path_indices[level] = Fp::from(position as u64); - - for i in level_start_index..level_end_index { - if i != current_index { - sibling_hashes[level] = nodes[level][i].hash; - sibling_sums[level] = nodes[level][i].balances; - } - } - current_index /= 2; - } - - Ok(MerkleProof { - leaf: leaf.clone(), - root_hash: root.hash, - sibling_hashes, - sibling_sums, - path_indices, - }) -} diff --git a/zk_prover/src/merkle_sum_tree/utils/index_of.rs b/zk_prover/src/merkle_sum_tree/utils/index_of.rs deleted file mode 100644 index 5ef7f191..00000000 --- a/zk_prover/src/merkle_sum_tree/utils/index_of.rs +++ /dev/null @@ -1,17 +0,0 @@ -use crate::merkle_sum_tree::{Entry, Node}; -use num_bigint::BigUint; - -pub fn index_of( - username: &str, - balances: [BigUint; N_ASSETS], - nodes: &[Vec>], -) -> Option -where - [usize; N_ASSETS + 1]: Sized, -{ - let entry: Entry = Entry::new(username.to_string(), balances).unwrap(); - let leaf = entry.compute_leaf(); - let leaf_hash = leaf.hash; - - nodes[0].iter().position(|node| node.hash == leaf_hash) -} diff --git a/zk_prover/src/merkle_sum_tree/utils/mod.rs b/zk_prover/src/merkle_sum_tree/utils/mod.rs index ccf3cf13..09d13579 100644 --- a/zk_prover/src/merkle_sum_tree/utils/mod.rs +++ b/zk_prover/src/merkle_sum_tree/utils/mod.rs @@ -1,17 +1,11 @@ mod build_tree; -mod create_proof; mod csv_parser; mod generate_leaf_hash; mod hash; -mod index_of; mod operation_helpers; -mod proof_verification; pub use build_tree::{build_merkle_tree_from_leaves, compute_leaves}; -pub use create_proof::create_proof; pub use csv_parser::parse_csv_to_entries; pub use generate_leaf_hash::generate_leaf_hash; pub use hash::{poseidon_entry, poseidon_node}; -pub use index_of::index_of; pub use operation_helpers::*; -pub use proof_verification::verify_proof; diff --git a/zk_prover/src/merkle_sum_tree/utils/proof_verification.rs b/zk_prover/src/merkle_sum_tree/utils/proof_verification.rs deleted file mode 100644 index 0098179d..00000000 --- a/zk_prover/src/merkle_sum_tree/utils/proof_verification.rs +++ /dev/null @@ -1,32 +0,0 @@ -use crate::merkle_sum_tree::{MerkleProof, Node}; - -pub fn verify_proof( - proof: &MerkleProof, -) -> bool -where - [usize; N_ASSETS + 1]: Sized, - [usize; 2 * (1 + N_ASSETS)]: Sized, -{ - let mut node = proof.leaf.clone(); - - let mut balances = proof.leaf.balances; - - for i in 0..proof.sibling_hashes.len() { - let sibling_node = Node { - hash: proof.sibling_hashes[i], - balances: proof.sibling_sums[i], - }; - - if proof.path_indices[i] == 0.into() { - node = Node::middle(&node, &sibling_node); - } else { - node = Node::middle(&sibling_node, &node); - } - - for (balance, sibling_balance) in balances.iter_mut().zip(sibling_node.balances.iter()) { - *balance += sibling_balance; - } - } - - proof.root_hash == node.hash && balances == node.balances -}