From 1808143987d2debe36f25516a89319ae0535ecb0 Mon Sep 17 00:00:00 2001 From: M Alghazwi Date: Thu, 24 Apr 2025 20:54:59 +0200 Subject: [PATCH] implement proof tracking and refactor --- codex-plonky2-circuits/Cargo.toml | 4 + codex-plonky2-circuits/src/bundle/mod.rs | 64 +++ codex-plonky2-circuits/src/lib.rs | 1 + .../src/recursion/{uniform => }/compress.rs | 3 +- .../src/recursion/dummy_gen.rs | 105 +++++ codex-plonky2-circuits/src/recursion/leaf.rs | 269 +++++++++++++ codex-plonky2-circuits/src/recursion/mod.rs | 7 +- codex-plonky2-circuits/src/recursion/node.rs | 370 ++++++++++++++++++ .../recursion/{uniform => }/pi_verifier.rs | 1 + .../src/recursion/{uniform => }/tree.rs | 101 ++++- .../src/recursion/uniform/leaf.rs | 139 ------- .../src/recursion/uniform/mod.rs | 5 - .../src/recursion/uniform/node.rs | 175 --------- codex-plonky2-circuits/src/recursion/utils.rs | 245 ++++++++++++ .../recursion/utils/conditional_verifier.rs | 121 ------ .../src/recursion/utils/dummy_gen.rs | 79 ---- .../src/recursion/utils/mod.rs | 2 - 17 files changed, 1153 insertions(+), 538 deletions(-) create mode 100644 codex-plonky2-circuits/src/bundle/mod.rs rename codex-plonky2-circuits/src/recursion/{uniform => }/compress.rs (97%) create mode 100644 codex-plonky2-circuits/src/recursion/dummy_gen.rs create mode 100644 codex-plonky2-circuits/src/recursion/leaf.rs create mode 100644 codex-plonky2-circuits/src/recursion/node.rs rename codex-plonky2-circuits/src/recursion/{uniform => }/pi_verifier.rs (99%) rename codex-plonky2-circuits/src/recursion/{uniform => }/tree.rs (73%) delete mode 100644 codex-plonky2-circuits/src/recursion/uniform/leaf.rs delete mode 100644 codex-plonky2-circuits/src/recursion/uniform/mod.rs delete mode 100644 codex-plonky2-circuits/src/recursion/uniform/node.rs create mode 100644 codex-plonky2-circuits/src/recursion/utils.rs delete mode 100644 codex-plonky2-circuits/src/recursion/utils/conditional_verifier.rs delete mode 100644 codex-plonky2-circuits/src/recursion/utils/dummy_gen.rs delete mode 100644 codex-plonky2-circuits/src/recursion/utils/mod.rs diff --git a/codex-plonky2-circuits/Cargo.toml b/codex-plonky2-circuits/Cargo.toml index 5eb96b7..3fd80ac 100644 --- a/codex-plonky2-circuits/Cargo.toml +++ b/codex-plonky2-circuits/Cargo.toml @@ -23,3 +23,7 @@ hashbrown = "0.14.5" criterion = { version = "0.5.1", default-features = false } tynm = { version = "0.1.6", default-features = false } +[features] +default = [] +parallel = ["plonky2/parallel"] + diff --git a/codex-plonky2-circuits/src/bundle/mod.rs b/codex-plonky2-circuits/src/bundle/mod.rs new file mode 100644 index 0000000..6d51996 --- /dev/null +++ b/codex-plonky2-circuits/src/bundle/mod.rs @@ -0,0 +1,64 @@ +use hashbrown::HashMap; +use plonky2::hash::hash_types::RichField; +use plonky2::plonk::circuit_data::{ProverCircuitData, VerifierCircuitData}; +use plonky2::plonk::config::{AlgebraicHasher, GenericConfig}; +use plonky2::plonk::proof::ProofWithPublicInputs; +use plonky2_field::extension::Extendable; +use plonky2_poseidon2::Poseidon2; +use crate::circuit_helper::Plonky2Circuit; +use crate::circuits::params::CircuitParams; +use crate::circuits::sample_cells::{SampleCircuit, SampleCircuitInput, SampleTargets}; +use crate::Result; + +pub struct Bundle< + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + const D: usize, + H: AlgebraicHasher, +>{ + pub bundle_size: usize, + pub circuit: SampleCircuit, + pub prover_data: ProverCircuitData, + pub verifier_data: VerifierCircuitData, + pub sample_targets: SampleTargets, + pub bundle_proofs: HashMap>, +} + +impl< + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + const D: usize, + H: AlgebraicHasher, +> Bundle { + pub fn new(bundle_size: usize, circuit_params: CircuitParams) -> Result<(Self)>{ + let samp_circ = SampleCircuit::::new(circuit_params.clone()); + let (sample_targets, circuit_data) = samp_circ.build_with_standard_config()?; + println!("sampling circuit built. Degree bits = {:?}", circuit_data.common.degree_bits()); + let verifier_data = circuit_data.verifier_data(); + let prover_data = circuit_data.prover_data(); + Ok(Self{ + bundle_size, + circuit: samp_circ, + prover_data, + verifier_data, + sample_targets, + bundle_proofs: HashMap::new(), + }) + } + + pub fn prove_all(&mut self, circ_inputs: Vec>) -> Result<()>{ + assert_eq!(circ_inputs.len(), self.bundle_size, "not enough circuit input provided"); + for (i, input) in circ_inputs.into_iter().enumerate(){ + self.prove(i, input)?; + } + Ok(()) + } + + pub fn prove(&mut self, index: usize, circ_input: SampleCircuitInput) -> Result<()>{ + let proof = self.circuit.prove(&self.sample_targets, &circ_input, &self.prover_data)?; + self.bundle_proofs.insert(index, proof); + Ok(()) + } + +} + diff --git a/codex-plonky2-circuits/src/lib.rs b/codex-plonky2-circuits/src/lib.rs index 7926659..e0dffed 100644 --- a/codex-plonky2-circuits/src/lib.rs +++ b/codex-plonky2-circuits/src/lib.rs @@ -2,5 +2,6 @@ pub mod circuits; pub mod recursion; pub mod error; pub mod circuit_helper; +mod bundle; pub type Result = core::result::Result; diff --git a/codex-plonky2-circuits/src/recursion/uniform/compress.rs b/codex-plonky2-circuits/src/recursion/compress.rs similarity index 97% rename from codex-plonky2-circuits/src/recursion/uniform/compress.rs rename to codex-plonky2-circuits/src/recursion/compress.rs index 4c38e8c..62f1fb6 100644 --- a/codex-plonky2-circuits/src/recursion/uniform/compress.rs +++ b/codex-plonky2-circuits/src/recursion/compress.rs @@ -10,6 +10,7 @@ use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; use crate::{error::CircuitError,Result}; use crate::circuit_helper::Plonky2Circuit; +//TODO: include the flag_buckets in the public input /// recursion compression circuit /// verifies 1 inner proof and as result should shrink it #[derive(Clone, Debug)] @@ -82,7 +83,7 @@ impl< let inner_pub_input = vir_proof.public_inputs.clone(); // take the public input from inner proof & make it public - assert_eq!(inner_pub_input.len(), 8); + assert!(inner_pub_input.len() >= 8); if register_pi { builder.register_public_inputs(&inner_pub_input[0..4]); } diff --git a/codex-plonky2-circuits/src/recursion/dummy_gen.rs b/codex-plonky2-circuits/src/recursion/dummy_gen.rs new file mode 100644 index 0000000..43b32ad --- /dev/null +++ b/codex-plonky2-circuits/src/recursion/dummy_gen.rs @@ -0,0 +1,105 @@ +use std::marker::PhantomData; +use plonky2::plonk::circuit_data::{CircuitData, CommonCircuitData, VerifierCircuitData, VerifierOnlyCircuitData}; +use plonky2::plonk::proof::{ProofWithPublicInputs}; +use plonky2::recursion::dummy_circuit::{dummy_proof}; +use hashbrown::HashMap; +use plonky2::gates::constant::ConstantGate; +use plonky2::gates::noop::NoopGate; +use plonky2::hash::hash_types::{ RichField}; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use plonky2::plonk::config::{AlgebraicHasher, GenericConfig}; +use plonky2_field::extension::Extendable; +use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; +use crate::{error::CircuitError, Result}; + +/// A generator for creating dummy proofs and verifier data. +pub struct DummyProofGen< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, +> + where + >::Hasher: AlgebraicHasher, +{ + phantom_data: PhantomData<(F, C)>, +} + +impl DummyProofGen + where + F: RichField + Extendable + Poseidon2, + C: GenericConfig, + >::Hasher: AlgebraicHasher, +{ + /// Builds a dummy circuit from the provided common circuit data. + pub fn gen_dummy_circ_data( + common_data: &CommonCircuitData, + ) -> CircuitData { + dummy_circuit::(common_data) + } + + /// Extracts the verifier-only data from the dummy circuit. + pub fn gen_dummy_verifier_data( + common_data: &CommonCircuitData, + ) -> VerifierOnlyCircuitData { + Self::gen_dummy_circ_data(common_data).verifier_only + } + + /// Generates a dummy proof and returns it along with the verifier-only data. + /// The `nonzero_public_inputs` argument allows you to supply a mapping of public input indices + /// to nonzero field element values. + pub fn gen_dummy_proof_and_vd_with_pi( + common_data: &CommonCircuitData, + nonzero_public_inputs: HashMap, + ) -> Result<(ProofWithPublicInputs, VerifierCircuitData)> { + let circuit_data = Self::gen_dummy_circ_data(common_data); + let proof = dummy_proof::(&circuit_data, nonzero_public_inputs) + .map_err(|e| CircuitError::ProofGenerationError(e.to_string()))?; + Ok((proof, circuit_data.verifier_data())) + } + + /// Generates a dummy proof and verifier data with zero public inputs. + pub fn gen_dummy_proof_and_vd_zero_pi( + common_data: &CommonCircuitData, + ) -> Result<(ProofWithPublicInputs, VerifierCircuitData)> { + Self::gen_dummy_proof_and_vd_with_pi(common_data, HashMap::new()) + } +} + +/// Generate a circuit matching a given `CommonCircuitData`. +/// This extends the Plonky2 one with Poseidon2 support +pub(crate) fn dummy_circuit + Poseidon2, C: GenericConfig, const D: usize>( + common_data: &CommonCircuitData, +) -> CircuitData { + let config = common_data.config.clone(); + assert!( + !common_data.config.zero_knowledge, + "Degree calculation can be off if zero-knowledge is on." + ); + + // Number of `NoopGate`s to add to get a circuit of size `degree` in the end. + // Need to account for public input hashing, a `PublicInputGate` and a `ConstantGate`. + let degree = common_data.degree(); + let num_noop_gate = degree - common_data.num_public_inputs.div_ceil(8) - 2; + + let mut builder = CircuitBuilder::::new(config.clone()); + + // add a ConstantGate + builder.add_gate( + ConstantGate::new(config.num_constants), + vec![], + ); + + for _ in 0..num_noop_gate { + builder.add_gate(NoopGate, vec![]); + } + for gate in &common_data.gates { + builder.add_gate_to_gate_set(gate.clone()); + } + for _ in 0..common_data.num_public_inputs { + builder.add_virtual_public_input(); + } + + let circuit = builder.build::(); + assert_eq!(&circuit.common, common_data); + circuit +} \ No newline at end of file diff --git a/codex-plonky2-circuits/src/recursion/leaf.rs b/codex-plonky2-circuits/src/recursion/leaf.rs new file mode 100644 index 0000000..890d570 --- /dev/null +++ b/codex-plonky2-circuits/src/recursion/leaf.rs @@ -0,0 +1,269 @@ +use std::marker::PhantomData; +use plonky2::hash::hash_types::{HashOut, RichField}; +use plonky2::iop::target::{BoolTarget, Target}; +use plonky2::iop::witness::{PartialWitness, WitnessWrite}; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use plonky2::plonk::circuit_data::{CommonCircuitData, VerifierOnlyCircuitData}; +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::{error::CircuitError,Result}; +use crate::circuit_helper::Plonky2Circuit; +use crate::recursion::dummy_gen::DummyProofGen; +use crate::recursion::utils::{bucket_count, compute_flag_buckets}; + +pub const BUCKET_SIZE: usize = 32; + +/// recursion leaf circuit - verifies N inner proof +/// N: number of inner proofs +/// T: total number of sampling proofs +#[derive(Clone, Debug)] +pub struct LeafCircuit< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, + const N: usize, + const T: usize, +> where + >::Hasher: AlgebraicHasher +{ + inner_common_data: CommonCircuitData, + inner_verifier_data: VerifierOnlyCircuitData, + phantom_data: PhantomData +} + +/// recursion leaf targets +/// inner_proof: inner (sampling) proofs +/// index: index of the node +/// flags: boolean target for each flag/signal for switching between real and dummy leaf proof +#[derive(Clone, Debug)] +pub struct LeafTargets < + const D: usize, +>{ + pub inner_proof: Vec>, + pub index: Target, // public input + // TODO: change this to vec of size N so that one flag per inner-proof + pub flag: BoolTarget, +} + +#[derive(Clone, Debug)] +pub struct LeafInput< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, +>{ + pub inner_proof: Vec>, + pub flag: bool, + pub index: usize +} + +impl< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, + const N: usize, + const T: usize, +> LeafCircuit where + >::Hasher: AlgebraicHasher +{ + pub fn new( + inner_common_data: CommonCircuitData, + inner_verifier_data: VerifierOnlyCircuitData, + ) -> Self { + Self { + inner_common_data, + inner_verifier_data, + phantom_data: PhantomData::default(), + } + } + +} + +impl< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, + const N: usize, + const T: usize, +> Plonky2Circuit for LeafCircuit where + >::Hasher: AlgebraicHasher +{ + type Targets = LeafTargets; + type Input = LeafInput; + + fn add_targets(&self, builder: &mut CircuitBuilder, register_pi: bool) -> Result> { + + let inner_common = self.inner_common_data.clone(); + let n_bucket: usize = bucket_count(T); + + // the proof virtual targets + let mut pub_input = vec![]; + let mut vir_proofs = vec![]; + for _i in 0..N { + 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); + if register_pi { + builder.register_public_inputs(&hash_inner_pub_input.elements); + } + + // pad the public input with constants so that it shares the same structure as the node + let zero_hash = builder.constant_hash(HashOut::::default()); + if register_pi { + builder.register_public_inputs(&zero_hash.elements); + } + + // virtual constant target for the verifier data + let const_verifier_data = builder.constant_verifier_data(&self.inner_verifier_data); + + // virtual constant target for dummy verifier data + let const_dummy_vd = builder.constant_verifier_data( + &DummyProofGen::::gen_dummy_verifier_data(&self.inner_common_data) + ); + + // index: 0 <= index < T where T = total number of proofs + let index = builder.add_virtual_public_input(); + let flag = builder.add_virtual_bool_target_safe(); + + // Instead of taking flag_buckets as external public inputs, + // compute them internally from the index and flag. + let computed_flag_buckets = compute_flag_buckets(builder, index, flag, BUCKET_SIZE, n_bucket)?; + // Then, for example, you could register these outputs as part of your public input vector: + if register_pi { + builder.register_public_inputs(&computed_flag_buckets); + } + + // verify the proofs in-circuit based on the + // true -> real proof, false -> dummy proof + let selected_vd = builder.select_verifier_data(flag.clone(), &const_verifier_data, &const_dummy_vd); + for i in 0..N { + builder.verify_proof::(&vir_proofs[i], &selected_vd, &inner_common); + } + + // Make sure we have every gate to match `common_data`. + for g in &inner_common.gates { + builder.add_gate_to_gate_set(g.clone()); + } + + // return targets + let t = LeafTargets { + inner_proof: vir_proofs, + index, + flag, + }; + Ok(t) + + } + + fn assign_targets( + &self, pw: &mut PartialWitness, + targets: &Self::Targets, + input: &Self::Input, + ) -> Result<()> { + assert_eq!(input.inner_proof.len(), N); + assert!(input.index <= T && input.index >= 0, "given index is not valid"); + // assign the proofs + for i in 0..N { + pw.set_proof_with_pis_target(&targets.inner_proof[i], &input.inner_proof[i]) + .map_err(|e| { + CircuitError::ProofTargetAssignmentError("inner-proof".to_string(), e.to_string()) + })?; + } + + // Assign the global index. + pw.set_target(targets.index, F::from_canonical_u64(input.index as u64)) + .map_err(|e| CircuitError::TargetAssignmentError(format!("index {}", input.index),e.to_string()))?; + // Assign the flag/condition for real/fake inner proof. + pw.set_bool_target(targets.flag, input.flag) + .map_err(|e| CircuitError::TargetAssignmentError(format!("flag {}", input.flag), e.to_string()))?; + + Ok(()) + } + +} + +#[cfg(test)] +mod tests { + use super::*; + use plonky2::plonk::config::PoseidonGoldilocksConfig; + use plonky2::plonk::circuit_builder::CircuitBuilder; + use plonky2::plonk::config::GenericConfig; + use plonky2_field::types::{Field, PrimeField64}; + use plonky2::plonk::circuit_data::{CircuitConfig, CommonCircuitData}; + use plonky2_poseidon2::poseidon2_hash::poseidon2::{Poseidon2, Poseidon2Hash}; + + // For our tests, we define: + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + type H = Poseidon2Hash; + + /// A helper to build a minimal circuit and return the common circuit data. + fn dummy_common_circuit_data() -> CommonCircuitData { + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + // Add one virtual public input so that the circuit has minimal structure. + builder.add_virtual_public_input(); + let circuit = builder.build::(); + circuit.common.clone() + } + + // ---------------------------------------------------------------------------- + // End-to-End test for the entire leaf circuit. + #[test] + fn test_full_leaf_circuit() -> anyhow::Result<()> { + const N: usize = 1; + + // get inner common + let common_data = dummy_common_circuit_data(); + + // Generate a dummy inner proof for the leaf using DummyProofGen + let (dummy_inner_proof, vd) = DummyProofGen::::gen_dummy_proof_and_vd_zero_pi(&common_data)?; + let dummy_verifier_data = DummyProofGen::::gen_dummy_verifier_data(&common_data); + + // the leaf circuit. + let leaf = LeafCircuit::::new(common_data.clone(), dummy_verifier_data); + + // Build the leaf circuit. + let (targets, circuit_data) = leaf.build_with_standard_config()?; + let verifier_data = circuit_data.verifier_data(); + let prover_data = circuit_data.prover_data(); + + // test leaf input + let input = LeafInput { + inner_proof: vec![dummy_inner_proof], + flag: true, + index: 45, + }; + + let proof = leaf.prove(&targets, &input, &prover_data)?; + + // Verify the proof. + assert!(verifier_data.verify(proof.clone()).is_ok(), "Proof verification failed"); + + println!("Public inputs: {:?}", proof.public_inputs); + + // the flag buckets appeared at positions 8..12. + let flag_buckets: Vec = proof.public_inputs[9..13] + .iter() + .map(|f| f.to_canonical_u64()) + .collect(); + + // With index = 45, we expect bucket[1] = 2^13 = 8192, and the rest 0. + let expected = vec![0, 8192, 0, 0]; + assert_eq!(flag_buckets, expected, "Flag bucket values mismatch"); + + Ok(()) + } +} + + diff --git a/codex-plonky2-circuits/src/recursion/mod.rs b/codex-plonky2-circuits/src/recursion/mod.rs index d4f1b3b..d2b82f5 100644 --- a/codex-plonky2-circuits/src/recursion/mod.rs +++ b/codex-plonky2-circuits/src/recursion/mod.rs @@ -1,2 +1,7 @@ +pub mod dummy_gen; pub mod utils; -pub mod uniform; +pub mod leaf; +pub mod node; +pub mod tree; +pub mod compress; +pub mod pi_verifier; diff --git a/codex-plonky2-circuits/src/recursion/node.rs b/codex-plonky2-circuits/src/recursion/node.rs new file mode 100644 index 0000000..43d7452 --- /dev/null +++ b/codex-plonky2-circuits/src/recursion/node.rs @@ -0,0 +1,370 @@ +use std::marker::PhantomData; +use plonky2::hash::hash_types::RichField; +use plonky2::iop::target::{BoolTarget, Target}; +use plonky2::iop::witness::{PartialWitness, WitnessWrite}; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use plonky2::plonk::circuit_data::{CommonCircuitData, VerifierCircuitTarget, VerifierOnlyCircuitData}; +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::{error::CircuitError,Result}; +use crate::circuit_helper::Plonky2Circuit; +use crate::recursion::leaf::BUCKET_SIZE; +use crate::recursion::dummy_gen::DummyProofGen; +use crate::recursion::utils::bucket_count; + +/// recursion node circuit +/// M: number of leaf proofs +/// T: total number of sampling proofs +#[derive(Clone, Debug)] +pub struct NodeCircuit< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, + const M: usize, + const T: usize, +> where + >::Hasher: AlgebraicHasher +{ + common_data: CommonCircuitData, + leaf_verifier_data: VerifierOnlyCircuitData, + phantom_data: PhantomData +} + +/// recursion node targets +/// leaf_proofs: leaf proofs +/// node_verifier_data: node verifier data, note: leaf verifier data is constant +/// condition: for switching between leaf and node verifier data +/// index: index of the node +/// flags: boolean target for each flag/signal for switching between real and dummy leaf proof +#[derive(Clone, Debug)] +pub struct NodeTargets< + const D: usize, +>{ + pub leaf_proofs: Vec>, + pub node_verifier_data: VerifierCircuitTarget, + pub condition: BoolTarget, + pub index: Target, + pub flags: Vec, +} + +#[derive(Clone, Debug)] +pub struct NodeInput< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, +>{ + pub node_proofs: Vec>, + pub verifier_only_data: VerifierOnlyCircuitData, + pub condition: bool, + pub flags: Vec, + pub index: usize +} + +impl< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, + const M: usize, + const T: usize, +> NodeCircuit where + >::Hasher: AlgebraicHasher +{ + + pub fn new( + common_data: CommonCircuitData, + leaf_verifier_data: VerifierOnlyCircuitData, + ) -> Self { + assert!(M.is_power_of_two(), "M is NOT a power of two"); + Self{ + common_data, + leaf_verifier_data, + phantom_data:PhantomData::default(), + } + } + +} + +impl< + F: RichField + Extendable + Poseidon2, + const D: usize, + C: GenericConfig, + H: AlgebraicHasher, + const M: usize, + const T: usize, +> Plonky2Circuit for NodeCircuit where + >::Hasher: AlgebraicHasher +{ + type Targets = NodeTargets; + type Input = NodeInput; + + fn add_targets(&self, builder: &mut CircuitBuilder, register_pi: bool) -> Result { + let inner_common = self.common_data.clone(); + let zero_target = builder.zero(); + + // assert public input is of size 8 + 1 (index) + B (flag buckets) + let n_bucket: usize = bucket_count(T); + assert_eq!(inner_common.num_public_inputs, 9+n_bucket); + + // the proof virtual targets - M proofs + let mut vir_proofs = vec![]; + let mut pub_input = vec![]; + let mut inner_flag_buckets = vec![]; + let mut inner_indexes = vec![]; + 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); + pub_input.extend_from_slice(&inner_pub_input[0..4]); + inner_indexes.push(inner_pub_input[8]); + inner_flag_buckets.push(inner_pub_input[9..(9+n_bucket)].to_vec()); + } + + // hash the public input & make it public + let hash_inner_pub_input = builder.hash_n_to_hash_no_pad::(pub_input); + if register_pi{ + builder.register_public_inputs(&hash_inner_pub_input.elements); + } + + // virtual target for the verifier data + let node_verifier_data = builder.add_virtual_verifier_data(inner_common.config.fri_config.cap_height); + + // virtual target for the verifier data + let const_leaf_verifier_data = builder.constant_verifier_data(&self.leaf_verifier_data); + + // virtual constant target for dummy verifier data + let const_dummy_vd = builder.constant_verifier_data( + &DummyProofGen::::gen_dummy_verifier_data(&self.common_data) + ); + + // register only the node verifier data hash as public input. + let mut vd_pub_input = vec![]; + vd_pub_input.extend_from_slice(&node_verifier_data.circuit_digest.elements); + for i in 0..builder.config.fri_config.num_cap_elements() { + vd_pub_input.extend_from_slice(&node_verifier_data.constants_sigmas_cap.0[i].elements); + } + let vd_hash = builder.hash_n_to_hash_no_pad::(vd_pub_input); + if register_pi { + builder.register_public_inputs(&vd_hash.elements); + } + + // condition for switching between node and leaf + let condition = builder.add_virtual_bool_target_safe(); + + // flag buckets targets + let mut flag_buckets: Vec = (0..n_bucket).map(|_i| zero_target.clone()).collect(); + // index: 0 <= index < T where T = total number of proofs + let index = builder.add_virtual_public_input(); + let flags: Vec = (0..M).map(|_i| builder.add_virtual_bool_target_safe()).collect(); + + // condition: true -> node, false -> leaf + let node_or_leaf_vd = builder.select_verifier_data(condition.clone(), &node_verifier_data, &const_leaf_verifier_data); + // verify the proofs in-circuit - M proofs + for i in 0..M { + // flag: true -> real, false -> dummy + let selected_vd = builder.select_verifier_data(flags[i].clone(), &node_or_leaf_vd, &const_dummy_vd); + builder.verify_proof::(&vir_proofs[i], &selected_vd, &inner_common); + } + + // Check flag buckets for dummy inner proofs: + // For each inner proof, if its corresponding flag `flags[i]` is false, + // then enforce that every bucket in inner_flag_buckets[i] is zero. + for i in 0..M { + let not_flag_i = builder.not(flags[i]); + let not_flag_val = not_flag_i.target; + for j in 0..n_bucket { + // Enforce: inner_flag_buckets[i][j] * (not_flag_val) = 0. + // If flag is false then not_flag_val = 1, forcing inner_flag_buckets[i][j] to be zero + let product = builder.mul(inner_flag_buckets[i][j], not_flag_val); + builder.connect(product, zero_target.clone()); + } + } + + // check inner proof indexes are correct + let m_const = builder.constant(F::from_canonical_u64(M as u64)); + let mut expected_inner_index = builder.mul(index, m_const); + for i in 0..M { + if i > 0 { + let i_const = builder.constant(F::from_canonical_u64(i as u64)); + expected_inner_index = builder.add(expected_inner_index, i_const); + } + builder.connect(expected_inner_index, inner_indexes[i]); + } + + // add flag buckets + for i in 0..flag_buckets.len(){ + for j in 0..inner_flag_buckets.len() { + flag_buckets[i] = builder.add(flag_buckets[i], inner_flag_buckets[j][i]); + } + } + // make flag buckets public + builder.register_public_inputs(&flag_buckets); + + // Make sure we have every gate + for g in &inner_common.gates { + builder.add_gate_to_gate_set(g.clone()); + } + + // return targets + let t = NodeTargets { + leaf_proofs: vir_proofs, + node_verifier_data, + condition, + index, + flags, + }; + + Ok(t) + } + + fn assign_targets(&self, pw: &mut PartialWitness, targets: &Self::Targets, input: &Self::Input) -> Result<()> { + // assert size of proofs vec + assert_eq!(input.node_proofs.len(), M); + assert_eq!(input.flags.len(), M); + assert!(input.index <= T && input.index >= 0, "given index is not valid"); + + // assign the proofs + for i in 0..M { + 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.node_verifier_data, &input.verifier_only_data) + .map_err(|e| { + CircuitError::VerifierDataTargetAssignmentError(e.to_string()) + })?; + + // assign the condition - for switching between leaf & node + pw.set_bool_target(targets.condition, input.condition) + .map_err(|e| CircuitError::BoolTargetAssignmentError("condition".to_string(), e.to_string()))?; + + // Assign the global index. + pw.set_target(targets.index, F::from_canonical_u64(input.index as u64)) + .map_err(|e| CircuitError::TargetAssignmentError(format!("index {}", input.index),e.to_string()))?; + // Assign the flags - switch between real & fake proof + for i in 0..M { + pw.set_bool_target(targets.flags[i], input.flags[i]) + .map_err(|e| CircuitError::TargetAssignmentError(format!("flag {}", input.flags[i]), e.to_string()))?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use plonky2::plonk::config::PoseidonGoldilocksConfig; + use plonky2::plonk::circuit_builder::CircuitBuilder; + use plonky2::plonk::config::GenericConfig; + use plonky2_field::types::{Field, PrimeField64}; + use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, CommonCircuitData, VerifierCircuitData}; + use plonky2_poseidon2::poseidon2_hash::poseidon2::{Poseidon2, Poseidon2Hash}; + use crate::recursion::leaf::{LeafCircuit, LeafInput}; + use crate::recursion::dummy_gen::DummyProofGen; + + // For our tests, we define: + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + type H = Poseidon2Hash; + + /// A helper to build a minimal leaf circuit (with 9+B public inputs) + /// and return the circuit data and targets + fn dummy_leaf() -> (CircuitData, Vec) { + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + let mut pub_input = vec![]; + for i in 0..9+B { + pub_input.push(builder.add_virtual_public_input()); + } + let data = builder.build::(); + (data, pub_input) + } + + /// A helper to generate test leaf proofs with given data, targets, and indices. + fn dummy_leaf_proofs(data: CircuitData, pub_input: Vec, indices: Vec) -> Vec> { + let mut proofs = vec![]; + for k in 0..indices.len() { + let mut pw = PartialWitness::new(); + for i in 0..8 { + pw.set_target(pub_input[i], F::ZERO).expect("assign error"); + } + pw.set_target(pub_input[8], F::from_canonical_u64(indices[k] as u64)).expect("assign error"); + let f_buckets = fill_buckets(indices[k], BUCKET_SIZE, B); + for i in 0..f_buckets.len() { + pw.set_target(pub_input[9 + i], f_buckets[i]).expect("assign error"); + } + // Run all the generators. (This method is typically called in the proving process.) + proofs.push(data.prove(pw).expect("prove failed")); + } + proofs + } + + /// helper: returns the flag buckets with the single bit at given `index` set to true `1` + fn fill_buckets(index: usize, bucket_size: usize, num_buckets: usize) -> Vec{ + assert!(index < bucket_size * num_buckets, "Index out of range"); + + let q = index / bucket_size; // bucket index + let r = index % bucket_size; // bucket bit + + let mut buckets = vec![F::ZERO; num_buckets]; + // Set the selected bucket to 2^r. + buckets[q] = F::from_canonical_u64(1 << r); + buckets + } + + /// End-to-End test for the entire node circuit. + #[test] + fn test_full_node_circuit() -> anyhow::Result<()> { + const M: usize = 2; + const B: usize = 4; // bucket size + const T: usize = 2; + + let (leaf_data, leaf_pi) = dummy_leaf::(); + let leaf_vd = leaf_data.verifier_data(); + + let indices = vec![0,1]; + let leaf_proofs = dummy_leaf_proofs::(leaf_data,leaf_pi,indices); + + let node = NodeCircuit::::new(leaf_vd.common.clone(), leaf_vd.verifier_only.clone()); + + // Build the node circuit. + let (targets, circuit_data) = node.build_with_standard_config()?; + let verifier_data = circuit_data.verifier_data(); + let prover_data = circuit_data.prover_data(); + + // node input + let input = NodeInput { + node_proofs: leaf_proofs, + verifier_only_data: leaf_vd.verifier_only.clone(), + condition: false, + flags: vec![true, true], + index: 0, + }; + + let proof = node.prove(&targets, &input, &prover_data)?; + + // Verify the proof. + assert!(verifier_data.verify(proof.clone()).is_ok(), "Proof verification failed"); + + println!("Public inputs: {:?}", proof.public_inputs); + + // the flag buckets appeared at positions 8..12. + let flag_buckets: Vec = proof.public_inputs[9..(9+B)] + .iter() + .map(|f| f.to_canonical_u64()) + .collect(); + + // With index = 45, we expect bucket 1 = 2^13 = 8192, and the rest 0. + let expected = vec![3, 0, 0, 0]; + assert_eq!(flag_buckets, expected, "Flag bucket values mismatch"); + + Ok(()) + } +} diff --git a/codex-plonky2-circuits/src/recursion/uniform/pi_verifier.rs b/codex-plonky2-circuits/src/recursion/pi_verifier.rs similarity index 99% rename from codex-plonky2-circuits/src/recursion/uniform/pi_verifier.rs rename to codex-plonky2-circuits/src/recursion/pi_verifier.rs index 790f5c2..3dd8415 100644 --- a/codex-plonky2-circuits/src/recursion/uniform/pi_verifier.rs +++ b/codex-plonky2-circuits/src/recursion/pi_verifier.rs @@ -11,6 +11,7 @@ use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; use crate::{error::CircuitError, Result}; use crate::circuit_helper::Plonky2Circuit; +// TODO: include the flag_buckets in the public input /// A circuit that verifies the aggregated public inputs from inner circuits. /// /// - `N`: Number of inner-proofs aggregated at the leaf level. diff --git a/codex-plonky2-circuits/src/recursion/uniform/tree.rs b/codex-plonky2-circuits/src/recursion/tree.rs similarity index 73% rename from codex-plonky2-circuits/src/recursion/uniform/tree.rs rename to codex-plonky2-circuits/src/recursion/tree.rs index c265cd8..1a8eb57 100644 --- a/codex-plonky2-circuits/src/recursion/uniform/tree.rs +++ b/codex-plonky2-circuits/src/recursion/tree.rs @@ -7,11 +7,12 @@ use plonky2::plonk::proof::ProofWithPublicInputs; use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; use plonky2_field::extension::Extendable; use crate::{error::CircuitError, Result}; +use crate::bundle::Bundle; use crate::circuit_helper::Plonky2Circuit; -use crate::recursion::uniform::{leaf::{LeafTargets,LeafCircuit},node::{NodeTargets,NodeCircuit}}; -use crate::recursion::uniform::compress::{CompressionCircuit, CompressionInput, CompressionTargets}; -use crate::recursion::uniform::leaf::LeafInput; -use crate::recursion::uniform::node::NodeInput; +use crate::recursion::{leaf::{LeafTargets, LeafCircuit}, node::{NodeTargets, NodeCircuit}}; +use crate::recursion::compress::{CompressionCircuit, CompressionInput, CompressionTargets}; +use crate::recursion::leaf::LeafInput; +use crate::recursion::node::NodeInput; /// tree recursion /// - `N`: Number of inner-proofs aggregated at the leaf level. @@ -23,11 +24,12 @@ pub struct TreeRecursion< H: AlgebraicHasher, const N: usize, const M: usize, + const T: usize, > where >::Hasher: AlgebraicHasher { - leaf: LeafCircuit, - node: NodeCircuit, + leaf: LeafCircuit, + node: NodeCircuit, compression: CompressionCircuit, leaf_circ_data: CircuitData, node_circ_data: CircuitData, @@ -45,7 +47,8 @@ impl< H: AlgebraicHasher, const N: usize, const M: usize, -> TreeRecursion where + const T: usize, +> TreeRecursion where >::Hasher: AlgebraicHasher { /// build with standard recursion config @@ -67,20 +70,20 @@ impl< config: CircuitConfig, ) -> Result { // build leaf with standard recursion config - let leaf = LeafCircuit::<_,D,_,_,N>::new(inner_common_data, inner_verifier_data); + let leaf = LeafCircuit::<_,D,_,_,N, T>::new(inner_common_data, inner_verifier_data); let (leaf_targets, leaf_circ_data) = leaf.build(config.clone())?; - // println!("leaf circuit size = {:?}", leaf_circ_data.common.degree_bits()); + println!("leaf circuit size = {:?}", leaf_circ_data.common.degree_bits()); // build node with standard recursion config - let node = NodeCircuit::<_,D,_,_,M>::new(leaf_circ_data.common.clone(), leaf_circ_data.verifier_only.clone()); + let node = NodeCircuit::<_,D,_,_,M, T>::new(leaf_circ_data.common.clone(), leaf_circ_data.verifier_only.clone()); let (node_targets, node_circ_data) = node.build(config.clone())?; - // println!("node circuit size = {:?}", node_circ_data.common.degree_bits()); + println!("node circuit size = {:?}", node_circ_data.common.degree_bits()); // compression build let node_common = node_circ_data.common.clone(); let compression_circ = CompressionCircuit::new(node_common, node_circ_data.verifier_only.clone()); let (compression_targets, compression_circ_data) = compression_circ.build(config.clone())?; - // println!("compress circuit size = {:?}", compression_circ_data.common.degree_bits()); + println!("compress circuit size = {:?}", compression_circ_data.common.degree_bits()); Ok(Self{ leaf, @@ -112,6 +115,10 @@ impl< self.node_circ_data.verifier_data() } + pub fn prove_bundle(bundle: Bundle){ + + } + pub fn prove_tree_and_compress( &mut self, proofs_with_pi: &[ProofWithPublicInputs], @@ -163,9 +170,11 @@ impl< let mut leaf_proofs = vec![]; - for proofs in proofs_with_pi.chunks(N){ + for (i, proofs) in proofs_with_pi.chunks(N).enumerate(){ let leaf_input = LeafInput{ inner_proof: proofs.to_vec(), + flag: true, + index: i, }; let mut pw = PartialWitness::::new(); @@ -196,7 +205,7 @@ impl< let condition = if level == 0 {false} else {true}; - for chunk in proofs_with_pi.chunks(M) { + for (i, chunk) in proofs_with_pi.chunks(M).enumerate() { let mut inner_pw = PartialWitness::new(); @@ -204,6 +213,8 @@ impl< node_proofs: chunk.to_vec().clone(), verifier_only_data: verifier_only_data.clone(), condition, + flags: [true; M].to_vec(), + index: i, }; self.node.assign_targets( @@ -257,7 +268,7 @@ impl< public_input: Vec, inner_public_input: Vec>, ) -> Result<()>{ - assert_eq!(public_input.len(), 8); + assert!(public_input.len() >= 8); let given_input_hash = &public_input[0..4]; let given_vd_hash = &public_input[4..8]; @@ -317,3 +328,63 @@ pub fn get_hash_of_verifier_data< H::hash_no_pad(&vd) } + +#[cfg(test)] +mod tests { + use plonky2::gates::noop::NoopGate; + use super::*; + use plonky2::plonk::config::PoseidonGoldilocksConfig; + use plonky2::plonk::circuit_builder::CircuitBuilder; + use plonky2::plonk::config::GenericConfig; + use plonky2_field::types::Field; + use plonky2::plonk::circuit_data::CircuitConfig; + use plonky2_poseidon2::poseidon2_hash::poseidon2::{Poseidon2, Poseidon2Hash}; + use plonky2::iop::witness::WitnessWrite; + + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + type H = Poseidon2Hash; + + // A helper to build a minimal circuit and returns T proofs & circuit data. + fn dummy_proofs() -> (CircuitData, Vec>) { + let config = CircuitConfig::standard_recursion_config(); + let mut builder = CircuitBuilder::::new(config); + for _ in 0..8192 { + builder.add_gate(NoopGate, vec![]); + } + // Add one virtual public input so that the circuit has minimal structure. + let t = builder.add_virtual_public_input(); + let circuit = builder.build::(); + println!("inner circuit size = {}", circuit.common.degree_bits()); + let mut pw = PartialWitness::::new(); + pw.set_target(t, F::ZERO).expect("faulty assign"); + let proofs = (0..T).map(|_i| circuit.prove(pw.clone()).unwrap()).collect(); + (circuit, proofs) + } + + + // End-to-End test for the entire Tree circuit. + #[test] + fn test_full_tree_circuit() -> anyhow::Result<()> { + const N: usize = 1; + const M: usize = 2; + const T: usize = 4; + + let (data, proofs) = dummy_proofs::(); + + let mut tree = TreeRecursion::::build_with_standard_config(data.common.clone(), data.verifier_only.clone())?; + + // aggregate - no compression + let root = tree.prove_tree(&proofs)?; + println!("pub input size = {}", root.public_inputs.len()); + println!("proof size = {:?} bytes", root.to_bytes().len()); + + assert!( + tree.verify_proof(root, false).is_ok(), + "proof verification failed" + ); + + Ok(()) + } +} \ No newline at end of file diff --git a/codex-plonky2-circuits/src/recursion/uniform/leaf.rs b/codex-plonky2-circuits/src/recursion/uniform/leaf.rs deleted file mode 100644 index 111f005..0000000 --- a/codex-plonky2-circuits/src/recursion/uniform/leaf.rs +++ /dev/null @@ -1,139 +0,0 @@ -use std::marker::PhantomData; -use plonky2::hash::hash_types::{HashOut, RichField}; -use plonky2::iop::witness::{PartialWitness, WitnessWrite}; -use plonky2::plonk::circuit_builder::CircuitBuilder; -use plonky2::plonk::circuit_data::{CommonCircuitData, VerifierOnlyCircuitData}; -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::{error::CircuitError,Result}; -use crate::circuit_helper::Plonky2Circuit; - -/// recursion leaf circuit - verifies N inner proof -#[derive(Clone, Debug)] -pub struct LeafCircuit< - F: RichField + Extendable + Poseidon2, - const D: usize, - C: GenericConfig, - H: AlgebraicHasher, - const N: usize, -> where - >::Hasher: AlgebraicHasher -{ - inner_common_data: CommonCircuitData, - inner_verifier_data: VerifierOnlyCircuitData, - phantom_data: PhantomData -} - -#[derive(Clone, Debug)] -pub struct LeafTargets < - const D: usize, ->{ - pub inner_proof: Vec>, -} - -#[derive(Clone, Debug)] -pub struct LeafInput< - F: RichField + Extendable + Poseidon2, - const D: usize, - C: GenericConfig, ->{ - pub inner_proof: Vec> -} - -impl< - F: RichField + Extendable + Poseidon2, - const D: usize, - C: GenericConfig, - H: AlgebraicHasher, - const N: usize, -> LeafCircuit where - >::Hasher: AlgebraicHasher -{ - pub fn new( - inner_common_data: CommonCircuitData, - inner_verifier_data: VerifierOnlyCircuitData, - ) -> Self { - Self { - inner_common_data, - inner_verifier_data, - phantom_data: PhantomData::default(), - } - } -} - -impl< - F: RichField + Extendable + Poseidon2, - const D: usize, - C: GenericConfig, - H: AlgebraicHasher, - const N: usize, -> Plonky2Circuit for LeafCircuit where - >::Hasher: AlgebraicHasher -{ - type Targets = LeafTargets; - type Input = LeafInput; - - fn add_targets(&self, builder: &mut CircuitBuilder, register_pi: bool) -> Result> { - - let inner_common = self.inner_common_data.clone(); - - // the proof virtual targets - let mut pub_input = vec![]; - let mut vir_proofs = vec![]; - for _i in 0..N { - 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); - if register_pi { - builder.register_public_inputs(&hash_inner_pub_input.elements); - } - - // pad the public input with constants so that it shares the same structure as the node - let zero_hash = builder.constant_hash(HashOut::::default()); - if register_pi { - builder.register_public_inputs(&zero_hash.elements); - } - - // virtual constant target for the verifier data - let const_verifier_data = builder.constant_verifier_data(&self.inner_verifier_data); - - // verify the proofs in-circuit - for i in 0..N { - builder.verify_proof::(&vir_proofs[i], &const_verifier_data, &inner_common); - } - - // return targets - let t = LeafTargets { - inner_proof: vir_proofs, - }; - Ok(t) - - } - - fn assign_targets( - &self, pw: &mut PartialWitness, - targets: &Self::Targets, - input: &Self::Input, - ) -> Result<()> { - assert_eq!(input.inner_proof.len(), N); - // assign the proofs - for i in 0..N { - pw.set_proof_with_pis_target(&targets.inner_proof[i], &input.inner_proof[i]) - .map_err(|e| { - CircuitError::ProofTargetAssignmentError("inner-proof".to_string(), e.to_string()) - })?; - } - - Ok(()) - } - -} - - diff --git a/codex-plonky2-circuits/src/recursion/uniform/mod.rs b/codex-plonky2-circuits/src/recursion/uniform/mod.rs deleted file mode 100644 index 08e3905..0000000 --- a/codex-plonky2-circuits/src/recursion/uniform/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -pub mod leaf; -pub mod node; -pub mod tree; -pub mod compress; -pub mod pi_verifier; \ 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 deleted file mode 100644 index 445ba62..0000000 --- a/codex-plonky2-circuits/src/recursion/uniform/node.rs +++ /dev/null @@ -1,175 +0,0 @@ -use std::marker::PhantomData; -use plonky2::hash::hash_types::RichField; -use plonky2::iop::target::BoolTarget; -use plonky2::iop::witness::{PartialWitness, WitnessWrite}; -use plonky2::plonk::circuit_builder::CircuitBuilder; -use plonky2::plonk::circuit_data::{CommonCircuitData, VerifierCircuitTarget, VerifierOnlyCircuitData}; -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::{error::CircuitError,Result}; -use crate::circuit_helper::Plonky2Circuit; - -/// recursion node circuit - verifies M leaf proofs -#[derive(Clone, Debug)] -pub struct NodeCircuit< - F: RichField + Extendable + Poseidon2, - const D: usize, - C: GenericConfig, - H: AlgebraicHasher, - const M: usize, -> where - >::Hasher: AlgebraicHasher -{ - common_data: CommonCircuitData, - leaf_verifier_data: VerifierOnlyCircuitData, - phantom_data: PhantomData -} - -#[derive(Clone, Debug)] -pub struct NodeTargets< - const D: usize, ->{ - pub leaf_proofs: Vec>, - pub node_verifier_data: VerifierCircuitTarget, - pub condition: BoolTarget, -} - -#[derive(Clone, Debug)] -pub struct NodeInput< - F: RichField + Extendable + Poseidon2, - const D: usize, - C: GenericConfig, ->{ - pub node_proofs: Vec>, - pub verifier_only_data: VerifierOnlyCircuitData, - pub condition: bool, -} - -impl< - F: RichField + Extendable + Poseidon2, - const D: usize, - C: GenericConfig, - H: AlgebraicHasher, - const M: usize, -> NodeCircuit where - >::Hasher: AlgebraicHasher -{ - - pub fn new( - common_data: CommonCircuitData, - leaf_verifier_data: VerifierOnlyCircuitData, - ) -> Self { - Self{ - common_data, - leaf_verifier_data, - phantom_data:PhantomData::default(), - } - } - -} - -impl< - F: RichField + Extendable + Poseidon2, - const D: usize, - C: GenericConfig, - H: AlgebraicHasher, - const M: usize, -> Plonky2Circuit for NodeCircuit where - >::Hasher: AlgebraicHasher -{ - type Targets = NodeTargets; - type Input = NodeInput; - - fn add_targets(&self, builder: &mut CircuitBuilder, register_pi: bool) -> Result { - let inner_common = self.common_data.clone(); - - // assert public input is of size 8 - 2 hashout - assert_eq!(inner_common.num_public_inputs, 8); - - // the proof virtual targets - M proofs - let mut vir_proofs = vec![]; - let mut pub_input = vec![]; - 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); - pub_input.extend_from_slice(&inner_pub_input[0..4]); - } - - // hash the public input & make it public - let hash_inner_pub_input = builder.hash_n_to_hash_no_pad::(pub_input); - if register_pi{ - builder.register_public_inputs(&hash_inner_pub_input.elements); - } - - // virtual target for the verifier data - let node_verifier_data = builder.add_virtual_verifier_data(inner_common.config.fri_config.cap_height); - - // virtual target for the verifier data - let const_leaf_verifier_data = builder.constant_verifier_data(&self.leaf_verifier_data); - - // register only the node verifier data hash as public input. - let mut vd_pub_input = vec![]; - vd_pub_input.extend_from_slice(&node_verifier_data.circuit_digest.elements); - for i in 0..builder.config.fri_config.num_cap_elements() { - vd_pub_input.extend_from_slice(&node_verifier_data.constants_sigmas_cap.0[i].elements); - } - let vd_hash = builder.hash_n_to_hash_no_pad::(vd_pub_input); - if register_pi { - builder.register_public_inputs(&vd_hash.elements); - } - - // condition for switching between node and leaf - let condition = builder.add_virtual_bool_target_safe(); - - // true -> node, false -> leaf - let selected_vd = builder.select_verifier_data(condition.clone(), &node_verifier_data, &const_leaf_verifier_data); - - // verify the proofs in-circuit - M proofs - for i in 0..M { - builder.verify_proof::(&vir_proofs[i], &selected_vd, &inner_common); - } - - // Make sure we have every gate to match `common_data`. - for g in &inner_common.gates { - builder.add_gate_to_gate_set(g.clone()); - } - - // return targets - let t = NodeTargets { - leaf_proofs: vir_proofs, - node_verifier_data, - condition, - }; - Ok(t) - } - - fn assign_targets(&self, pw: &mut PartialWitness, targets: &Self::Targets, input: &Self::Input) -> Result<()> { - // assert size of proofs vec - assert_eq!(input.node_proofs.len(), M); - - // assign the proofs - for i in 0..M { - 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.node_verifier_data, &input.verifier_only_data) - .map_err(|e| { - CircuitError::VerifierDataTargetAssignmentError(e.to_string()) - })?; - - // assign the condition - pw.set_bool_target(targets.condition, input.condition) - .map_err(|e| CircuitError::BoolTargetAssignmentError("condition".to_string(), e.to_string()))?; - - Ok(()) - } -} - - diff --git a/codex-plonky2-circuits/src/recursion/utils.rs b/codex-plonky2-circuits/src/recursion/utils.rs new file mode 100644 index 0000000..8b0d840 --- /dev/null +++ b/codex-plonky2-circuits/src/recursion/utils.rs @@ -0,0 +1,245 @@ +use plonky2::hash::hash_types::RichField; +use plonky2::iop::target::{BoolTarget, Target}; +use plonky2::plonk::circuit_builder::CircuitBuilder; +use plonky2_field::extension::Extendable; +use plonky2_poseidon2::Poseidon2; +use crate::recursion::leaf::{BUCKET_SIZE, LeafCircuit}; + +/// Splits a target `index` which is known to lie in the range [0, T) +/// where T = bucket_size * num_buckets +/// into two components (q, r) such that: +/// +/// index = q * bucket_size + r, +/// +/// where: +/// - `r` is in the range [0, bucket_size), +/// - `q` is in the range [0, num_buckets), +/// +/// requires that the total range T = (bucket_size * num_buckets) is a power of 2. +pub fn split_index< + F: RichField + Extendable + Poseidon2, + const D: usize, +>( + builder: &mut CircuitBuilder, + index: Target, + bucket_size: usize, + num_buckets: usize, +) -> crate::Result<(Target, Target)> +{ + // T = bucket_size * num_buckets + let total = bucket_size * num_buckets; + // check total is a power of two + assert!(total.is_power_of_two(), "Total must be a power of two for split_index to work."); + + let total_bits = total.trailing_zeros() as usize; + let log_bucket = bucket_size.trailing_zeros() as usize; + + // Decompose the index into total_bits bits (little-endian). + let bits: Vec = builder.split_le(index, total_bits); + + // Recompose the remainder (r) from the lower log_bucket bits. + let mut r_val = builder.zero(); + for i in 0..log_bucket { + let bit_val = bits[i].target; + let weight = builder.constant(F::from_canonical_u64(1 << i)); + let bit_mul_weight = builder.mul(bit_val, weight); + r_val = builder.add(r_val, bit_mul_weight); + } + + // Recompose the quotient (q) from the remaining log_q bits. + let mut q_val = builder.zero(); + for i in log_bucket..total_bits { + // The weight here is 2^(i - log_bucket). + let bit_val = bits[i].target; + let weight = builder.constant(F::from_canonical_u64(1 << (i - log_bucket))); + let bit_mul_weight = builder.mul(bit_val, weight); + q_val = builder.add(q_val, bit_mul_weight); + } + + Ok((q_val, r_val)) +} + +/// A helper that computes 2^r for a target r in [0, 32) using selection over 32 constants. +pub fn compute_power_of_two< + F: RichField + Extendable + Poseidon2, + const D: usize, +>( + builder: &mut CircuitBuilder, + r: Target, +) -> crate::Result +{ + // First range-check r so it is in [0, 32). + builder.range_check(r, BUCKET_SIZE); + let mut result = builder.zero(); + for i in 0..BUCKET_SIZE { + let i_const = builder.constant(F::from_canonical_u64(i as u64)); + let eq_bool = builder.is_equal(r, i_const); + let eq_val = eq_bool.target; + let two_i = builder.constant(F::from_canonical_u64(1 << i)); + let eq_val_mul_two_i = builder.mul(eq_val, two_i); + result = builder.add(result, eq_val_mul_two_i); + } + Ok(result) +} + +/// Computes the flag buckets from a given index and flag. +/// +/// Given: +/// - `index` is a Target representing a number in T = [0, bucket_size * num_buckets), +/// - `flag` is a BoolTarget (true if the proof is real, false if dummy), +/// - `bucket_size` (e.g. 32 for Goldilocks) and `num_buckets` (e.g. 4 to fit 128 proofs), +/// this function returns a vector of Targets representing the computed flag buckets. +/// For bucket i, the value is: +/// - flag * 2^(r) if i is the selected bucket (i.e. i == q), where (q, r) = split_index(index), +/// - 0 otherwise. +pub fn compute_flag_buckets< + F: RichField + Extendable + Poseidon2, + const D: usize, +>( + builder: &mut CircuitBuilder, + index: Target, + flag: BoolTarget, + bucket_size: usize, + num_buckets: usize, +) -> crate::Result> +{ + let total = bucket_size * num_buckets; + // Range-check the index. + builder.range_check(index, total); + + // Use your split_index helper to get (q, r) + let (q, r) = split_index::(builder, index, bucket_size, num_buckets)?; + // Compute 2^(r) + let power_of_two = compute_power_of_two::(builder, r)?; + // flag target from Boolean target. + let flag_val = flag.target; + // computed_value equals flag * 2^(r) + let computed_value = builder.mul(flag_val, power_of_two); + + // For each bucket, if the bucket is the selected one (i.e. equals q), then its value is computed_value; otherwise 0. + let mut buckets = Vec::with_capacity(num_buckets); + for i in 0..num_buckets { + let bucket_const = builder.constant(F::from_canonical_u64(i as u64)); + let is_selected = builder.is_equal(q, bucket_const); + let is_selected_val = is_selected.target; + // bucket value = is_selected * computed_value. + let bucket_value = builder.mul(is_selected_val, computed_value); + buckets.push(bucket_value); + } + Ok(buckets) +} + +/// Returns the number of buckets required to hold `t` flags, +/// where each bucket can hold up to BUCKET_SIZE flags. +pub fn bucket_count(t: usize) -> usize { + (t + BUCKET_SIZE -1) / BUCKET_SIZE +} + + +#[cfg(test)] +mod tests { + use super::*; + use plonky2::plonk::config::PoseidonGoldilocksConfig; + use plonky2::plonk::circuit_builder::CircuitBuilder; + use plonky2::plonk::config::GenericConfig; + use plonky2_field::types::{Field, PrimeField64}; + use plonky2::plonk::circuit_data::CircuitConfig; + use plonky2_poseidon2::poseidon2_hash::poseidon2::{Poseidon2, Poseidon2Hash}; + use plonky2::iop::witness::PartialWitness; + + // For our tests, we define: + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + type H = Poseidon2Hash; + + // Helper: Build, prove, and return public inputs --- + fn build_and_prove(builder: CircuitBuilder) -> Vec { + // Build the circuit. + let circuit = builder.build::(); + let pw = PartialWitness::new(); + // prove + let p= circuit.prove(pw).expect("prove failed"); + + p.public_inputs + } + + #[test] + fn test_split_index() -> anyhow::Result<()> { + // Create a circuit where we register the outputs q and r of split_index. + let mut builder = CircuitBuilder::::new(CircuitConfig::standard_recursion_config()); + // Let index = 45. + let index_val: u64 = 45; + let index_target = builder.constant(F::from_canonical_u64(index_val)); + // Call split_index with bucket_size=32 and num_buckets=4. We expect q = 1 and r = 13. + let (q_target, r_target) = + split_index::(&mut builder, index_target, BUCKET_SIZE, 4)?; + // Register outputs as public inputs. + builder.register_public_input(q_target); + builder.register_public_input(r_target); + // Build and prove the circuit. + let pub_inputs = build_and_prove(builder); + // We expect the first public input to be q = 1 and the second r = 13. + assert_eq!(pub_inputs[0].to_canonical_u64(), 1, "q should be 1"); + assert_eq!(pub_inputs[1].to_canonical_u64(), 13, "r should be 13"); + Ok(()) + } + + #[test] + fn test_compute_power_of_two() -> anyhow::Result<()> { + // Create a circuit to compute 2^r. + let mut builder = CircuitBuilder::::new(CircuitConfig::standard_recursion_config()); + // Let r = 13. + let r_val: u64 = 13; + let r_target = builder.constant(F::from_canonical_u64(r_val)); + let pow_target = + compute_power_of_two::(&mut builder, r_target)?; + builder.register_public_input(pow_target); + let pub_inputs = build_and_prove(builder); + // Expect 2^13 = 8192. + assert_eq!( + pub_inputs[0].to_canonical_u64(), + 1 << 13, + "2^13 should be 8192" + ); + Ok(()) + } + + #[test] + fn test_compute_flag_buckets() -> anyhow::Result<()> { + // Create a circuit to compute flag buckets. + // Let index = 45 and flag = true. + let mut builder = CircuitBuilder::::new(CircuitConfig::standard_recursion_config()); + let index_val: u64 = 45; + let index_target = builder.constant(F::from_canonical_u64(index_val)); + // Create a boolean constant target for flag = true. + let flag_target = builder.constant_bool(true); + // Compute the flag buckets with bucket_size = 32 and num_buckets = 4. + let buckets = compute_flag_buckets::( + &mut builder, + index_target, + flag_target, + BUCKET_SIZE, + 4, + )?; + // Register each bucket as a public input. + for bucket in buckets.iter() { + builder.register_public_input(*bucket); + } + let pub_inputs = build_and_prove(builder); + // With index = 45, we expect: + // q = 45 / 32 = 1 and r = 45 % 32 = 13, so bucket 1 should be 2^13 = 8192 and the others 0. + let expected = vec![0, 8192, 0, 0]; + for (i, &expected_val) in expected.iter().enumerate() { + let computed = pub_inputs[i].to_canonical_u64(); + assert_eq!( + computed, expected_val, + "Bucket {}: expected {} but got {}", + i, expected_val, computed + ); + } + Ok(()) + } +} + + diff --git a/codex-plonky2-circuits/src/recursion/utils/conditional_verifier.rs b/codex-plonky2-circuits/src/recursion/utils/conditional_verifier.rs deleted file mode 100644 index 2956e67..0000000 --- a/codex-plonky2-circuits/src/recursion/utils/conditional_verifier.rs +++ /dev/null @@ -1,121 +0,0 @@ -use plonky2::gates::noop::NoopGate; -use plonky2::hash::hash_types::RichField; -use plonky2::iop::target::BoolTarget; -use plonky2::plonk::circuit_builder::CircuitBuilder; -use plonky2::plonk::circuit_data::{CircuitData, CommonCircuitData, VerifierCircuitTarget}; -use plonky2::plonk::config::{AlgebraicHasher, GenericConfig}; -use plonky2::plonk::proof::ProofWithPublicInputsTarget; -use plonky2_field::extension::Extendable; -use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; - -/// this takes verifier data (not public) and doesn't check the verifier data for consistency -pub fn conditionally_verify_recursion_proof_or_dummy + Poseidon2, const D: usize ,C: GenericConfig + 'static>( - builder: &mut CircuitBuilder, - condition: BoolTarget, - cyclic_proof_with_pis: &ProofWithPublicInputsTarget, - verifier_data: &VerifierCircuitTarget, - common_data: &CommonCircuitData, -) -> anyhow::Result<()> - where - C::Hasher: AlgebraicHasher, -{ - let (dummy_proof_with_pis_target, dummy_verifier_data_target) = - dummy_proof_and_vk_no_generator::(builder, common_data)?; - - // TODO: make verifier data public - // // Connect previous verifier data to current one. This guarantees that every proof in the cycle uses the same verifier data. - // self.connect_hashes( - // inner_cyclic_pis.circuit_digest, - // verifier_data.circuit_digest, - // ); - // self.connect_merkle_caps( - // &inner_cyclic_pis.constants_sigmas_cap, - // &verifier_data.constants_sigmas_cap, - // ); - - // Verify the cyclic proof if `condition` is set to true, otherwise verify the other proof. - builder.conditionally_verify_proof::( - condition, - cyclic_proof_with_pis, - verifier_data, - &dummy_proof_with_pis_target, - &dummy_verifier_data_target, - common_data, - ); - - // Make sure we have every gate to match `common_data`. - for g in &common_data.gates { - builder.add_gate_to_gate_set(g.clone()); - } - - Ok(()) -} - -/// Conditionally verify a proof with a new generated dummy proof. -pub fn conditionally_verify_proof_or_dummy + Poseidon2, const D: usize ,C: GenericConfig + 'static>( - builder: &mut CircuitBuilder, - condition: BoolTarget, - proof_with_pis: &ProofWithPublicInputsTarget, - inner_verifier_data: &VerifierCircuitTarget, - inner_common_data: &CommonCircuitData, -) -> anyhow::Result<()> - where - C::Hasher: AlgebraicHasher, -{ - let (dummy_proof_with_pis_target, dummy_verifier_data_target) = - dummy_proof_and_vk_no_generator::(builder, inner_common_data)?; - builder.conditionally_verify_proof::( - condition, - proof_with_pis, - inner_verifier_data, - &dummy_proof_with_pis_target, - &dummy_verifier_data_target, - inner_common_data, - ); - Ok(()) -} - -/// Generate a circuit matching a given `CommonCircuitData`. -pub(crate) fn dummy_circuit + Poseidon2, C: GenericConfig, const D: usize>( - common_data: &CommonCircuitData, -) -> CircuitData { - let config = common_data.config.clone(); - assert!( - !common_data.config.zero_knowledge, - "Degree calculation can be off if zero-knowledge is on." - ); - - // Number of `NoopGate`s to add to get a circuit of size `degree` in the end. - // Need to account for public input hashing, a `PublicInputGate` and a `ConstantGate`. - let degree = common_data.degree(); - let num_noop_gate = degree - common_data.num_public_inputs.div_ceil(8) - 2; - - let mut builder = CircuitBuilder::::new(config); - for _ in 0..num_noop_gate { - builder.add_gate(NoopGate, vec![]); - } - for gate in &common_data.gates { - builder.add_gate_to_gate_set(gate.clone()); - } - for _ in 0..common_data.num_public_inputs { - builder.add_virtual_public_input(); - } - - let circuit = builder.build::(); - assert_eq!(&circuit.common, common_data); - circuit -} - -pub(crate) fn dummy_proof_and_vk_no_generator + Poseidon2, const D: usize ,C: GenericConfig + 'static> ( - builder: &mut CircuitBuilder, - common_data: &CommonCircuitData, -) -> anyhow::Result<(ProofWithPublicInputsTarget, VerifierCircuitTarget)> - where - C::Hasher: AlgebraicHasher, -{ - let dummy_circuit = dummy_circuit::(common_data); - let dummy_proof_with_pis_target = builder.add_virtual_proof_with_pis(common_data); - let dummy_verifier_data_target = builder.constant_verifier_data(&dummy_circuit.verifier_only); - - Ok((dummy_proof_with_pis_target, dummy_verifier_data_target)) -} diff --git a/codex-plonky2-circuits/src/recursion/utils/dummy_gen.rs b/codex-plonky2-circuits/src/recursion/utils/dummy_gen.rs deleted file mode 100644 index 7b412d6..0000000 --- a/codex-plonky2-circuits/src/recursion/utils/dummy_gen.rs +++ /dev/null @@ -1,79 +0,0 @@ -use std::marker::PhantomData; -use plonky2::plonk::circuit_data::{CommonCircuitData, VerifierOnlyCircuitData}; -use plonky2::plonk::proof::{ProofWithPublicInputs}; -use plonky2::recursion::dummy_circuit::{dummy_proof}; -use crate::recursion::utils::conditional_verifier::dummy_circuit; -use hashbrown::HashMap; -use plonky2::hash::hash_types::{ RichField}; -use plonky2::plonk::config::{AlgebraicHasher, GenericConfig}; -use plonky2_field::extension::Extendable; -use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2; -use crate::{error::CircuitError, Result}; -use crate::circuits::utils::vec_to_array; - -/// A generator for creating dummy proofs. -pub struct DummyProofGen< - F: RichField + Extendable + Poseidon2, - const D: usize, - C: GenericConfig, - > where - >::Hasher: AlgebraicHasher -{ - phantom_data: PhantomData<(F,C)>, -} - -impl< - F: RichField + Extendable + Poseidon2, - const D: usize, - C: GenericConfig, -> DummyProofGen - where - >::Hasher: AlgebraicHasher -{ - - /// Generates a single dummy leaf proof. - pub fn gen_dummy_leaf_proof( - common_data: &CommonCircuitData, - ) -> Result> { - dummy_proof::(&dummy_circuit::(common_data), HashMap::new()) - .map_err(|e| CircuitError::DummyProofGenerationError(e.to_string())) - } - - /// Generates a single dummy node proof. - pub fn get_dummy_node_proof( - node_common: &CommonCircuitData, - node_verifier_only_data: &VerifierOnlyCircuitData, - ) -> ProofWithPublicInputs { - Self::recursion_base_proof(node_common, HashMap::new()) - } - - fn recursion_base_proof( - common_data: &CommonCircuitData, - mut nonzero_public_inputs: HashMap - ) -> ProofWithPublicInputs{ - dummy_proof::( - &dummy_circuit::(common_data), - nonzero_public_inputs, - ) - .unwrap() - } - - /// Generates an array of `N` dummy leaf proofs. - pub fn gen_n_dummy_leaf_proofs( - common_data: &CommonCircuitData, - ) -> Result<[ProofWithPublicInputs; N]> { - let dummy_proof = Self::gen_dummy_leaf_proof(common_data)?; - let n_dummy_vec = (0..N).map(|_| dummy_proof.clone()).collect::>(); - vec_to_array::>(n_dummy_vec) - } - - /// Generates an array of `N` dummy node proofs. - pub fn gen_n_dummy_node_proofs( - node_common: &CommonCircuitData, - node_verifier_only_data: &VerifierOnlyCircuitData, - ) -> Result<[ProofWithPublicInputs; N]> { - let dummy_proof = Self::get_dummy_node_proof(node_common, node_verifier_only_data); - let n_dummy_vec = (0..N).map(|_| dummy_proof.clone()).collect::>(); - vec_to_array::>(n_dummy_vec) - } -} diff --git a/codex-plonky2-circuits/src/recursion/utils/mod.rs b/codex-plonky2-circuits/src/recursion/utils/mod.rs deleted file mode 100644 index f8e0bc8..0000000 --- a/codex-plonky2-circuits/src/recursion/utils/mod.rs +++ /dev/null @@ -1,2 +0,0 @@ -pub mod dummy_gen; -pub mod conditional_verifier; \ No newline at end of file