refactor tests

proof-input/src/recursion/leaf_test.rs

@@ -1,43 +1,20 @@
 // some tests for the leaf in tree recursion
 
 #[cfg(test)]
-mod tests {
+pub mod tests {
     use plonky2::plonk::circuit_data::{ProverCircuitData, VerifierCircuitData};
     use plonky2::plonk::proof::ProofWithPublicInputs;
     use plonky2_field::types::{Field, PrimeField64};
     use codex_plonky2_circuits::circuit_helper::Plonky2Circuit;
-    use codex_plonky2_circuits::circuits::sample_cells::SampleCircuit;
     use codex_plonky2_circuits::recursion::leaf::{LeafCircuit, LeafInput};
-    use codex_plonky2_circuits::recursion::node::{NodeCircuit, NodeInput};
     use codex_plonky2_circuits::recursion::dummy_gen::DummyProofGen;
     use crate::params::{F, D, C, HF};
-    use crate::gen_input::gen_testing_circuit_input;
+    use crate::recursion::run_sampling_circ;
-    use crate::params::Params;
 
-    fn run_sampling_circ() -> anyhow::Result<(ProofWithPublicInputs<F, C, D>, ProverCircuitData<F, C, D>, VerifierCircuitData<F, C, D>)> {
+    pub fn run_leaf_circ<const T: usize>(inner_proof: ProofWithPublicInputs<F, C, D>, inner_verifier_data: VerifierCircuitData<F, C, D>, flag: bool, index: usize) -> anyhow::Result<(ProofWithPublicInputs<F, C, D>, ProverCircuitData<F, C, D>, VerifierCircuitData<F, C, D>)> {
-        //------------ sampling inner circuit ----------------------
-        // Circuit that does the sampling - 100 samples
-        let mut params = Params::default();
-        params.set_n_samples(100);
-        let one_circ_input = gen_testing_circuit_input::<F,D>(&params.input_params);
-        let samp_circ = SampleCircuit::<F,D,HF>::new(params.circuit_params);
-        let (inner_tar, inner_data) = samp_circ.build_with_standard_config()?;
-
-        let inner_verifier_data = inner_data.verifier_data();
-        let inner_prover_data = inner_data.prover_data();
-
-        println!("sampling circuit degree bits = {:?}", inner_verifier_data.common.degree_bits());
-        let inner_proof = samp_circ.prove(&inner_tar, &one_circ_input, &inner_prover_data)?;
-
-        Ok((inner_proof, inner_prover_data, inner_verifier_data))
-    }
-
-    fn run_leaf_circ<const N: usize, const T: usize>(inner_proof: ProofWithPublicInputs<F, C, D>, inner_verifier_data: VerifierCircuitData<F, C, D>, flag: bool, index: usize) -> anyhow::Result<()> {
 
         // ------------------- leaf --------------------
-        // N inner proofs
+        let leaf = LeafCircuit::<F,D,C,HF,T>::new(inner_verifier_data.clone());
-
-        let leaf = LeafCircuit::<F,D,C,HF,N, T>::new(inner_verifier_data.common.clone(),inner_verifier_data.verifier_only.clone());
 
         // build
         let (targets, data) = leaf.build_with_standard_config()?;
@@ -47,7 +24,7 @@ mod tests {
 
         // prove
         let input = LeafInput{
-            inner_proof: vec![inner_proof],
+            inner_proof,
             flag,
             index,
         };
@@ -71,7 +48,7 @@ mod tests {
             }
         }
 
-        Ok(())
+        Ok((proof, prover_data, verifier_data))
     }
 
     fn check_flag_buckets(index: usize, flag_buckets: Vec<F>) {
@@ -89,62 +66,20 @@ mod tests {
         }
     }
 
-    fn run_node_circ<const M: usize, const T: usize>(leaf_proof: ProofWithPublicInputs<F, C, D>, leaf_verifier_data: VerifierCircuitData<F, C, D>, flag: bool, index: usize) -> anyhow::Result<()> {
-
-        // ------------------- Node --------------------
-        // M leaf proofs
-
-        let node = NodeCircuit::<F,D,C,HF,M, T>::new(leaf_verifier_data.common.clone(), leaf_verifier_data.verifier_only.clone());
-
-        // build
-        let (targets, data) = node.build_with_standard_config()?;
-        let verifier_data: VerifierCircuitData<F,C,D> = data.verifier_data();
-        let prover_data = data.prover_data();
-        println!("node circuit degree bits = {:?}", prover_data.common.degree_bits());
-
-        // prove
-        let input = NodeInput{
-            node_proofs: vec![leaf_proof.clone(), leaf_proof.clone()],
-            verifier_only_data: leaf_verifier_data.verifier_only,
-            condition: false,
-            flags: [true; M].to_vec(),
-            index,
-        };
-        let proof = node.prove(&targets, &input, &prover_data)?;
-        println!("pub input size = {}", proof.public_inputs.len());
-        println!("proof size = {:?} bytes", proof.to_bytes().len());
-        println!("pub input = {:?}", proof.public_inputs);
-
-        // verify
-        assert!(
-            verifier_data.verify(proof.clone()).is_ok(),
-            "proof verification failed"
-        );
-
-        // let flag_buckets: Vec<F> = proof.public_inputs[9..13].to_vec();
-        // if flag {
-        //     check_flag_buckets(index, flag_buckets);
-        // } else {
-        //     for i in 0..flag_buckets.len() {
-        //         assert_eq!(flag_buckets[i], F::ZERO, "bucket not valid");
-        //     }
-        // }
-
-        Ok(())
-    }
-
     #[test]
     fn test_real_leaf_circ() -> anyhow::Result<()> {
         let (inner_proof, _, inner_verifier) = run_sampling_circ()?;
 
-        run_leaf_circ::<1, 1>(inner_proof, inner_verifier, true, 1)
+        run_leaf_circ::<128>(inner_proof, inner_verifier, true, 1)?;
+        Ok(())
     }
 
     #[test]
     fn test_dummy_leaf_circ() -> anyhow::Result<()> {
         let (_, _, inner_verifier) = run_sampling_circ()?;
         let (dummy_proof, dummy_vd) = DummyProofGen::gen_dummy_proof_and_vd_zero_pi(&inner_verifier.common)?;
-        run_leaf_circ::<1, 1>(dummy_proof, dummy_vd, false, 0)
+        run_leaf_circ::<128>(dummy_proof, dummy_vd, false, 0)?;
+        Ok(())
     }
 
 }

proof-input/src/recursion/mod.rs

@@ -1,2 +1,29 @@
+use plonky2::plonk::circuit_data::{ProverCircuitData, VerifierCircuitData};
+use plonky2::plonk::proof::ProofWithPublicInputs;
+use codex_plonky2_circuits::circuit_helper::Plonky2Circuit;
+use codex_plonky2_circuits::circuits::sample_cells::SampleCircuit;
+use crate::gen_input::gen_testing_circuit_input;
+use crate::params::{C, D, F, HF, Params};
+
 pub mod tree_test;
-pub mod leaf_test;
+pub mod leaf_test;
+pub mod node_test;
+
+
+pub fn run_sampling_circ() -> anyhow::Result<(ProofWithPublicInputs<F, C, D>, ProverCircuitData<F, C, D>, VerifierCircuitData<F, C, D>)> {
+    //------------ sampling inner circuit ----------------------
+    // Circuit that does the sampling - 100 samples
+    let mut params = Params::default();
+    params.set_n_samples(100);
+    let one_circ_input = gen_testing_circuit_input::<F, D>(&params.input_params);
+    let samp_circ = SampleCircuit::<F,D,HF>::new(params.circuit_params);
+    let (inner_tar, inner_data) = samp_circ.build_with_standard_config()?;
+
+    let inner_verifier_data = inner_data.verifier_data();
+    let inner_prover_data = inner_data.prover_data();
+
+    println!("sampling circuit degree bits = {:?}", inner_verifier_data.common.degree_bits());
+    let inner_proof = samp_circ.prove(&inner_tar, &one_circ_input, &inner_prover_data)?;
+
+    Ok((inner_proof, inner_prover_data, inner_verifier_data))
+}

proof-input/src/recursion/node_test.rs

@@ -0,0 +1,62 @@
+// some tests for the leaf in tree recursion
+
+#[cfg(test)]
+mod tests {
+    use plonky2::plonk::circuit_data::VerifierCircuitData;
+    use plonky2::plonk::proof::ProofWithPublicInputs;
+    use codex_plonky2_circuits::circuit_helper::Plonky2Circuit;
+    use codex_plonky2_circuits::recursion::node::{NodeCircuit, NodeInput};
+    use crate::params::{F, D, C, HF};
+    use crate::recursion::leaf_test::tests::run_leaf_circ;
+    use crate::recursion::run_sampling_circ;
+
+    fn run_node_circ<const N: usize, const T: usize>(leaf_proofs: Vec<ProofWithPublicInputs<F, C, D>>, leaf_verifier_data: VerifierCircuitData<F, C, D>, flag: bool, index: usize) -> anyhow::Result<()> {
+
+        // ------------------- Node --------------------
+        // N leaf proofs
+        assert_eq!(leaf_proofs.len(), N);
+        let node = NodeCircuit::<F,D,C,HF, N, T>::new(leaf_verifier_data.clone());
+
+        // build
+        let (targets, data) = node.build_with_standard_config()?;
+        let verifier_data: VerifierCircuitData<F,C,D> = data.verifier_data();
+        let prover_data = data.prover_data();
+        println!("node circuit degree bits = {:?}", prover_data.common.degree_bits());
+
+        // prove
+        let input = NodeInput{
+            inner_proofs: leaf_proofs,
+            verifier_only_data: leaf_verifier_data.verifier_only,
+            condition: false,
+            flags: [true; N].to_vec(),
+            index,
+        };
+
+        let proof = node.prove(&targets, &input, &prover_data)?;
+        println!("pub input size = {}", proof.public_inputs.len());
+        println!("proof size = {:?} bytes", proof.to_bytes().len());
+        println!("pub input = {:?}", proof.public_inputs);
+
+        // verify
+        assert!(
+            verifier_data.verify(proof.clone()).is_ok(),
+            "proof verification failed"
+        );
+
+        // TODO: check flags
+
+        Ok(())
+    }
+
+
+    #[test]
+    fn test_real_node_circ() -> anyhow::Result<()> {
+        let (inner_proof, _, inner_verifier) = run_sampling_circ()?;
+        // this is a bit wasteful to build leaf twice, TODO: fix this
+        let (leaf_proof_1, _, leaf_verifier) = run_leaf_circ::<128>(inner_proof.clone(), inner_verifier.clone(), true, 0)?;
+        let (leaf_proof_2, _, leaf_verifier) = run_leaf_circ::<128>(inner_proof, inner_verifier, true, 1)?;
+        let leaf_proofs = vec![leaf_proof_1,leaf_proof_2];
+        run_node_circ::<2,128>(leaf_proofs, leaf_verifier, true, 0)
+    }
+
+}

proof-input/src/recursion/tree_test.rs

@@ -4,155 +4,56 @@
 mod tests {
     use plonky2::plonk::proof::{ProofWithPublicInputs};
     use codex_plonky2_circuits::circuit_helper::Plonky2Circuit;
-    use codex_plonky2_circuits::circuits::sample_cells::SampleCircuit;
     use crate::params::{F, D, C, HF};
-    use crate::gen_input::gen_testing_circuit_input;
-    use crate::params::Params;
     use codex_plonky2_circuits::recursion::{tree::TreeRecursion};
-    use codex_plonky2_circuits::recursion::pi_verifier::{PublicInputVerificationCircuit, PublicInputVerificationInput};
+    use crate::recursion::run_sampling_circ;
-    use codex_plonky2_circuits::recursion::tree::get_hash_of_verifier_data;
 
-    #[test]
+    fn run_tree_recursion<const N: usize, const T: usize>(compress: bool) -> anyhow::Result<()> {
-    fn test_uniform_recursion() -> anyhow::Result<()> {
-
-        // total number of proofs to aggregate
-        const T:usize = 4;
 
         //------------ sampling inner circuit ----------------------
         // Circuit that does the sampling - 100 samples
-        let mut params = Params::default();
+        let (inner_proof, inner_prover_data, inner_verifier_data) = run_sampling_circ()?;
-        params.input_params.n_samples = 100;
-        params.circuit_params.n_samples = 100;
-        let one_circ_input = gen_testing_circuit_input::<F,D>(&params.input_params);
-        let samp_circ = SampleCircuit::<F,D,HF>::new(params.circuit_params);
-        let (inner_tar, inner_data) = samp_circ.build_with_standard_config()?;
-
-        let inner_verifier_data = inner_data.verifier_data();
-        let inner_prover_data = inner_data.prover_data();
-
-        println!("sampling circuit degree bits = {:?}", inner_verifier_data.common.degree_bits());
-        let inner_proof = samp_circ.prove(&inner_tar, &one_circ_input, &inner_prover_data)?;
 
         let proofs: Vec<ProofWithPublicInputs<F, C, D>> = (0..T).map(|_i| inner_proof.clone()).collect();
 
         // ------------------- tree --------------------
-        // 2-to-1 tree aggregation
+        // N-to-1 tree aggregation
-        const N: usize = 1;
-        const M: usize = 2;
 
-        let mut tree = TreeRecursion::<F,D,C,HF, N, M, T>::build_with_standard_config(inner_verifier_data.common.clone(), inner_verifier_data.verifier_only.clone())?;
+        let mut tree = TreeRecursion::<F,D,C,HF, N, T>::build_with_standard_config(inner_verifier_data.clone())?;
 
-        // aggregate - no compression
+        // aggregate
-        let root = tree.prove_tree(&proofs)?;
+        let root = if !compress {
+            tree.prove_tree(&proofs)?
+        } else {
+            println!("Mode: tree with compression");
+            tree.prove_tree_and_compress(&proofs)?
+        };
         println!("pub input size = {}", root.public_inputs.len());
         println!("pub input = {:?}", root.public_inputs);
         println!("proof size = {:?} bytes", root.to_bytes().len());
 
-        // aggregate with compression
-        // let root_compressed = tree.prove_tree_and_compress(&proofs)?;
-        // println!("pub input size (compressed) = {}", root_compressed.public_inputs.len());
-        // println!("proof size compressed = {:?} bytes", root_compressed.to_bytes().len());
-
         let inner_pi: Vec<Vec<F>> = proofs.iter().map(|p| p.public_inputs.clone()).collect();
 
         assert!(
-            tree.verify_proof_and_public_input(root,inner_pi.clone(), false).is_ok(),
+            tree.verify_proof_and_public_input(root,inner_pi.clone(), compress).is_ok(),
             "proof verification failed"
         );
 
-        // assert!(
-        //     tree.verify_proof_and_public_input(root_compressed,inner_pi, true).is_ok(),
-        //     "compressed proof verification failed"
-        // );
-
-
         Ok(())
     }
 
     #[test]
-    fn test_pi_verifier() -> anyhow::Result<()> {
+    fn test_tree_recursion() -> anyhow::Result<()> {
-
         // total number of proofs to aggregate
         const T:usize = 4;
-        // 9 field elems as public inputs in the sampling circuit
+        run_tree_recursion::<2, T>(false)
-        const K:usize = 9;
-
-        //------------ sampling inner circuit ----------------------
-        // Circuit that does the sampling - 100 samples
-        let mut params = Params::default();
-        params.input_params.n_samples = 100;
-        params.circuit_params.n_samples = 100;
-        let one_circ_input = gen_testing_circuit_input::<F,D>(&params.input_params);
-        let samp_circ = SampleCircuit::<F,D,HF>::new(params.circuit_params);
-        let (inner_tar, inner_data) = samp_circ.build_with_standard_config()?;
-
-        let inner_verifier_data = inner_data.verifier_data();
-        let inner_prover_data = inner_data.prover_data();
-
-        // get generate a sampling proof
-        println!("sampling circuit degree bits = {:?}", inner_verifier_data.common.degree_bits());
-        let inner_proof = samp_circ.prove(&inner_tar, &one_circ_input, &inner_prover_data)?;
-
-        let proofs: Vec<ProofWithPublicInputs<F, C, D>> = (0..T).map(|_i| inner_proof.clone()).collect();
-
-        // ------------------- tree --------------------
-        const N: usize = 1;
-        const M: usize = 2;
-
-        let mut tree = TreeRecursion::<F,D,C,HF, N, M, T>::build_with_standard_config(inner_verifier_data.common.clone(), inner_verifier_data.verifier_only.clone())?;
-
-        let root = tree.prove_tree(&proofs)?;
-        println!("pub input size = {}", root.public_inputs.len());
-        println!("proof size = {:?} bytes", root.to_bytes().len());
-
-        let inner_pi: Vec<Vec<F>> = proofs.iter().map(|p| p.public_inputs.clone()).collect();
-
-        assert!(
-            tree.verify_proof_and_public_input(root.clone(),inner_pi.clone(), false).is_ok(),
-            "proof verification failed"
-        );
-
-        // ------------------- Public input verifier Circuit --------------------
-
-        let pi_verifier_circ = PublicInputVerificationCircuit::<F, D, C, HF, N, M, T, K>::new(tree.get_node_common_data(), tree.get_node_verifier_data().verifier_only);
-
-        let (pi_tarq, pi_circ_data) = pi_verifier_circ.build_with_standard_config()?;
-        println!("PI verifier circuit degree bits = {:?}", pi_circ_data.common.degree_bits());
-
-        let pi_circ_input = PublicInputVerificationInput{
-            inner_proof:root,
-            inner_pub_inputs_vals: inner_pi.clone()
-        };
-
-        let pi_circ_verifier_data = pi_circ_data.verifier_data();
-        let pi_circ_prover_data = pi_circ_data.prover_data();
-
-        let proof =pi_verifier_circ.prove(&pi_tarq, &pi_circ_input, &pi_circ_prover_data)?;
-
-        println!("pub input size = {}", proof.public_inputs.len());
-        println!("proof size = {:?} bytes", proof.to_bytes().len());
-
-        let pub_input_flat: Vec<F> = inner_pi.iter().cloned().flatten().collect();
-
-        // sanity check on public input
-        for (i, e) in proof.public_inputs.iter().enumerate(){
-            if i < pub_input_flat.len() {
-                assert_eq!(*e, pub_input_flat[i])
-            }
-        }
-
-        // sanity check on the verifier data
-        let hashed_node_vd = get_hash_of_verifier_data::<F,D,C,HF>(&tree.get_node_verifier_data());
-        for (i, &e) in proof.public_inputs[proof.public_inputs.len()-4 ..].iter().enumerate(){
-            assert_eq!(e, hashed_node_vd.elements[i])
-        }
-
-
-        assert!(
-            pi_circ_verifier_data.verify(proof).is_ok(),
-            "pi-verifier proof verification failed"
-        );
-
-        Ok(())
     }
+
+    #[test]
+    fn test_tree_recursion_with_compression() -> anyhow::Result<()> {
+        // total number of proofs to aggregate
+        const T:usize = 4;
+        run_tree_recursion::<2, T>(true)
+    }
+
 }