From 9a57b660f237a06a6b041cfc10cef6e558d4ef5b Mon Sep 17 00:00:00 2001 From: M Alghazwi Date: Thu, 30 Jan 2025 13:46:37 +0100 Subject: [PATCH] impl uniform recursion circuits --- codex-plonky2-circuits/src/recursion/mod.rs | 1 + .../src/recursion/uniform/leaf.rs | 136 +++++++++++++++ .../src/recursion/uniform/mod.rs | 3 + .../src/recursion/uniform/node.rs | 147 +++++++++++++++++ .../src/recursion/uniform/tree.rs | 156 ++++++++++++++++++ 5 files changed, 443 insertions(+) create mode 100644 codex-plonky2-circuits/src/recursion/uniform/leaf.rs create mode 100644 codex-plonky2-circuits/src/recursion/uniform/mod.rs create mode 100644 codex-plonky2-circuits/src/recursion/uniform/node.rs create mode 100644 codex-plonky2-circuits/src/recursion/uniform/tree.rs diff --git a/codex-plonky2-circuits/src/recursion/mod.rs b/codex-plonky2-circuits/src/recursion/mod.rs index 070e087..dbc2d24 100644 --- a/codex-plonky2-circuits/src/recursion/mod.rs +++ b/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; diff --git a/codex-plonky2-circuits/src/recursion/uniform/leaf.rs b/codex-plonky2-circuits/src/recursion/uniform/leaf.rs new file mode 100644 index 0000000..9ef988a --- /dev/null +++ b/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 + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, +> where + >::Hasher: AlgebraicHasher +{ + inner_common_data: CommonCircuitData, + phantom_data: PhantomData<(C,H)> +} + +impl< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, +> LeafCircuit where + >::Hasher: AlgebraicHasher +{ + pub fn new(inner_common_data: CommonCircuitData) -> Self { + Self{ + inner_common_data, + phantom_data:PhantomData::default(), + } + } +} +#[derive(Clone, Debug)] +pub struct LeafTargets < + const D: usize, +>{ + pub inner_proof: ProofWithPublicInputsTarget, + pub verifier_data: VerifierCircuitTarget, +} +#[derive(Clone, Debug)] +pub struct LeafInput< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, +>{ + pub inner_proof: ProofWithPublicInputs, + pub verifier_data: VerifierCircuitData +} + +impl< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, +> LeafCircuit where + >::Hasher: AlgebraicHasher +{ + + /// build the leaf circuit + pub fn build(&self, builder: &mut CircuitBuilder) -> Result> { + + 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::(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::(&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, + targets: &LeafTargets, + input: &LeafInput + ) -> 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> + where + >::Hasher: AlgebraicHasher + { + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config.clone()); + + self.build(&mut builder)?; + + let circ_data = builder.build::(); + + Ok(circ_data) + } + +} + + diff --git a/codex-plonky2-circuits/src/recursion/uniform/mod.rs b/codex-plonky2-circuits/src/recursion/uniform/mod.rs new file mode 100644 index 0000000..6daaa29 --- /dev/null +++ b/codex-plonky2-circuits/src/recursion/uniform/mod.rs @@ -0,0 +1,3 @@ +pub mod leaf; +pub mod node; +pub mod tree; \ No newline at end of file diff --git a/codex-plonky2-circuits/src/recursion/uniform/node.rs b/codex-plonky2-circuits/src/recursion/uniform/node.rs new file mode 100644 index 0000000..c9d89d3 --- /dev/null +++ b/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 + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, +> where + >::Hasher: AlgebraicHasher +{ + leaf_common_data: CommonCircuitData, + phantom_data: PhantomData<(C,H)> +} + +impl< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, +> NodeCircuit where + >::Hasher: AlgebraicHasher +{ + pub fn new(inner_common_data: CommonCircuitData) -> 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; 2], + pub verifier_data: VerifierCircuitTarget, +} +#[derive(Clone, Debug)] +pub struct NodeInput< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, +>{ + pub node_proofs: [ProofWithPublicInputs;2], + pub verifier_data: VerifierCircuitData +} + +impl< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, +> NodeCircuit where + >::Hasher: AlgebraicHasher +{ + + /// build the leaf circuit + pub fn build(&self, builder: &mut CircuitBuilder) -> Result> { + + 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::(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::(&vir_proofs[i], &inner_verifier_data, &inner_common); + } + + let proofs = vec_to_array::<2, ProofWithPublicInputsTarget>(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, + targets: &NodeTargets, + input: &NodeInput + ) -> 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> + where + >::Hasher: AlgebraicHasher + { + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config.clone()); + + self.build(&mut builder)?; + + let circ_data = builder.build::(); + + Ok(circ_data) + } + +} + + diff --git a/codex-plonky2-circuits/src/recursion/uniform/tree.rs b/codex-plonky2-circuits/src/recursion/uniform/tree.rs new file mode 100644 index 0000000..d42be53 --- /dev/null +++ b/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 + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, +> where + >::Hasher: AlgebraicHasher +{ + leaf: LeafCircuit, + node: NodeCircuit, + leaf_circ_data: CircuitData, + node_circ_data: CircuitData, + leaf_targets: LeafTargets, + node_targets: NodeTargets, + phantom_data: PhantomData<(H)> +} + +impl< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, +> TreeRecursion where + >::Hasher: AlgebraicHasher +{ + + pub fn build( + inner_common_data: CommonCircuitData + ) -> Result { + // build leaf with standard recursion config + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + + let leaf = LeafCircuit::new(inner_common_data.clone()); + let leaf_targets = leaf.build(&mut builder)?; + let leaf_circ_data = builder.build::(); + // 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::::new(config); + + let node = NodeCircuit::new(leaf_circ_data.common.clone()); + let node_targets = node.build(&mut builder)?; + let node_circ_data = builder.build::(); + // 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], + inner_verifier_data: VerifierCircuitData, + ) -> Result<(ProofWithPublicInputs)> + { + 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], + inner_verifier_data: VerifierCircuitData, + ) -> Result<(Vec>)> { + + let mut leaf_proofs = vec![]; + + for proof in proofs_with_pi{ + let mut pw = PartialWitness::::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], + verifier_data: VerifierCircuitData, + ) -> Result<(ProofWithPublicInputs, VerifierCircuitData)> where + >::Hasher: AlgebraicHasher + { + + 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()) + } + +} +