feat: Tree trait

summa-dev · Nov 6, 2023 · 022dde6 · 022dde6
1 parent 400fc49
commit 022dde6
Show file tree

Hide file tree

Showing 13 changed files with 245 additions and 110 deletions.
diff --git a/zk_prover/benches/full_solvency_flow.rs b/zk_prover/benches/full_solvency_flow.rs
@@ -11,7 +11,7 @@ use summa_solvency::{
         solvency::SolvencyCircuit,
         utils::{full_prover, full_verifier, generate_setup_artifacts},
     },
-    merkle_sum_tree::MerkleSumTree,
+    merkle_sum_tree::{MerkleSumTree, Tree},
 };
 
 const SAMPLE_SIZE: usize = 10;
@@ -101,7 +101,7 @@ fn generate_zk_proof_mst_inclusion_circuit(_c: &mut Criterion) {
 
     let empty_circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init_empty();
 
-    let (params, pk, vk) = generate_setup_artifacts(13, None, empty_circuit).unwrap();
+    let (params, pk, _) = generate_setup_artifacts(13, None, empty_circuit).unwrap();
 
     let csv_file = format!(
         "benches/csv/{}/{}_entry_2_{}.csv",
@@ -111,7 +111,14 @@ fn generate_zk_proof_mst_inclusion_circuit(_c: &mut Criterion) {
     let merkle_sum_tree = MerkleSumTree::<N_ASSETS, N_BYTES>::new(&csv_file).unwrap();
 
     // Only now we can instantiate the circuit with the actual inputs
-    let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+
+    let user_index = 0;
+
+    let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+    let user_entry = merkle_sum_tree.get_entry(user_index);
+
+    let circuit =
+        MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_proof, user_entry.clone());
 
     let bench_name = format!(
         "generate zk proof - tree of 2 power of {} entries with {} assets mst inclusion circuit",
@@ -139,7 +146,14 @@ fn verify_zk_proof_mst_inclusion_circuit(_c: &mut Criterion) {
     let merkle_sum_tree = MerkleSumTree::<N_ASSETS, N_BYTES>::new(&csv_file).unwrap();
 
     // Only now we can instantiate the circuit with the actual inputs
-    let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+
+    let user_index = 0;
+
+    let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+    let user_entry = merkle_sum_tree.get_entry(user_index);
+
+    let circuit =
+        MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_proof, user_entry.clone());
 
     let proof = full_prover(&params, &pk, circuit.clone(), circuit.instances());
 
@@ -197,7 +211,7 @@ fn generate_zk_proof_solvency_circuit(_c: &mut Criterion) {
 
     let empty_circuit = SolvencyCircuit::<N_ASSETS, N_BYTES>::init_empty();
 
-    let (params, pk, vk) = generate_setup_artifacts(11, None, empty_circuit).unwrap();
+    let (params, pk, _) = generate_setup_artifacts(11, None, empty_circuit).unwrap();
 
     let csv_file = format!(
         "benches/csv/{}/{}_entry_2_{}.csv",

diff --git a/zk_prover/examples/gen_inclusion_verifier.rs b/zk_prover/examples/gen_inclusion_verifier.rs
@@ -14,7 +14,7 @@ use summa_solvency::{
             gen_proof_solidity_calldata, generate_setup_artifacts, write_verifier_sol_from_yul,
         },
     },
-    merkle_sum_tree::MerkleSumTree,
+    merkle_sum_tree::{MerkleSumTree, Tree},
 };
 
 const LEVELS: usize = 4;
@@ -49,7 +49,15 @@ fn main() {
 
     // In order to generate a proof for testing purpose we create the circuit using the init() method
     // which takes as input the merkle sum tree and the index of the leaf we are generating the proof for.
-    let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+
+    let user_index = 0;
+
+    let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+    let user_entry = merkle_sum_tree.get_entry(user_index);
+
+    // Generate the circuit with the actual inputs
+    let mut circuit =
+        MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_proof, user_entry.clone());
 
     let instances = circuit.instances();
 

diff --git a/zk_prover/src/circuits/merkle_sum_tree.rs b/zk_prover/src/circuits/merkle_sum_tree.rs
@@ -3,7 +3,7 @@ use crate::chips::poseidon::hash::{PoseidonChip, PoseidonConfig};
 use crate::chips::poseidon::poseidon_spec::PoseidonSpec;
 use crate::chips::range::range_check::{RangeCheckChip, RangeCheckConfig};
 use crate::circuits::traits::CircuitBase;
-use crate::merkle_sum_tree::{big_uint_to_fp, Entry, MerkleSumTree};
+use crate::merkle_sum_tree::{big_uint_to_fp, Entry, MerkleProof};
 use halo2_proofs::circuit::{AssignedCell, Layouter, SimpleFloorPlanner};
 use halo2_proofs::halo2curves::bn256::Fr as Fp;
 use halo2_proofs::plonk::{
@@ -68,23 +68,24 @@ impl<const LEVELS: usize, const N_ASSETS: usize, const N_BYTES: usize>
         }
     }
 
-    /// Initializes the circuit with the merkle sum tree and the index of the user of which the inclusion is to be verified.
-    pub fn init(merkle_sum_tree: MerkleSumTree<N_ASSETS, N_BYTES>, user_index: usize) -> Self
+    /// Initializes the circuit with the merkle proof and the entry of the user of which the inclusion is to be verified.
+    pub fn init(merkle_proof: MerkleProof<N_ASSETS, N_BYTES>, entry: Entry<N_ASSETS>) -> Self
     where
         [usize; N_ASSETS + 1]:,
     {
-        let proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+        assert_eq!(merkle_proof.path_indices.len(), LEVELS);
+        assert_eq!(merkle_proof.sibling_hashes.len(), LEVELS);
+        assert_eq!(merkle_proof.sibling_sums.len(), LEVELS);
 
-        assert_eq!(proof.path_indices.len(), LEVELS);
-        assert_eq!(proof.sibling_hashes.len(), LEVELS);
-        assert_eq!(proof.sibling_sums.len(), LEVELS);
+        // assert that the entry leaf hash matches the leaf hash in the merkle proof
+        assert_eq!(merkle_proof.leaf.hash, entry.compute_leaf().hash);
 
         Self {
-            entry: proof.entry,
-            path_element_hashes: proof.sibling_hashes,
-            path_element_balances: proof.sibling_sums,
-            path_indices: proof.path_indices,
-            root_hash: proof.root_hash,
+            entry,
+            path_element_hashes: merkle_proof.sibling_hashes,
+            path_element_balances: merkle_proof.sibling_sums,
+            path_indices: merkle_proof.path_indices,
+            root_hash: merkle_proof.root_hash,
         }
     }
 }
@@ -222,8 +223,7 @@ where
                 config.poseidon_middle_config.clone(),
             );
 
-        let range_check_chip =
-            RangeCheckChip::<N_BYTES>::construct(config.range_check_config.clone());
+        let range_check_chip = RangeCheckChip::<N_BYTES>::construct(config.range_check_config);
 
         // Assign the entry username
         let username = self.assign_value_to_witness(

diff --git a/zk_prover/src/circuits/solvency.rs b/zk_prover/src/circuits/solvency.rs
@@ -3,7 +3,7 @@ use crate::chips::merkle_sum_tree::{MerkleSumTreeChip, MerkleSumTreeConfig};
 use crate::chips::poseidon::hash::{PoseidonChip, PoseidonConfig};
 use crate::chips::poseidon::poseidon_spec::PoseidonSpec;
 use crate::circuits::traits::CircuitBase;
-use crate::merkle_sum_tree::MerkleSumTree;
+use crate::merkle_sum_tree::{MerkleSumTree, Tree};
 use halo2_proofs::circuit::{AssignedCell, Layouter, SimpleFloorPlanner, Value};
 use halo2_proofs::halo2curves::bn256::Fr as Fp;
 use halo2_proofs::plonk::{

diff --git a/zk_prover/src/circuits/tests.rs b/zk_prover/src/circuits/tests.rs
@@ -1,7 +1,7 @@
 #[cfg(test)]
 mod test {
 
-    use crate::merkle_sum_tree::MerkleSumTree;
+    use crate::merkle_sum_tree::{MerkleSumTree, Tree};
     use crate::{
         circuits::{
             merkle_sum_tree::MstInclusionCircuit,
@@ -33,9 +33,13 @@ mod test {
                 .unwrap();
 
         for user_index in 0..16 {
+            // get proof for entry ˆuser_indexˆ
+            let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+            let user_entry = merkle_sum_tree.get_entry(user_index);
+
             let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(
-                merkle_sum_tree.clone(),
-                user_index,
+                merkle_proof,
+                user_entry.clone(),
             );
 
             let valid_prover = MockProver::run(K, &circuit, circuit.instances()).unwrap();
@@ -62,8 +66,16 @@ mod test {
             MerkleSumTree::<N_ASSETS, N_BYTES>::new("src/merkle_sum_tree/csv/entry_16.csv")
                 .unwrap();
 
+        let user_index = 0;
+
+        let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+        let user_entry = merkle_sum_tree.get_entry(user_index);
+
         // Only now we can instantiate the circuit with the actual inputs
-        let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+        let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(
+            merkle_proof,
+            user_entry.clone(),
+        );
 
         // Generate the proof
         let proof = full_prover(&params, &pk, circuit.clone(), circuit.instances());
@@ -104,8 +116,15 @@ mod test {
         let merkle_sum_tree =
             MerkleSumTree::<N_ASSETS, N_BYTES>::new("src/merkle_sum_tree/csv/entry_16.csv")
                 .unwrap();
+        let user_index = 0;
 
-        let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+        let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+        let user_entry = merkle_sum_tree.get_entry(user_index);
+
+        let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(
+            merkle_proof,
+            user_entry.clone(),
+        );
 
         let mut instances = circuit.instances();
         let invalid_root_hash = Fp::from(1000u64);
@@ -142,8 +161,16 @@ mod test {
             MerkleSumTree::<N_ASSETS, N_BYTES>::new("src/merkle_sum_tree/csv/entry_16.csv")
                 .unwrap();
 
+        let user_index = 0;
+
+        let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+        let user_entry = merkle_sum_tree.get_entry(user_index);
+
         // Only now we can instantiate the circuit with the actual inputs
-        let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+        let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(
+            merkle_proof,
+            user_entry.clone(),
+        );
 
         let invalid_root_hash = Fp::from(1000u64);
 
@@ -166,8 +193,16 @@ mod test {
             MerkleSumTree::<N_ASSETS, N_BYTES>::new("src/merkle_sum_tree/csv/entry_16.csv")
                 .unwrap();
 
-        let mut circuit =
-            MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+        let user_index = 0;
+
+        let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+        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::<LEVELS, N_ASSETS, N_BYTES>::init(
+            merkle_proof,
+            user_entry.clone(),
+        );
 
         let instances = circuit.instances();
 
@@ -216,7 +251,16 @@ mod test {
             MerkleSumTree::<N_ASSETS, N_BYTES>::new("src/merkle_sum_tree/csv/entry_16.csv")
                 .unwrap();
 
-        let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+        let user_index = 0;
+
+        let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+        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::<LEVELS, N_ASSETS, N_BYTES>::init(
+            merkle_proof,
+            user_entry.clone(),
+        );
 
         let mut instances = circuit.instances();
         let invalid_leaf_hash = Fp::from(1000u64);
@@ -249,8 +293,16 @@ mod test {
             MerkleSumTree::<N_ASSETS, N_BYTES>::new("src/merkle_sum_tree/csv/entry_16.csv")
                 .unwrap();
 
-        let mut circuit =
-            MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+        let user_index = 0;
+
+        let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+        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::<LEVELS, N_ASSETS, N_BYTES>::init(
+            merkle_proof,
+            user_entry.clone(),
+        );
 
         let instances = circuit.instances();
 
@@ -366,8 +418,16 @@ mod test {
             MerkleSumTree::<N_ASSETS, N_BYTES>::new("src/merkle_sum_tree/csv/entry_16.csv")
                 .unwrap();
 
-        let mut circuit =
-            MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+        let user_index = 0;
+
+        let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+        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::<LEVELS, N_ASSETS, N_BYTES>::init(
+            merkle_proof,
+            user_entry.clone(),
+        );
 
         let instances = circuit.instances();
 
@@ -401,8 +461,16 @@ mod test {
             MerkleSumTree::<N_ASSETS, N_BYTES>::new("src/merkle_sum_tree/csv/entry_16.csv")
                 .unwrap();
 
-        let mut circuit =
-            MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree.clone(), 0);
+        let user_index = 0;
+
+        let merkle_proof = merkle_sum_tree.generate_proof(user_index).unwrap();
+        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::<LEVELS, N_ASSETS, N_BYTES>::init(
+            merkle_proof,
+            user_entry.clone(),
+        );
 
         let balance = BigUint::from(2u64).pow(N_BYTES as u32 * 8) - BigUint::from(1u64);
 
@@ -668,7 +736,15 @@ mod test {
             MerkleSumTree::<N_ASSETS, N_BYTES>::new("src/merkle_sum_tree/csv/entry_16.csv")
                 .unwrap();
 
-        let circuit = MstInclusionCircuit::<LEVELS, N_ASSETS, N_BYTES>::init(merkle_sum_tree, 0);
+        let user_index = 0;
+
+        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::<LEVELS, N_ASSETS, N_BYTES>::init(
+            merkle_proof,
+            user_entry.clone(),
+        );
 
         let root = BitMapBackend::new("prints/mst-inclusion-layout.png", (2048, 32768))
             .into_drawing_area();

diff --git a/zk_prover/src/merkle_sum_tree/mod.rs b/zk_prover/src/merkle_sum_tree/mod.rs
@@ -2,13 +2,14 @@ mod entry;
 mod mst;
 mod node;
 mod tests;
+mod tree;
 pub mod utils;
 use halo2_proofs::halo2curves::bn256::Fr as Fp;
 
 #[derive(Clone, Debug)]
-pub struct MerkleProof<const N_ASSETS: usize> {
+pub struct MerkleProof<const N_ASSETS: usize, const N_BYTES: usize> {
+    pub leaf: Node<N_ASSETS>,
     pub root_hash: Fp,
-    pub entry: Entry<N_ASSETS>,
     pub sibling_hashes: Vec<Fp>,
     pub sibling_sums: Vec<[Fp; N_ASSETS]>,
     pub path_indices: Vec<Fp>,
@@ -17,4 +18,5 @@ pub struct MerkleProof<const N_ASSETS: usize> {
 pub use entry::Entry;
 pub use mst::MerkleSumTree;
 pub use node::Node;
+pub use tree::Tree;
 pub use utils::{big_intify_username, big_uint_to_fp};