impl uniform recursion circuits

codex-plonky2-circuits/src/recursion/mod.rs

@@ -5,3 +5,4 @@ pub mod tree1;
 pub mod tree2;
 pub mod hybrid;
 pub mod utils;
+pub mod uniform;

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

@@ -0,0 +1,136 @@
+use std::marker::PhantomData;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::witness::{PartialWitness, WitnessWrite};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, CommonCircuitData, VerifierCircuitData, VerifierCircuitTarget};
+use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
+use plonky2::plonk::proof::{ProofWithPublicInputs, ProofWithPublicInputsTarget};
+use plonky2_field::extension::Extendable;
+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
+#[derive(Clone, Debug)]
+pub struct LeafCircuit<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    inner_common_data: CommonCircuitData<F, D>,
+    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 verifier_data: VerifierCircuitTarget,
+}
+#[derive(Clone, Debug)]
+pub struct LeafInput<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+>{
+    pub inner_proof: ProofWithPublicInputs<F, C, D>,
+    pub verifier_data: VerifierCircuitData<F, C, D>
+}
+
+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>
+{
+
+    /// 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
+        let mut pub_input = vec![];
+        let vir_proof = builder.add_virtual_proof_with_pis(&inner_common);
+        let inner_pub_input = vir_proof.public_inputs.clone();
+        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);
+        builder.register_public_inputs(&hash_inner_pub_input.elements);
+
+        // virtual target for the verifier data
+        let inner_verifier_data = builder.add_virtual_verifier_data(inner_common.config.fri_config.cap_height);
+
+        // verify the proofs in-circuit (only one )
+        builder.verify_proof::<C>(&vir_proof, &inner_verifier_data, &inner_common);
+
+        // return targets
+        let t = LeafTargets {
+            inner_proof: vir_proof,
+            verifier_data: inner_verifier_data,
+        };
+        Ok(t)
+
+    }
+
+    /// assign the leaf targets with given input
+    pub fn assign_targets(
+        &self, pw: &mut PartialWitness<F>,
+        targets: &LeafTargets<D>,
+        input: &LeafInput<F, D, C>
+    ) -> Result<()> {
+        // assign the proofs
+        pw.set_proof_with_pis_target(&targets.inner_proof, &input.inner_proof)
+            .map_err(|e| {
+                CircuitError::ProofTargetAssignmentError("inner-proof".to_string(), e.to_string())
+            })?;
+
+        // assign the verifier data
+        pw.set_verifier_data_target(&targets.verifier_data, &input.verifier_data.verifier_only)
+            .map_err(|e| {
+                CircuitError::VerifierDataTargetAssignmentError(e.to_string())
+            })?;
+
+        Ok(())
+    }
+
+    /// returns the leaf circuit data
+    pub fn get_circuit_data (&self) -> Result<CircuitData<F, C, D>>
+        where
+            <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+    {
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config.clone());
+
+        self.build(&mut builder)?;
+
+        let circ_data = builder.build::<C>();
+
+        Ok(circ_data)
+    }
+
+}
+
+

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

@@ -0,0 +1,3 @@
+pub mod leaf;
+pub mod node;
+pub mod tree;

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

@@ -0,0 +1,147 @@
+use std::marker::PhantomData;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::witness::{PartialWitness, WitnessWrite};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, CommonCircuitData, VerifierCircuitData, VerifierCircuitTarget};
+use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
+use plonky2::plonk::proof::{ProofWithPublicInputs, ProofWithPublicInputsTarget};
+use plonky2_field::extension::Extendable;
+use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
+use crate::recursion::circuits::inner_circuit::InnerCircuit;
+use crate::{error::CircuitError,Result};
+use crate::circuits::utils::vec_to_array;
+
+/// recursion node circuit - verifies 2 leaf proofs
+#[derive(Clone, Debug)]
+pub struct NodeCircuit<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    leaf_common_data: CommonCircuitData<F, D>,
+    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 verifier_data: VerifierCircuitTarget,
+}
+#[derive(Clone, Debug)]
+pub struct NodeInput<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+>{
+    pub node_proofs: [ProofWithPublicInputs<F, C, D>;2],
+    pub verifier_data: VerifierCircuitData<F, C, D>
+}
+
+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>
+{
+
+    /// build the leaf circuit
+    pub fn build(&self, builder: &mut CircuitBuilder<F, D>) -> Result<NodeTargets<D>> {
+
+        let inner_common = self.leaf_common_data.clone();
+
+        // the proof virtual targets - 2 proofs
+        let mut vir_proofs = vec![];
+        let mut pub_input = vec![];
+        for _i in 0..2 {
+            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);
+        builder.register_public_inputs(&hash_inner_pub_input.elements);
+
+        // virtual target for the verifier data
+        let inner_verifier_data = builder.add_virtual_verifier_data(inner_common.config.fri_config.cap_height);
+
+        // verify the proofs in-circuit  - 2 proofs
+        for i in 0..2 {
+        builder.verify_proof::<C>(&vir_proofs[i], &inner_verifier_data, &inner_common);
+        }
+
+        let proofs = vec_to_array::<2, ProofWithPublicInputsTarget<D>>(vir_proofs)?;
+
+        // return targets
+        let t = NodeTargets {
+            leaf_proofs: proofs,
+            verifier_data: inner_verifier_data,
+        };
+        Ok(t)
+
+    }
+
+    /// assign the leaf targets with given input
+    pub fn assign_targets(
+        &self, pw: &mut PartialWitness<F>,
+        targets: &NodeTargets<D>,
+        input: &NodeInput<F, D, C>
+    ) -> Result<()> {
+        // assign the proofs
+        for i in 0..2 {
+            pw.set_proof_with_pis_target(&targets.leaf_proofs[i], &input.node_proofs[i])
+                .map_err(|e| {
+                    CircuitError::ProofTargetAssignmentError("inner-proof".to_string(), e.to_string())
+                })?;
+        }
+
+        // assign the verifier data
+        pw.set_verifier_data_target(&targets.verifier_data, &input.verifier_data.verifier_only)
+            .map_err(|e| {
+                CircuitError::VerifierDataTargetAssignmentError(e.to_string())
+            })?;
+
+        Ok(())
+    }
+
+    /// returns the leaf circuit data
+    pub fn get_circuit_data (&self) -> Result<CircuitData<F, C, D>>
+        where
+            <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+    {
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config.clone());
+
+        self.build(&mut builder)?;
+
+        let circ_data = builder.build::<C>();
+
+        Ok(circ_data)
+    }
+
+}
+
+

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

@@ -0,0 +1,156 @@
+use std::marker::PhantomData;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::witness::PartialWitness;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, CommonCircuitData, VerifierCircuitData};
+use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
+use plonky2::plonk::proof::ProofWithPublicInputs;
+use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
+use crate::recursion::circuits::inner_circuit::InnerCircuit;
+use plonky2_field::extension::Extendable;
+use crate::{error::CircuitError, Result};
+use crate::recursion::uniform::{leaf::{LeafTargets,LeafInput,LeafCircuit},node::{NodeTargets,NodeInput,NodeCircuit}};
+
+/// tree recursion
+pub struct TreeRecursion<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    leaf: LeafCircuit<F, D, C, H>,
+    node: NodeCircuit<F, D, C, H>,
+    leaf_circ_data: CircuitData<F, C, D>,
+    node_circ_data: CircuitData<F, C, D>,
+    leaf_targets: LeafTargets<D>,
+    node_targets: NodeTargets<D>,
+    phantom_data: PhantomData<(H)>
+}
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+> TreeRecursion<F, D, C, H> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+
+    pub fn build(
+        inner_common_data: CommonCircuitData<F,D>
+    ) -> Result<Self> {
+        // build leaf with standard recursion config
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+
+        let leaf = LeafCircuit::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());
+
+        // build node with standard recursion config
+        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_targets = node.build(&mut builder)?;
+        let node_circ_data = builder.build::<C>();
+        // println!("node circuit size = {:?}", node_circ_data.common.degree_bits());
+
+        Ok(Self{
+            leaf,
+            node,
+            leaf_circ_data,
+            node_circ_data,
+            leaf_targets,
+            node_targets,
+            phantom_data: Default::default(),
+        })
+    }
+
+    pub fn prove_tree
+    (
+        &mut self,
+        proofs_with_pi: &[ProofWithPublicInputs<F, C, D>],
+        inner_verifier_data: VerifierCircuitData<F, C, D>,
+    ) -> Result<(ProofWithPublicInputs<F, C, D>)>
+    {
+        if proofs_with_pi.len() % 2 != 0 {
+            return
+                Err(CircuitError::RecursionTreeError(format!(
+                    "input proofs must be divisible by {}, got {}", 2, proofs_with_pi.len())
+                ))
+        }
+        // process leaves
+        let leaf_proofs = self.get_leaf_proofs(
+            proofs_with_pi,
+            inner_verifier_data,
+        )?;
+
+        // process nodes
+        let (root_proof, vd) =
+            self.prove(&leaf_proofs,self.leaf_circ_data.verifier_data())?;
+
+        Ok(root_proof)
+    }
+
+
+    fn get_leaf_proofs
+    (
+        &mut self,
+        proofs_with_pi: &[ProofWithPublicInputs<F, C, D>],
+        inner_verifier_data: VerifierCircuitData<F, C, D>,
+    ) -> Result<(Vec<ProofWithPublicInputs<F, C, D>>)> {
+
+        let mut leaf_proofs = vec![];
+
+        for proof in proofs_with_pi{
+            let mut pw = PartialWitness::<F>::new();
+            let leaf_in = LeafInput{
+                inner_proof: proof.clone(),
+                verifier_data: inner_verifier_data.clone(),
+            };
+            self.leaf.assign_targets(&mut pw,&self.leaf_targets,&leaf_in)?;
+            let proof = self.leaf_circ_data.prove(pw).unwrap();
+            leaf_proofs.push(proof);
+        }
+
+        Ok(leaf_proofs)
+    }
+
+    /// generates a proof - only one node
+    fn prove(
+        &mut self,
+        proofs_with_pi: &[ProofWithPublicInputs<F, C, D>],
+        verifier_data: VerifierCircuitData<F, C, D>,
+    ) -> Result<(ProofWithPublicInputs<F, C, D>, VerifierCircuitData<F, C, D>)> where
+        <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+    {
+
+        if proofs_with_pi.len() == 1 {
+            return Ok((proofs_with_pi[0].clone(), verifier_data));
+        }
+
+        let mut new_proofs = vec![];
+
+        for chunk in proofs_with_pi.chunks(2) {
+
+            let mut inner_pw = PartialWitness::new();
+            let node_in = NodeInput{
+                node_proofs: [chunk[0].clone(), chunk[1].clone()],
+                verifier_data: verifier_data.clone() ,
+            };
+            self.node.assign_targets(&mut inner_pw,&self.node_targets,&node_in)?;
+
+            let proof = self.node_circ_data.prove(inner_pw)
+                .map_err(|e| CircuitError::ProofGenerationError(e.to_string()))?;
+            new_proofs.push(proof);
+        }
+
+        self.prove(&new_proofs, self.node_circ_data.verifier_data())
+    }
+
+}
+