add support N-to-1 leaf and node aggregation

codex-plonky2-circuits/src/recursion/uniform/leaf.rs

@@ -10,13 +10,14 @@ use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
 use crate::recursion::circuits::inner_circuit::InnerCircuit;
 use crate::{error::CircuitError,Result};
 
-/// recursion leaf circuit - verifies one inner proof
+/// recursion leaf circuit - verifies N inner proof
 #[derive(Clone, Debug)]
 pub struct LeafCircuit<
     F: RichField + Extendable<D> + Poseidon2,
     const D: usize,
     C: GenericConfig<D, F = F>,
     H: AlgebraicHasher<F>,
+    const N: usize,
 > where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
 {
@@ -24,26 +25,11 @@ pub struct LeafCircuit<
     phantom_data: PhantomData<(C,H)>
 }
 
-impl<
-    F: RichField + Extendable<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
-> LeafCircuit<F,D,C,H> where
-    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-    pub fn new(inner_common_data: CommonCircuitData<F,D>) -> Self {
-        Self{
-            inner_common_data,
-            phantom_data:PhantomData::default(),
-        }
-    }
-}
 #[derive(Clone, Debug)]
 pub struct LeafTargets <
     const D: usize,
 >{
-    pub inner_proof: ProofWithPublicInputsTarget<D>,
+    pub inner_proof: Vec<ProofWithPublicInputsTarget<D>>,
     pub verifier_data: VerifierCircuitTarget,
 }
 
@@ -52,20 +38,31 @@ impl<
     const D: usize,
     C: GenericConfig<D, F = F>,
     H: AlgebraicHasher<F>,
-> LeafCircuit<F,D,C,H> where
+    const N: usize,
+> LeafCircuit<F,D,C,H,N> where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
 {
+    pub fn new(inner_common_data: CommonCircuitData<F,D>) -> Self {
+        Self{
+            inner_common_data,
+            phantom_data:PhantomData::default(),
+        }
+    }
 
     /// build the leaf circuit
     pub fn build(&self, builder: &mut CircuitBuilder<F, D>) -> Result<LeafTargets<D>> {
 
         let inner_common = self.inner_common_data.clone();
 
-        // the proof virtual targets - only one
+        // the proof virtual targets
         let mut pub_input = vec![];
-        let vir_proof = builder.add_virtual_proof_with_pis(&inner_common);
+        let mut vir_proofs = vec![];
-        let inner_pub_input = vir_proof.public_inputs.clone();
+        for _i in 0..N {
-        pub_input.extend_from_slice(&inner_pub_input);
+            let vir_proof = builder.add_virtual_proof_with_pis(&inner_common);
+            let inner_pub_input = vir_proof.public_inputs.clone();
+            vir_proofs.push(vir_proof);
+            pub_input.extend_from_slice(&inner_pub_input);
+        }
 
         // hash the public input & make it public
         let hash_inner_pub_input = builder.hash_n_to_hash_no_pad::<H>(pub_input);
@@ -83,12 +80,14 @@ impl<
         let hash_inner_vd_pub_input = builder.hash_n_to_hash_no_pad::<H>(vd_pub_input);
         builder.register_public_inputs(&hash_inner_vd_pub_input.elements);
 
-        // verify the proofs in-circuit (only one )
+        // verify the proofs in-circuit
-        builder.verify_proof::<C>(&vir_proof, &inner_verifier_data, &inner_common);
+        for i in 0..N {
+            builder.verify_proof::<C>(&vir_proofs[i], &inner_verifier_data, &inner_common);
+        }
 
         // return targets
         let t = LeafTargets {
-            inner_proof: vir_proof,
+            inner_proof: vir_proofs,
             verifier_data: inner_verifier_data,
         };
         Ok(t)
@@ -99,14 +98,17 @@ impl<
     pub fn assign_targets(
         &self, pw: &mut PartialWitness<F>,
         targets: &LeafTargets<D>,
-        inner_proof: &ProofWithPublicInputs<F, C, D>,
+        inner_proof: &[ProofWithPublicInputs<F, C, D>],
         verifier_only_data: &VerifierOnlyCircuitData<C, D>,
     ) -> Result<()> {
+        assert_eq!(inner_proof.len(), N);
         // assign the proofs
-        pw.set_proof_with_pis_target(&targets.inner_proof, inner_proof)
+        for i in 0..N {
-            .map_err(|e| {
+            pw.set_proof_with_pis_target(&targets.inner_proof[i], &inner_proof[i])
-                CircuitError::ProofTargetAssignmentError("inner-proof".to_string(), e.to_string())
+                .map_err(|e| {
-            })?;
+                    CircuitError::ProofTargetAssignmentError("inner-proof".to_string(), e.to_string())
+                })?;
+        }
 
         // assign the verifier data
         pw.set_verifier_data_target(&targets.verifier_data, verifier_only_data)

codex-plonky2-circuits/src/recursion/uniform/node.rs

@@ -18,6 +18,7 @@ pub struct NodeCircuit<
     const D: usize,
     C: GenericConfig<D, F = F>,
     H: AlgebraicHasher<F>,
+    const M: usize,
 > where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
 {
@@ -25,26 +26,11 @@ pub struct NodeCircuit<
     phantom_data: PhantomData<(C,H)>
 }
 
-impl<
-    F: RichField + Extendable<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
-> NodeCircuit<F,D,C,H> where
-    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-    pub fn new(inner_common_data: CommonCircuitData<F,D>) -> Self {
-        Self{
-            leaf_common_data: inner_common_data,
-            phantom_data:PhantomData::default(),
-        }
-    }
-}
 #[derive(Clone, Debug)]
 pub struct NodeTargets<
     const D: usize,
 >{
-    pub leaf_proofs: [ProofWithPublicInputsTarget<D>; 2],
+    pub leaf_proofs: Vec<ProofWithPublicInputsTarget<D>>,
     pub verifier_data: VerifierCircuitTarget,
 }
 
@@ -53,10 +39,18 @@ impl<
     const D: usize,
     C: GenericConfig<D, F = F>,
     H: AlgebraicHasher<F>,
-> NodeCircuit<F,D,C,H> where
+    const M: usize,
+> NodeCircuit<F,D,C,H,M> where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
 {
 
+    pub fn new(inner_common_data: CommonCircuitData<F,D>) -> Self {
+        Self{
+            leaf_common_data: inner_common_data,
+            phantom_data:PhantomData::default(),
+        }
+    }
+
     /// build the leaf circuit
     pub fn build(&self, builder: &mut CircuitBuilder<F, D>) -> Result<NodeTargets<D>> {
 
@@ -65,11 +59,11 @@ impl<
         // assert public input is of size 8 - 2 hashout
         assert_eq!(inner_common.num_public_inputs, 8);
 
-        // the proof virtual targets - 2 proofs
+        // the proof virtual targets - M proofs
         let mut vir_proofs = vec![];
         let mut pub_input = vec![];
         let mut inner_vd_hashes = vec![];
-        for _i in 0..2 {
+        for _i in 0..M {
             let vir_proof = builder.add_virtual_proof_with_pis(&inner_common);
             let inner_pub_input = vir_proof.public_inputs.clone();
             vir_proofs.push(vir_proof);
@@ -95,16 +89,16 @@ impl<
         let vd_hash_all = builder.hash_n_to_hash_no_pad::<H>(inner_vd_hashes);
         builder.register_public_inputs(&vd_hash_all.elements);
 
-        // verify the proofs in-circuit  - 2 proofs
+        // verify the proofs in-circuit  - M proofs
-        for i in 0..2 {
+        for i in 0..M {
         builder.verify_proof::<C>(&vir_proofs[i], &inner_verifier_data, &inner_common);
         }
 
-        let proofs = vec_to_array::<2, ProofWithPublicInputsTarget<D>>(vir_proofs)?;
+        // let proofs = vec_to_array::<2, ProofWithPublicInputsTarget<D>>(vir_proofs)?;
 
         // return targets
         let t = NodeTargets {
-            leaf_proofs: proofs,
+            leaf_proofs: vir_proofs,
             verifier_data: inner_verifier_data,
         };
         Ok(t)
@@ -119,10 +113,10 @@ impl<
         verifier_only_data: &VerifierOnlyCircuitData<C, D>,
     ) -> Result<()> {
         // assert size of proofs vec
-        assert_eq!(node_proofs.len(), 2);
+        assert_eq!(node_proofs.len(), M);
 
         // assign the proofs
-        for i in 0..2 {
+        for i in 0..M {
             pw.set_proof_with_pis_target(&targets.leaf_proofs[i], &node_proofs[i])
                 .map_err(|e| {
                     CircuitError::ProofTargetAssignmentError("inner-proof".to_string(), e.to_string())

codex-plonky2-circuits/src/recursion/uniform/tree.rs

@@ -17,11 +17,13 @@ pub struct TreeRecursion<
     const D: usize,
     C: GenericConfig<D, F = F>,
     H: AlgebraicHasher<F>,
+    const N: usize,
+    const M: usize,
 > where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
 {
-    leaf: LeafCircuit<F, D, C, H>,
+    leaf: LeafCircuit<F, D, C, H, N>,
-    node: NodeCircuit<F, D, C, H>,
+    node: NodeCircuit<F, D, C, H, M>,
     leaf_circ_data: CircuitData<F, C, D>,
     node_circ_data: CircuitData<F, C, D>,
     leaf_targets: LeafTargets<D>,
@@ -34,7 +36,9 @@ impl<
     const D: usize,
     C: GenericConfig<D, F = F>,
     H: AlgebraicHasher<F>,
-> TreeRecursion<F, D, C, H> where
+    const N: usize,
+    const M: usize,
+> TreeRecursion<F, D, C, H, N, M> where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
 {
 
@@ -45,7 +49,7 @@ impl<
         let config = CircuitConfig::standard_recursion_config();
         let mut builder = CircuitBuilder::<F, D>::new(config);
 
-        let leaf = LeafCircuit::new(inner_common_data.clone());
+        let leaf = LeafCircuit::<_,D,_,_,N>::new(inner_common_data.clone());
         let leaf_targets = leaf.build(&mut builder)?;
         let leaf_circ_data = builder.build::<C>();
         // println!("leaf circuit size = {:?}", leaf_circ_data.common.degree_bits());
@@ -54,7 +58,7 @@ impl<
         let config = CircuitConfig::standard_recursion_config();
         let mut builder = CircuitBuilder::<F, D>::new(config);
 
-        let node = NodeCircuit::new(leaf_circ_data.common.clone());
+        let node = NodeCircuit::<_,D,_,_,M>::new(leaf_circ_data.common.clone());
         let node_targets = node.build(&mut builder)?;
         let node_circ_data = builder.build::<C>();
         // println!("node circuit size = {:?}", node_circ_data.common.degree_bits());
@@ -113,7 +117,7 @@ impl<
 
         let mut leaf_proofs = vec![];
 
-        for proof in proofs_with_pi{
+        for proof in proofs_with_pi.chunks(N){
             let mut pw = PartialWitness::<F>::new();
 
             self.leaf.assign_targets(&mut pw,&self.leaf_targets,proof,inner_verifier_only_data)?;
@@ -124,7 +128,7 @@ impl<
         Ok(leaf_proofs)
     }
 
-    /// generates a proof - only one node
+    /// generates a proof
     fn prove(
         &self,
         proofs_with_pi: &[ProofWithPublicInputs<F, C, D>],
@@ -139,7 +143,7 @@ impl<
 
         let mut new_proofs = vec![];
 
-        for chunk in proofs_with_pi.chunks(2) {
+        for chunk in proofs_with_pi.chunks(M) {
 
             let mut inner_pw = PartialWitness::new();
 
@@ -185,8 +189,12 @@ impl<
 
         let mut pub_in_hashes = vec![];
         let mut inner_vd_hashes = vec![];
-        for pub_in in inner_public_input{
+        for pub_in in inner_public_input.chunks(N){
-            let hash = H::hash_no_pad(&pub_in);
+            let pub_in_flat: Vec<F> = pub_in
+                .iter()
+                .flat_map(|v| v.iter().cloned())
+                .collect();
+            let hash = H::hash_no_pad(&pub_in_flat);
             pub_in_hashes.push(hash);
             inner_vd_hashes.push(inner_hash.clone());
         }
@@ -195,7 +203,7 @@ impl<
         while pub_in_hashes.len() > 1 {
             let mut next_level_pi_hashes = Vec::new();
             let mut next_level_vd_hashes = Vec::new();
-            for (pi_chunk, vd_chunk) in pub_in_hashes.chunks(2).zip(inner_vd_hashes.chunks(2)) {
+            for (pi_chunk, vd_chunk) in pub_in_hashes.chunks(M).zip(inner_vd_hashes.chunks(M)) {
                 // collect field elements
                 let pi_chunk_f: Vec<F> = pi_chunk.iter()
                     .flat_map(|h| h.elements.iter().cloned())

proof-input/src/recursion/uniform.rs

@@ -35,11 +35,22 @@ mod tests {
         println!("sampling circuit degree bits = {:?}", inner_data.common.degree_bits());
         let inner_proof = inner_data.prove(pw)?;
 
-        let proofs: Vec<ProofWithPublicInputs<F, C, D>> = (0..4).map(|i| inner_proof.clone()).collect();
+        let proofs: Vec<ProofWithPublicInputs<F, C, D>> = (0..16).map(|i| inner_proof.clone()).collect();
 
         // ------------------- tree --------------------
+        const N: usize = 1;
+        const M: usize = 4;
+
+        let mut tree = TreeRecursion::<F,D,C,HF, N, M>::build(inner_data.common.clone())?;
+
+        // serialize circuit into JSON
+        let common_circuit_data_serialized        = serde_json::to_string(&tree.get_leaf_verifier_data().common       ).unwrap();
+        fs::write("leaf_common.json" , common_circuit_data_serialized)       .expect("Unable to write file");
+
+        // serialize circuit into JSON
+        let common_circuit_data_serialized        = serde_json::to_string(&tree.get_node_verifier_data().common       ).unwrap();
+        fs::write("node_common.json" , common_circuit_data_serialized)       .expect("Unable to write file");
 
-        let mut tree = TreeRecursion::<F,D,C,HF>::build(inner_data.common.clone())?;
 
         let root = tree.prove_tree(&proofs, &inner_data.verifier_only)?;
         println!("pub input size = {}", root.public_inputs.len());