Merge pull request #1 from codex-storage/tracking-proofs

Tracking proofs

codex-plonky2-circuits/.gitignore

@@ -11,3 +11,7 @@ pgo-data.profdata
 
 # MacOS nuisances
 .DS_Store
+
+# circuit files
+/prover_data
+/verifier_data

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"]
+

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<D> + Poseidon2,
+    C: GenericConfig<D, F = F>,
+    const D: usize,
+    H: AlgebraicHasher<F>,
+>{
+    pub bundle_size: usize,
+    pub circuit: SampleCircuit<F, D, H>,
+    pub prover_data: ProverCircuitData<F, C, D>,
+    pub verifier_data: VerifierCircuitData<F, C, D>,
+    pub sample_targets: SampleTargets,
+    pub bundle_proofs: HashMap<usize, ProofWithPublicInputs<F, C, D>>,
+}
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    C: GenericConfig<D, F=F>,
+    const D: usize,
+    H: AlgebraicHasher<F>,
+> Bundle<F, C, D, H> {
+    pub fn new(bundle_size: usize, circuit_params: CircuitParams) -> Result<Self>{
+        let samp_circ = SampleCircuit::<F, D, H>::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<SampleCircuitInput<F, D>>) -> 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<F, D>) -> Result<()>{
+        let proof = self.circuit.prove(&self.sample_targets, &circ_input, &self.prover_data)?;
+        self.bundle_proofs.insert(index, proof);
+        Ok(())
+    }
+
+}
+

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<T> = core::result::Result<T, error::CircuitError>;

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

@@ -1,8 +1,7 @@
-use std::marker::PhantomData;
-use plonky2::hash::hash_types::{HashOutTarget, RichField};
+use plonky2::hash::hash_types::RichField;
 use plonky2::iop::witness::{PartialWitness, WitnessWrite};
 use plonky2::plonk::circuit_builder::CircuitBuilder;
-use plonky2::plonk::circuit_data::{CommonCircuitData, VerifierOnlyCircuitData};
+use plonky2::plonk::circuit_data::VerifierCircuitData;
 use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
 use plonky2::plonk::proof::{ProofWithPublicInputs, ProofWithPublicInputsTarget};
 use plonky2_field::extension::Extendable;
@@ -17,13 +16,10 @@ pub struct CompressionCircuit<
     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>,
-    inner_verifier_data: VerifierOnlyCircuitData<C, D>,
-    phantom_data: PhantomData<H>
+    inner_verifier_data: VerifierCircuitData<F, C, D>,
 }
 
 #[derive(Clone, Debug)]
@@ -46,18 +42,14 @@ impl<
     F: RichField + Extendable<D> + Poseidon2,
     const D: usize,
     C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
-> CompressionCircuit<F,D,C,H> where
+> CompressionCircuit<F,D,C> where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
 {
     pub fn new(
-        inner_common_data: CommonCircuitData<F,D>,
-        inner_verifier_data: VerifierOnlyCircuitData<C, D>,
+        inner_verifier_data: VerifierCircuitData<F, C, D>,
     ) -> Self {
         Self{
-            inner_common_data,
             inner_verifier_data,
-            phantom_data:PhantomData::default(),
         }
     }
 
@@ -67,44 +59,25 @@ impl<
     F: RichField + Extendable<D> + Poseidon2,
     const D: usize,
     C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
-> Plonky2Circuit<F, C, D> for CompressionCircuit<F, D, C, H> where
+> Plonky2Circuit<F, C, D> for CompressionCircuit<F, D, C> where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
 {
     type Targets = CompressionTargets<D>;
     type Input = CompressionInput<F, D, C>;
 
     fn add_targets(&self, builder: &mut CircuitBuilder<F, D>, register_pi: bool) -> Result<Self::Targets> {
-        let inner_common = self.inner_common_data.clone();
+        let inner_common = self.inner_verifier_data.common.clone();
 
         // the proof virtual targets
         let vir_proof = builder.add_virtual_proof_with_pis(&inner_common);
-        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);
         if register_pi {
-            builder.register_public_inputs(&inner_pub_input[0..4]);
+            builder.register_public_inputs(&vir_proof.public_inputs);
         }
 
         // constant target for the verifier data
-        let const_verifier_data = builder.constant_verifier_data(&self.inner_verifier_data);
-
-        // register verifier data hash as public input.
-        let mut vd_pub_input = vec![];
-        vd_pub_input.extend_from_slice(&const_verifier_data.circuit_digest.elements);
-        for i in 0..builder.config.fri_config.num_cap_elements() {
-            vd_pub_input.extend_from_slice(&const_verifier_data.constants_sigmas_cap.0[i].elements);
-        }
-
-        let hash_inner_vd_pub_input = builder.hash_n_to_hash_no_pad::<H>(vd_pub_input);
-
-        // make sure the VerifierData we use is the same as the tree root hash of the VerifierData
-        builder.connect_hashes(hash_inner_vd_pub_input,HashOutTarget::from_vec(inner_pub_input[4..8].to_vec()));
-
-        if register_pi {
-            builder.register_public_inputs(&hash_inner_vd_pub_input.elements);
-        }
+        let const_verifier_data = builder.constant_verifier_data(&self.inner_verifier_data.verifier_only);
 
         // verify the proofs in-circuit
         builder.verify_proof::<C>(&vir_proof, &const_verifier_data, &inner_common);

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<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+>
+    where
+        <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>,
+{
+    phantom_data: PhantomData<(F, C)>,
+}
+
+impl<F, const D: usize, C> DummyProofGen<F, D, C>
+    where
+        F: RichField + Extendable<D> + Poseidon2,
+        C: GenericConfig<D, F = F>,
+        <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>,
+{
+    /// Builds a dummy circuit from the provided common circuit data.
+    pub fn gen_dummy_circ_data(
+        common_data: &CommonCircuitData<F, D>,
+    ) -> CircuitData<F, C, D> {
+        dummy_circuit::<F, C, D>(common_data)
+    }
+
+    /// Extracts the verifier-only data from the dummy circuit.
+    pub fn gen_dummy_verifier_data(
+        common_data: &CommonCircuitData<F, D>,
+    ) -> VerifierOnlyCircuitData<C, D> {
+        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<F, D>,
+        nonzero_public_inputs: HashMap<usize, F>,
+    ) -> Result<(ProofWithPublicInputs<F, C, D>, VerifierCircuitData<F, C, D>)> {
+        let circuit_data = Self::gen_dummy_circ_data(common_data);
+        let proof = dummy_proof::<F, C, D>(&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<F, D>,
+    ) -> Result<(ProofWithPublicInputs<F, C, D>, VerifierCircuitData<F, C, D>)> {
+        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<F: RichField + Extendable<D> + Poseidon2, C: GenericConfig<D, F = F>, const D: usize>(
+    common_data: &CommonCircuitData<F, D>,
+) -> CircuitData<F, C, D> {
+    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::<F, D>::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::<C>();
+    assert_eq!(&circuit.common, common_data);
+    circuit
+}

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

@@ -0,0 +1,248 @@
+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::VerifierCircuitData;
+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 1 inner proof
+/// T: total number of sampling proofs
+#[derive(Clone, Debug)]
+pub struct LeafCircuit<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+    const T: usize,
+> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    inner_verifier_data: VerifierCircuitData<F, C, D>,
+    phantom_data: PhantomData<H>
+}
+
+/// 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: ProofWithPublicInputsTarget<D>,
+    pub index: Target, // public input
+    pub flag: BoolTarget,
+}
+
+#[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 flag: bool,
+    pub index: usize
+}
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+    const T: usize,
+> LeafCircuit<F,D,C,H,T> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    pub fn new(
+        inner_verifier_data: VerifierCircuitData<F, C, D>,
+    ) -> Self {
+        Self {
+            inner_verifier_data,
+            phantom_data: PhantomData::default(),
+        }
+    }
+
+}
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+    const T: usize,
+> Plonky2Circuit<F, C, D> for LeafCircuit<F,D,C,H,T> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    type Targets = LeafTargets<D>;
+    type Input = LeafInput<F, D, C>;
+
+    fn add_targets(&self, builder: &mut CircuitBuilder<F, D>, register_pi: bool) -> Result<LeafTargets<D>> {
+
+        let inner_common = self.inner_verifier_data.common.clone();
+        let n_bucket: usize = bucket_count(T);
+
+        // the proof virtual target
+        let vir_proof = builder.add_virtual_proof_with_pis(&inner_common);
+        let inner_pub_input = vir_proof.public_inputs.clone();
+
+        // hash the public input & make it public
+        let hash_inner_pub_input = builder.hash_n_to_hash_no_pad::<H>(inner_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::<F>::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.verifier_only);
+
+        // virtual constant target for dummy verifier data
+        let const_dummy_vd = builder.constant_verifier_data(
+            &DummyProofGen::<F,D,C>::gen_dummy_verifier_data(&inner_common)
+        );
+
+        // 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)?;
+        // 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);
+        builder.verify_proof::<C>(&vir_proof, &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_proof,
+            index,
+            flag,
+        };
+        Ok(t)
+
+    }
+
+    fn assign_targets(
+        &self, pw: &mut PartialWitness<F>,
+        targets: &Self::Targets,
+        input: &Self::Input,
+    ) -> Result<()> {
+        assert!(input.index <= T, "given index is not valid");
+        // 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 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::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 = <C as GenericConfig<D>>::F;
+    type H = Poseidon2Hash;
+
+    /// A helper to build a minimal circuit and return the common circuit data.
+    fn dummy_common_circuit_data() -> CommonCircuitData<F, D> {
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+        // Add one virtual public input so that the circuit has minimal structure.
+        builder.add_virtual_public_input();
+        let circuit = builder.build::<C>();
+        circuit.common.clone()
+    }
+
+    // ----------------------------------------------------------------------------
+    // End-to-End test for the entire leaf circuit.
+    #[test]
+    fn test_full_leaf_circuit() -> anyhow::Result<()> {
+        // 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::<F, D, C>::gen_dummy_proof_and_vd_zero_pi(&common_data)?;
+        // let dummy_verifier_data = DummyProofGen::<F, D, C>::gen_dummy_verifier_data(&common_data);
+
+        // the leaf circuit.
+        let leaf = LeafCircuit::<F, D, C, H, 128>::new(vd.clone());
+
+        // 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: 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<u64> = 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(())
+    }
+}
+
+

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;

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

@@ -0,0 +1,366 @@
+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::{VerifierCircuitData, 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::dummy_gen::DummyProofGen;
+use crate::recursion::utils::bucket_count;
+
+/// recursion node circuit
+/// N: number of leaf proofs
+/// T: total number of sampling proofs
+#[derive(Clone, Debug)]
+pub struct NodeCircuit<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+    const N: usize,
+    const T: usize,
+> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    leaf_verifier_data: VerifierCircuitData<F, C, D>,
+    phantom_data: PhantomData<H>
+}
+
+/// 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 inner_proofs: Vec<ProofWithPublicInputsTarget<D>>,
+    pub inner_verifier_data: VerifierCircuitTarget,
+    pub condition: BoolTarget,
+    pub index: Target,
+    pub flags: Vec<BoolTarget>,
+}
+
+#[derive(Clone, Debug)]
+pub struct NodeInput<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+>{
+    pub inner_proofs: Vec<ProofWithPublicInputs<F, C, D>>,
+    pub verifier_only_data: VerifierOnlyCircuitData<C, D>,
+    pub condition: bool,
+    pub flags: Vec<bool>,
+    pub index: usize
+}
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+    const N: usize,
+    const T: usize,
+> NodeCircuit<F,D,C,H, N,T> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+
+    pub fn new(
+        leaf_verifier_data: VerifierCircuitData<F, C, D>,
+    ) -> Self {
+        assert!(N.is_power_of_two(), "M is NOT a power of two");
+        Self{
+            leaf_verifier_data,
+            phantom_data:PhantomData::default(),
+        }
+    }
+
+}
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+    const N: usize,
+    const T: usize,
+> Plonky2Circuit<F, C, D> for NodeCircuit<F, D, C, H, N, T> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    type Targets = NodeTargets<D>;
+    type Input = NodeInput<F, D, C>;
+
+    fn add_targets(&self, builder: &mut CircuitBuilder<F, D>, register_pi: bool) -> Result<Self::Targets> {
+        let inner_common = self.leaf_verifier_data.common.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..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[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::<H>(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.verifier_only);
+
+        // virtual constant target for dummy verifier data
+        let const_dummy_vd = builder.constant_verifier_data(
+            &DummyProofGen::<F,D,C>::gen_dummy_verifier_data(&inner_common)
+        );
+
+        // 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::<H>(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<Target> = (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<BoolTarget> = (0..N).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..N {
+            // 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::<C>(&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..N {
+            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(N as u64));
+        let mut expected_inner_index = builder.mul(index, m_const);
+        for i in 0..N {
+            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 {
+            inner_proofs: vir_proofs,
+            inner_verifier_data: node_verifier_data,
+            condition,
+            index,
+            flags,
+        };
+
+        Ok(t)
+    }
+
+    fn assign_targets(&self, pw: &mut PartialWitness<F>, targets: &Self::Targets, input: &Self::Input) -> Result<()> {
+        // assert size of proofs vec
+        assert_eq!(input.inner_proofs.len(), N);
+        assert_eq!(input.flags.len(), N);
+        assert!(input.index <= T, "given index is not valid");
+
+        // assign the proofs
+        for i in 0..N {
+            pw.set_proof_with_pis_target(&targets.inner_proofs[i], &input.inner_proofs[i])
+                .map_err(|e| {
+                    CircuitError::ProofTargetAssignmentError("inner-proof".to_string(), e.to_string())
+                })?;
+        }
+
+        // assign the verifier data
+        pw.set_verifier_data_target(&targets.inner_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..N {
+            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};
+    use plonky2_poseidon2::poseidon2_hash::poseidon2::{Poseidon2, Poseidon2Hash};
+    use crate::recursion::leaf::BUCKET_SIZE;
+
+
+    // For our tests, we define:
+    const D: usize = 2;
+    type C = PoseidonGoldilocksConfig;
+    type F = <C as GenericConfig<D>>::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<const B: usize>() -> (CircuitData<F, C, D>, Vec<Target>) {
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::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::<C>();
+        (data, pub_input)
+    }
+
+    /// A helper to generate test leaf proofs with given data, targets, and indices.
+    fn dummy_leaf_proofs<const B: usize>(data: CircuitData<F, C, D>, pub_input: Vec<Target>, indices: Vec<usize>) -> Vec<ProofWithPublicInputs<F, C, D>> {
+        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<F>{
+        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 N: usize = 2;
+        const B: usize = 4; // bucket size
+        const T: usize = 128;
+
+        let (leaf_data, leaf_pi) = dummy_leaf::<B>();
+        let leaf_vd = leaf_data.verifier_data();
+
+        let indices = vec![0,1];
+        let leaf_proofs = dummy_leaf_proofs::<B>(leaf_data,leaf_pi,indices);
+
+        let node = NodeCircuit::<F, D, C, H, N, T>::new(leaf_vd.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 {
+            inner_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<u64> = 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(())
+    }
+}

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

@@ -1,9 +1,9 @@
 use std::marker::PhantomData;
-use plonky2::hash::hash_types::{ HashOutTarget, RichField};
+use plonky2::hash::hash_types::RichField;
 use plonky2::iop::target::Target;
 use plonky2::iop::witness::{PartialWitness, WitnessWrite};
 use plonky2::plonk::circuit_builder::{CircuitBuilder};
-use plonky2::plonk::circuit_data::{CommonCircuitData, VerifierOnlyCircuitData};
+use plonky2::plonk::circuit_data::VerifierCircuitData;
 use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
 use plonky2::plonk::proof::{ProofWithPublicInputs, ProofWithPublicInputsTarget};
 use plonky2_field::extension::Extendable;
@@ -12,9 +12,7 @@ use crate::{error::CircuitError, Result};
 use crate::circuit_helper::Plonky2Circuit;
 
 /// A circuit that verifies the aggregated public inputs from inner circuits.
-///
-/// - `N`: Number of inner-proofs aggregated at the leaf level.
-/// - `M`: Number of leaf proofs aggregated at the node level.
+/// - `N`: Number of leaf proofs aggregated at the node level.
 /// - `T`: Total Number of inner-proofs.
 /// - `K`: Number of public input field elements per inner-proof (sampling proof).
 #[derive(Clone, Debug)]
@@ -24,14 +22,12 @@ pub struct PublicInputVerificationCircuit<
     C: GenericConfig<D, F = F>,
     H: AlgebraicHasher<F>,
     const N: usize,
-    const M: usize,
     const T: usize,
     const K: usize,
 > where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>,
 {
-    pub node_common_data: CommonCircuitData<F, D>,
-    pub node_verifier_data: VerifierOnlyCircuitData<C, D>,
+    pub node_verifier_data: VerifierCircuitData<F, C, D>,
     phantom: PhantomData<H>,
 }
 
@@ -45,7 +41,7 @@ pub struct PublicInputVerificationTargets<const D: usize> {
 }
 
 /// input to the circuit for public input verification
-/// - `inner_proof`: The tree root proof with 2 hash digests (8 Goldilocks field elements) public inputs [pi_hash, vd_hash].
+/// - `inner_proof`: The tree root proof with public inputs: [pi_hash, vd_hash, ...].
 /// - `inner_pub_inputs_vals`: T×K public input values from inner proofs.
 #[derive(Clone, Debug)]
 pub struct PublicInputVerificationInput<
@@ -57,8 +53,8 @@ pub struct PublicInputVerificationInput<
     pub inner_pub_inputs_vals: Vec<Vec<F>>,
 }
 
-impl<F, const D: usize, C, H, const N: usize, const M: usize, const T: usize, const K: usize>
-PublicInputVerificationCircuit<F, D, C, H, N, M, T, K>
+impl<F, const D: usize, C, H, const N: usize, const T: usize, const K: usize>
+PublicInputVerificationCircuit<F, D, C, H, N, T, K>
     where
         F: RichField + Extendable<D> + Poseidon2,
         C: GenericConfig<D, F = F>,
@@ -67,22 +63,20 @@ PublicInputVerificationCircuit<F, D, C, H, N, M, T, K>
 {
     /// Create a new instance of the circuit.
     pub fn new(
-        node_common_data: CommonCircuitData<F, D>,
-        node_verifier_data: VerifierOnlyCircuitData<C, D>,
+        node_verifier_data: VerifierCircuitData<F, C, D>,
     ) -> Self {
-        // we expect exactly 8 public inputs from the tree root proof
-        // 4 for the final aggregated public-input hash, 4 for the node verifier-data hash
-        assert_eq!(node_common_data.num_public_inputs, 8);
+        // we expect at least 4 public inputs from the tree root proof
+        // 1 hash digest (4 Goldilocks) for the final aggregated public-input hash
+        assert!(node_verifier_data.common.num_public_inputs >= 4);
 
         Self {
-            node_common_data,
             node_verifier_data,
             phantom: PhantomData,
         }
     }
 }
-impl<F, const D: usize, C, H, const N: usize, const M: usize, const T: usize, const K: usize>
-Plonky2Circuit<F, C, D> for PublicInputVerificationCircuit<F, D, C, H, N, M, T, K>
+impl<F, const D: usize, C, H, const N: usize, const T: usize, const K: usize>
+Plonky2Circuit<F, C, D> for PublicInputVerificationCircuit<F, D, C, H, N, T, K>
     where
         F: RichField + Extendable<D> + Poseidon2,
         C: GenericConfig<D, F = F>,
@@ -93,19 +87,19 @@ Plonky2Circuit<F, C, D> for PublicInputVerificationCircuit<F, D, C, H, N, M, T,
     type Input = PublicInputVerificationInput<F, D, C>;
 
     /// Builds the circuit by:
-    /// 1. Verifies a proof target with 8 public inputs (the final [pi_hash, vd_hash]).
+    /// 1. Verifies a proof target with public inputs (the final [pi_hash, vd_hash, ...]).
     /// 2. verifies correct tree hashing of all T×K targets to represent all inner public inputs.
-    /// 3. verifies correct node_verifier_date used is the same as in public input (last 4 field elements).
+    /// 3. register the un-hashed inner public input as this circuit public input + the rest of the inner public input
     fn add_targets(&self, builder: &mut CircuitBuilder<F, D>, register_pi: bool) -> Result<PublicInputVerificationTargets<D>> {
         // Add a virtual proof with 8 public inputs. This is the final root proof whose
         // public inputs we want to check in-circuit.
-        let inner_proof = builder.add_virtual_proof_with_pis(&self.node_common_data);
+        let inner_proof = builder.add_virtual_proof_with_pis(&self.node_verifier_data.common);
 
         // Create a constant VerifierCircuitTarget for the node's verifier data.
-        let const_node_vd = builder.constant_verifier_data(&self.node_verifier_data);
+        let const_node_vd = builder.constant_verifier_data(&self.node_verifier_data.verifier_only);
 
         // verify the proof
-        builder.verify_proof::<C>(&inner_proof, &const_node_vd, &self.node_common_data);
+        builder.verify_proof::<C>(&inner_proof, &const_node_vd, &self.node_verifier_data.common);
 
         // create T×K targets for all inner public inputs from the base level.
         let mut inner_pub_inputs = Vec::with_capacity(T);
@@ -125,59 +119,39 @@ Plonky2Circuit<F, C, D> for PublicInputVerificationCircuit<F, D, C, H, N, M, T,
         //   Summary of the logic:
         //
         //   let final_pi = proof.public_inputs[0..4];
-        //   let node_vd = proof.public_inputs[4..8];
         //   ...
-        //   leaf-level pub inputs tree hashing: chunks of N -> hash
-        //   node-level pub inputs tree hashing: chunks of M -> hash
+        //   leaf-level pub inputs tree hashing
+        //   node-level pub inputs tree hashing: chunks of N -> hash
         //   ...
         //   check final result matches final_pi
         // ------------------------------------------------------------------
 
-        // Extract the final 4 field elements for the public-input hash & next 4 for the verifier-data hash.
-        let final_pi_hash_t = &inner_proof.public_inputs[0..4];
-        let node_vd_hash_t = &inner_proof.public_inputs[4..8];
-
-        // Compute node_hash in-circuit
-        let mut node_vd_input_t = Vec::new();
-        node_vd_input_t.extend_from_slice(&const_node_vd.circuit_digest.elements);
-        for cap_elem in const_node_vd.constants_sigmas_cap.0.iter() {
-            node_vd_input_t.extend_from_slice(&cap_elem.elements);
-        }
-        let node_hash_t = builder.hash_n_to_hash_no_pad::<H>(node_vd_input_t);
-        // make sure the VerifierData we use is the same as the tree root hash of the VerifierData
-        builder.connect_hashes(node_hash_t,HashOutTarget::from_vec(node_vd_hash_t.to_vec()));
-        if register_pi {
-            builder.register_public_inputs(&node_hash_t.elements); // public input
-        }
+        // Extract the final 4 field elements for the public-input hash & the rest for the verifier-data hash, index, and flags.
+        let final_pi_hash_target = &inner_proof.public_inputs[0..4];
+        let rest_of_inner_pi = &inner_proof.public_inputs[4..];
+        builder.register_public_inputs(&rest_of_inner_pi); // public input
 
         let mut pub_in_hashes_t = Vec::new();
 
-        // Leaf level hashing: chunks of N
-        let base_chunks = T / N; // T is assumed to be multiple of N
-        for i in 0..base_chunks {
-            // flatten the inputs from i*N .. i*N + N
-            let mut chunk_targets = Vec::with_capacity(N * K);
-            for row_idx in (i * N)..(i * N + N) {
-                chunk_targets.extend_from_slice(&inner_pub_inputs[row_idx]);
-            }
+        // Leaf level hashing - hash each row i = 0..T of inner_pub_inputs matrix
+        for i in 0..T {
             // hash
-            let pi_hash_chunk = builder.hash_n_to_hash_no_pad::<H>(chunk_targets);
-
-            // track these in vectors
+            let pi_hash_chunk = builder.hash_n_to_hash_no_pad::<H>(inner_pub_inputs[i].clone());
+            // track these in hash digests
             pub_in_hashes_t.push(pi_hash_chunk);
         }
 
         // Now at the node level:
-        let mut current_len = base_chunks;
+        let mut current_len = 0;
         while current_len > 1 {
 
-            let next_len = (current_len + (M - 1)) / M;
+            let next_len = (current_len + (N - 1)) / N;
 
             let mut next_pub_in_hashes_t = Vec::with_capacity(next_len);
 
             for i in 0..next_len {
-                let start_idx = i * M;
-                let end_idx = (start_idx + M).min(current_len);
+                let start_idx = i * N;
+                let end_idx = (start_idx + N).min(current_len);
 
                 // flatten all pub_in_hashes in [start_idx..end_idx]
                 let mut pi_flat = Vec::with_capacity((end_idx - start_idx) * 4);
@@ -197,7 +171,7 @@ Plonky2Circuit<F, C, D> for PublicInputVerificationCircuit<F, D, C, H, N, M, T,
 
         // connect them to the final 4 public inputs of `inner_proof`.
         for i in 0..4 {
-            builder.connect(final_pi_hash_t[i], final_computed_pi_t.elements[i]);
+            builder.connect(final_pi_hash_target[i], final_computed_pi_t.elements[i]);
         }
 
         // return all the targets
@@ -213,7 +187,7 @@ Plonky2Circuit<F, C, D> for PublicInputVerificationCircuit<F, D, C, H, N, M, T,
         targets: &Self::Targets,
         input: &Self::Input,
     ) -> Result<()> {
-        // Assign the final proof - it should have 8 public inputs
+        // Assign the tree root proof
         pw.set_proof_with_pis_target(&targets.inner_proof, &input.inner_proof)
             .map_err(|e| {
                 CircuitError::ProofTargetAssignmentError("final-proof".to_string(), e.to_string())

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

@@ -1,5 +1,5 @@
 use std::marker::PhantomData;
-use plonky2::hash::hash_types::{HashOut, RichField};
+use plonky2::hash::hash_types::RichField;
 use plonky2::iop::witness::PartialWitness;
 use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, CommonCircuitData, VerifierCircuitData, VerifierOnlyCircuitData};
 use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
@@ -7,28 +7,29 @@ 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;
+use crate::recursion::utils::get_hash_of_verifier_data;
 
 /// tree recursion
-/// - `N`: Number of inner-proofs aggregated at the leaf level.
-/// - `M`: Number of leaf proofs aggregated at the node level.
+/// - `N`: Number of leaf proofs aggregated at the node level. set to 2 for 2-to-1 tree
 pub struct TreeRecursion<
     F: RichField + Extendable<D> + Poseidon2,
     const D: usize,
     C: GenericConfig<D, F = F>,
     H: AlgebraicHasher<F>,
     const N: usize,
-    const M: usize,
+    const T: usize,
 > where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
 {
-    leaf: LeafCircuit<F, D, C, H, N>,
-    node: NodeCircuit<F, D, C, H, M>,
-    compression: CompressionCircuit<F, D, C, H>,
+    leaf: LeafCircuit<F, D, C, H, T>,
+    node: NodeCircuit<F, D, C, H, N, T>,
+    compression: CompressionCircuit<F, D, C>,
     leaf_circ_data: CircuitData<F, C, D>,
     node_circ_data: CircuitData<F, C, D>,
     compression_circ_data: CircuitData<F, C, D>,
@@ -44,17 +45,15 @@ impl<
     C: GenericConfig<D, F = F>,
     H: AlgebraicHasher<F>,
     const N: usize,
-    const M: usize,
-> TreeRecursion<F, D, C, H, N, M> where
+    const T: usize,
+> TreeRecursion<F, D, C, H, N, T> where
     <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
 {
     /// build with standard recursion config
     pub fn build_with_standard_config(
-        inner_common_data: CommonCircuitData<F,D>,
-        inner_verifier_data: VerifierOnlyCircuitData<C, D>,
+        inner_verifier_data: VerifierCircuitData<F, C, D>,
     ) -> Result<Self> {
         Self::build(
-            inner_common_data,
             inner_verifier_data,
             CircuitConfig::standard_recursion_config()
         )
@@ -62,25 +61,23 @@ impl<
 
     /// build the tree with given config
     pub fn build(
-        inner_common_data: CommonCircuitData<F,D>,
-        inner_verifier_data: VerifierOnlyCircuitData<C, D>,
+        inner_verifier_data: VerifierCircuitData<F, C, D>,
         config: CircuitConfig,
     ) -> Result<Self> {
         // build leaf with standard recursion config
-        let leaf = LeafCircuit::<_,D,_,_,N>::new(inner_common_data, inner_verifier_data);
+        let leaf = LeafCircuit::<_,D,_,_,T>::new(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,_,_,N, T>::new(leaf_circ_data.verifier_data());
         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_circ = CompressionCircuit::new(node_circ_data.verifier_data());
         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 +109,10 @@ impl<
         self.node_circ_data.verifier_data()
     }
 
+    pub fn prove_bundle(_bundle: Bundle<F, C, D, H>){
+        todo!()
+    }
+
     pub fn prove_tree_and_compress(
         &mut self,
         proofs_with_pi: &[ProofWithPublicInputs<F, C, D>],
@@ -163,9 +164,11 @@ impl<
 
         let mut leaf_proofs = vec![];
 
-        for proofs in proofs_with_pi.chunks(N){
+        for (i, proof) in proofs_with_pi.iter().enumerate(){
             let leaf_input = LeafInput{
-                inner_proof: proofs.to_vec(),
+                inner_proof: proof.clone(),
+                flag: true,
+                index: i,
             };
 
             let mut pw = PartialWitness::<F>::new();
@@ -196,14 +199,16 @@ 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(N).enumerate() {
 
             let mut inner_pw = PartialWitness::new();
 
             let node_input = NodeInput{
-                node_proofs: chunk.to_vec().clone(),
+                inner_proofs: chunk.to_vec().clone(),
                 verifier_only_data: verifier_only_data.clone(),
                 condition,
+                flags: [true; N].to_vec(),
+                index: i,
             };
 
             self.node.assign_targets(
@@ -257,7 +262,7 @@ impl<
         public_input: Vec<F>,
         inner_public_input: Vec<Vec<F>>,
     ) -> 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];
@@ -265,18 +270,14 @@ impl<
         let node_hash = get_hash_of_verifier_data::<F,D,C,H>(&self.node_circ_data.verifier_data());
 
         let mut pub_in_hashes = vec![];
-        for pub_in in inner_public_input.chunks(N){
-            let pub_in_flat: Vec<F> = pub_in
-                .iter()
-                .flat_map(|v| v.iter().cloned())
-                .collect();
-            let hash = H::hash_no_pad(&pub_in_flat);
+        for pub_in in inner_public_input{
+            let hash = H::hash_no_pad(&pub_in);
             pub_in_hashes.push(hash);
         }
 
         while pub_in_hashes.len() > 1 {
             let mut next_level_pi_hashes = Vec::new();
-            for pi_chunk in pub_in_hashes.chunks(M) {
+            for pi_chunk in pub_in_hashes.chunks(N) {
                 // collect field elements
                 let pi_chunk_f: Vec<F> = pi_chunk.iter()
                     .flat_map(|h| h.elements.iter().cloned())
@@ -297,23 +298,61 @@ impl<
     }
 }
 
-/// helper fn to generate hash of verifier data
-pub fn get_hash_of_verifier_data<
-    F: RichField + Extendable<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
->(verifier_data: &VerifierCircuitData<F, C, D>) -> HashOut<F> where
-    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-    let mut vd = vec![];
-    let digest: &HashOut<F> = &verifier_data.verifier_only.circuit_digest;
-    let caps = &verifier_data.verifier_only.constants_sigmas_cap;
-    vd.extend_from_slice(&digest.elements);
-    for i in 0..verifier_data.common.config.fri_config.num_cap_elements() {
-        let cap_hash = caps.0[i] as HashOut<F>;
-        vd.extend_from_slice(&cap_hash.elements);
+#[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 = <C as GenericConfig<D>>::F;
+    type H = Poseidon2Hash;
+
+    // A helper to build a minimal circuit and returns T proofs & circuit data.
+    fn dummy_proofs<const T: usize>() -> (CircuitData<F, C, D>, Vec<ProofWithPublicInputs<F, C, D>>) {
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+        for _ in 0..(4096+10) {
+            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::<C>();
+        println!("inner circuit size = {}", circuit.common.degree_bits());
+        let mut pw = PartialWitness::<F>::new();
+        pw.set_target(t, F::ZERO).expect("faulty assign");
+        let proofs = (0..T).map(|_i| circuit.prove(pw.clone()).unwrap()).collect();
+        (circuit, proofs)
     }
 
-    H::hash_no_pad(&vd)
-}
+
+    // End-to-End test for the entire Tree circuit.
+    #[test]
+    fn test_full_tree_circuit() -> anyhow::Result<()> {
+        const N: usize = 2;
+        const T: usize = 128;
+
+        let (data, proofs) = dummy_proofs::<T>();
+
+        let mut tree = TreeRecursion::<F,D,C,H, N, T>::build_with_standard_config(data.verifier_data())?;
+
+        // 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(())
+    }
+}

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

@@ -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<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
-    const N: usize,
-> where
-    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-    inner_common_data: CommonCircuitData<F, D>,
-    inner_verifier_data: VerifierOnlyCircuitData<C, D>,
-    phantom_data: PhantomData<H>
-}
-
-#[derive(Clone, Debug)]
-pub struct LeafTargets <
-    const D: usize,
->{
-    pub inner_proof: Vec<ProofWithPublicInputsTarget<D>>,
-}
-
-#[derive(Clone, Debug)]
-pub struct LeafInput<
-    F: RichField + Extendable<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
->{
-    pub inner_proof: Vec<ProofWithPublicInputs<F, C, D>>
-}
-
-impl<
-    F: RichField + Extendable<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
-    const N: usize,
-> LeafCircuit<F,D,C,H,N> where
-    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-    pub fn new(
-        inner_common_data: CommonCircuitData<F, D>,
-        inner_verifier_data: VerifierOnlyCircuitData<C, D>,
-    ) -> Self {
-        Self {
-            inner_common_data,
-            inner_verifier_data,
-            phantom_data: PhantomData::default(),
-        }
-    }
-}
-
-impl<
-    F: RichField + Extendable<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
-    const N: usize,
-> Plonky2Circuit<F, C, D> for LeafCircuit<F,D,C,H,N> where
-    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-    type Targets = LeafTargets<D>;
-    type Input = LeafInput<F, D, C>;
-
-    fn add_targets(&self, builder: &mut CircuitBuilder<F, D>, register_pi: bool) -> Result<LeafTargets<D>> {
-
-        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::<H>(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::<F>::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::<C>(&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<F>,
-        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(())
-    }
-
-}
-
-

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

@@ -1,5 +0,0 @@
-pub mod leaf;
-pub mod node;
-pub mod tree;
-pub mod compress;
-pub mod pi_verifier;

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

@@ -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<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
-    const M: usize,
-> where
-    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-    common_data: CommonCircuitData<F, D>,
-    leaf_verifier_data: VerifierOnlyCircuitData<C, D>,
-    phantom_data: PhantomData<H>
-}
-
-#[derive(Clone, Debug)]
-pub struct NodeTargets<
-    const D: usize,
->{
-    pub leaf_proofs: Vec<ProofWithPublicInputsTarget<D>>,
-    pub node_verifier_data: VerifierCircuitTarget,
-    pub condition: BoolTarget,
-}
-
-#[derive(Clone, Debug)]
-pub struct NodeInput<
-    F: RichField + Extendable<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
->{
-    pub node_proofs: Vec<ProofWithPublicInputs<F, C, D>>,
-    pub verifier_only_data: VerifierOnlyCircuitData<C, D>,
-    pub condition: bool,
-}
-
-impl<
-    F: RichField + Extendable<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
-    const M: usize,
-> NodeCircuit<F,D,C,H,M> where
-    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-
-    pub fn new(
-        common_data: CommonCircuitData<F,D>,
-        leaf_verifier_data: VerifierOnlyCircuitData<C, D>,
-    ) -> Self {
-        Self{
-            common_data,
-            leaf_verifier_data,
-            phantom_data:PhantomData::default(),
-        }
-    }
-
-}
-
-impl<
-    F: RichField + Extendable<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-    H: AlgebraicHasher<F>,
-    const M: usize,
-> Plonky2Circuit<F, C, D> for NodeCircuit<F, D, C, H, M> where
-    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-    type Targets = NodeTargets<D>;
-    type Input = NodeInput<F, D, C>;
-
-    fn add_targets(&self, builder: &mut CircuitBuilder<F, D>, register_pi: bool) -> Result<Self::Targets> {
-        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::<H>(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::<H>(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::<C>(&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<F>, 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(())
-    }
-}
-
-

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

@@ -0,0 +1,267 @@
+use plonky2::hash::hash_types::{HashOut, RichField};
+use plonky2::iop::target::{BoolTarget, Target};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::VerifierCircuitData;
+use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
+use plonky2_field::extension::Extendable;
+use plonky2_poseidon2::Poseidon2;
+use crate::recursion::leaf::BUCKET_SIZE;
+
+/// 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<D> + Poseidon2,
+    const D: usize,
+>(
+    builder: &mut CircuitBuilder<F, D>,
+    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<BoolTarget> = 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<D> + Poseidon2,
+    const D: usize,
+>(
+    builder: &mut CircuitBuilder<F, D>,
+    r: Target,
+) -> crate::Result<Target>
+{
+    // 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<D> + Poseidon2,
+    const D: usize,
+>(
+    builder: &mut CircuitBuilder<F, D>,
+    index: Target,
+    flag: BoolTarget,
+    bucket_size: usize,
+    num_buckets: usize,
+) -> crate::Result<Vec<Target>>
+{
+    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::<F,D>(builder, index, bucket_size, num_buckets)?;
+    // Compute 2^(r)
+    let power_of_two = compute_power_of_two::<F,D>(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
+}
+
+/// helper fn to generate hash of verifier data
+pub fn get_hash_of_verifier_data<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+>(verifier_data: &VerifierCircuitData<F, C, D>) -> HashOut<F> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    let mut vd = vec![];
+    let digest: &HashOut<F> = &verifier_data.verifier_only.circuit_digest;
+    let caps = &verifier_data.verifier_only.constants_sigmas_cap;
+    vd.extend_from_slice(&digest.elements);
+    for i in 0..verifier_data.common.config.fri_config.num_cap_elements() {
+        let cap_hash = caps.0[i] as HashOut<F>;
+        vd.extend_from_slice(&cap_hash.elements);
+    }
+
+    H::hash_no_pad(&vd)
+}
+
+#[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 = <C as GenericConfig<D>>::F;
+    type H = Poseidon2Hash;
+
+    // Helper: Build, prove, and return public inputs ---
+    fn build_and_prove(builder: CircuitBuilder<F, D>) -> Vec<F> {
+        // Build the circuit.
+        let circuit = builder.build::<C>();
+        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::<F, D>::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::<F,D>(&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::<F, D>::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::<F,D>(&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::<F, D>::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::<F,D>(
+            &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(())
+    }
+}
+
+

codex-plonky2-circuits/src/recursion/utils/conditional_verifier.rs

@@ -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<F: RichField + Extendable<D> + Poseidon2, const D: usize ,C: GenericConfig<D, F = F> + 'static>(
-    builder: &mut CircuitBuilder<F, D>,
-    condition: BoolTarget,
-    cyclic_proof_with_pis: &ProofWithPublicInputsTarget<D>,
-    verifier_data: &VerifierCircuitTarget,
-    common_data: &CommonCircuitData<F, D>,
-) -> anyhow::Result<()>
-    where
-        C::Hasher: AlgebraicHasher<F>,
-{
-    let (dummy_proof_with_pis_target, dummy_verifier_data_target) =
-        dummy_proof_and_vk_no_generator::<F, D, C>(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::<C>(
-        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<F: RichField + Extendable<D> + Poseidon2, const D: usize ,C: GenericConfig<D, F = F> + 'static>(
-    builder: &mut CircuitBuilder<F, D>,
-    condition: BoolTarget,
-    proof_with_pis: &ProofWithPublicInputsTarget<D>,
-    inner_verifier_data: &VerifierCircuitTarget,
-    inner_common_data: &CommonCircuitData<F, D>,
-) -> anyhow::Result<()>
-    where
-        C::Hasher: AlgebraicHasher<F>,
-{
-    let (dummy_proof_with_pis_target, dummy_verifier_data_target) =
-        dummy_proof_and_vk_no_generator::<F, D, C>(builder, inner_common_data)?;
-    builder.conditionally_verify_proof::<C>(
-        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<F: RichField + Extendable<D> + Poseidon2, C: GenericConfig<D, F = F>, const D: usize>(
-    common_data: &CommonCircuitData<F, D>,
-) -> CircuitData<F, C, D> {
-    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::<F, D>::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::<C>();
-    assert_eq!(&circuit.common, common_data);
-    circuit
-}
-
-pub(crate) fn dummy_proof_and_vk_no_generator<F: RichField + Extendable<D> + Poseidon2, const D: usize ,C: GenericConfig<D, F = F> + 'static> (
-    builder: &mut CircuitBuilder<F, D>,
-    common_data: &CommonCircuitData<F, D>,
-) -> anyhow::Result<(ProofWithPublicInputsTarget<D>, VerifierCircuitTarget)>
-    where
-        C::Hasher: AlgebraicHasher<F>,
-{
-    let dummy_circuit = dummy_circuit::<F, C, D>(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))
-}

codex-plonky2-circuits/src/recursion/utils/dummy_gen.rs

@@ -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<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-    > where
-        <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-    phantom_data: PhantomData<(F,C)>,
-}
-
-impl<
-    F: RichField + Extendable<D> + Poseidon2,
-    const D: usize,
-    C: GenericConfig<D, F = F>,
-> DummyProofGen<F, D, C>
-    where
-        <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
-{
-
-    /// Generates a single dummy leaf proof.
-    pub fn gen_dummy_leaf_proof(
-        common_data: &CommonCircuitData<F, D>,
-    ) -> Result<ProofWithPublicInputs<F, C, D>> {
-        dummy_proof::<F, C, D>(&dummy_circuit::<F, C, D>(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<F, D>,
-        node_verifier_only_data: &VerifierOnlyCircuitData<C, D>,
-    ) -> ProofWithPublicInputs<F, C, D> {
-        Self::recursion_base_proof(node_common, HashMap::new())
-    }
-
-    fn recursion_base_proof(
-        common_data: &CommonCircuitData<F, D>,
-        mut nonzero_public_inputs: HashMap<usize, F>
-    ) -> ProofWithPublicInputs<F, C, D>{
-        dummy_proof::<F, C, D>(
-            &dummy_circuit::<F, C, D>(common_data),
-            nonzero_public_inputs,
-        )
-            .unwrap()
-    }
-
-    /// Generates an array of `N` dummy leaf proofs.
-    pub fn gen_n_dummy_leaf_proofs<const N: usize>(
-        common_data: &CommonCircuitData<F, D>,
-    ) -> Result<[ProofWithPublicInputs<F, C, D>; N]> {
-        let dummy_proof = Self::gen_dummy_leaf_proof(common_data)?;
-        let n_dummy_vec = (0..N).map(|_| dummy_proof.clone()).collect::<Vec<_>>();
-        vec_to_array::<N, ProofWithPublicInputs<F, C, D>>(n_dummy_vec)
-    }
-
-    /// Generates an array of `N` dummy node proofs.
-    pub fn gen_n_dummy_node_proofs<const N: usize>(
-        node_common: &CommonCircuitData<F, D>,
-        node_verifier_only_data: &VerifierOnlyCircuitData<C, D>,
-    ) -> Result<[ProofWithPublicInputs<F, C, D>; 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<_>>();
-        vec_to_array::<N, ProofWithPublicInputs<F, C, D>>(n_dummy_vec)
-    }
-}

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

@@ -1,2 +0,0 @@
-pub mod dummy_gen;
-pub mod conditional_verifier;

proof-input/src/recursion/leaf_test.rs

@@ -0,0 +1,85 @@
+// some tests for the leaf in tree recursion
+
+#[cfg(test)]
+pub mod tests {
+    use plonky2::plonk::circuit_data::{ProverCircuitData, VerifierCircuitData};
+    use plonky2::plonk::proof::ProofWithPublicInputs;
+    use plonky2_field::types::{Field, PrimeField64};
+    use codex_plonky2_circuits::circuit_helper::Plonky2Circuit;
+    use codex_plonky2_circuits::recursion::leaf::{LeafCircuit, LeafInput};
+    use codex_plonky2_circuits::recursion::dummy_gen::DummyProofGen;
+    use crate::params::{F, D, C, HF};
+    use crate::recursion::run_sampling_circ;
+
+    pub fn run_leaf_circ<const T: usize>(inner_proof: ProofWithPublicInputs<F, C, D>, inner_verifier_data: VerifierCircuitData<F, C, D>, flag: bool, index: usize) -> anyhow::Result<(ProofWithPublicInputs<F, C, D>, ProverCircuitData<F, C, D>, VerifierCircuitData<F, C, D>)> {
+
+        // ------------------- leaf --------------------
+        let leaf = LeafCircuit::<F,D,C,HF,T>::new(inner_verifier_data.clone());
+
+        // build
+        let (targets, data) = leaf.build_with_standard_config()?;
+        let verifier_data: VerifierCircuitData<F,C,D> = data.verifier_data();
+        let prover_data = data.prover_data();
+        println!("leaf circuit degree bits = {:?}", prover_data.common.degree_bits());
+
+        // prove
+        let input = LeafInput{
+            inner_proof,
+            flag,
+            index,
+        };
+        let proof = leaf.prove(&targets, &input, &prover_data)?;
+        println!("pub input size = {}", proof.public_inputs.len());
+        println!("proof size = {:?} bytes", proof.to_bytes().len());
+        println!("pub input = {:?}", proof.public_inputs);
+
+        // verify
+        assert!(
+            verifier_data.verify(proof.clone()).is_ok(),
+            "proof verification failed"
+        );
+
+        let flag_buckets: Vec<F> = proof.public_inputs[9..13].to_vec();
+        if flag {
+            check_flag_buckets(index, flag_buckets);
+        } else {
+            for i in 0..flag_buckets.len() {
+                assert_eq!(flag_buckets[i], F::ZERO, "bucket not valid");
+            }
+        }
+
+        Ok((proof, prover_data, verifier_data))
+    }
+
+    fn check_flag_buckets(index: usize, flag_buckets: Vec<F>) {
+        // Compute the bucket and bit position from the input index.
+        let bucket = index / 32;
+        let bit = index % 32;
+        // For each flag target (bucket), assign the appropriate 32-bit one-hot value.
+        for (i, &flag_bucket) in flag_buckets.iter().enumerate() {
+            let value: u64 = if i == bucket {
+                    1 << bit
+            } else {
+                0
+            };
+            assert_eq!(value, flag_bucket.to_canonical_u64(), "bucket value mismatch");
+        }
+    }
+
+    #[test]
+    fn test_real_leaf_circ() -> anyhow::Result<()> {
+        let (inner_proof, _, inner_verifier) = run_sampling_circ()?;
+
+        run_leaf_circ::<128>(inner_proof, inner_verifier, true, 1)?;
+        Ok(())
+    }
+
+    #[test]
+    fn test_dummy_leaf_circ() -> anyhow::Result<()> {
+        let (_, _, inner_verifier) = run_sampling_circ()?;
+        let (dummy_proof, dummy_vd) = DummyProofGen::gen_dummy_proof_and_vd_zero_pi(&inner_verifier.common)?;
+        run_leaf_circ::<128>(dummy_proof, dummy_vd, false, 0)?;
+        Ok(())
+    }
+
+}

proof-input/src/recursion/mod.rs

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

proof-input/src/recursion/node_test.rs

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

proof-input/src/recursion/tree_test.rs

@@ -0,0 +1,59 @@
+// some tests for the tree recursion
+
+#[cfg(test)]
+mod tests {
+    use plonky2::plonk::proof::{ProofWithPublicInputs};
+    use codex_plonky2_circuits::circuit_helper::Plonky2Circuit;
+    use crate::params::{F, D, C, HF};
+    use codex_plonky2_circuits::recursion::{tree::TreeRecursion};
+    use crate::recursion::run_sampling_circ;
+
+    fn run_tree_recursion<const N: usize, const T: usize>(compress: bool) -> anyhow::Result<()> {
+
+        //------------ sampling inner circuit ----------------------
+        // Circuit that does the sampling - 100 samples
+        let (inner_proof, inner_prover_data, inner_verifier_data) = run_sampling_circ()?;
+
+        let proofs: Vec<ProofWithPublicInputs<F, C, D>> = (0..T).map(|_i| inner_proof.clone()).collect();
+
+        // ------------------- tree --------------------
+        // N-to-1 tree aggregation
+
+        let mut tree = TreeRecursion::<F,D,C,HF, N, T>::build_with_standard_config(inner_verifier_data.clone())?;
+
+        // aggregate
+        let root = if !compress {
+            tree.prove_tree(&proofs)?
+        } else {
+            println!("Mode: tree with compression");
+            tree.prove_tree_and_compress(&proofs)?
+        };
+        println!("pub input size = {}", root.public_inputs.len());
+        println!("pub input = {:?}", root.public_inputs);
+        println!("proof size = {:?} bytes", root.to_bytes().len());
+
+        let inner_pi: Vec<Vec<F>> = proofs.iter().map(|p| p.public_inputs.clone()).collect();
+
+        assert!(
+            tree.verify_proof_and_public_input(root,inner_pi.clone(), compress).is_ok(),
+            "proof verification failed"
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_tree_recursion() -> anyhow::Result<()> {
+        // total number of proofs to aggregate
+        const T:usize = 4;
+        run_tree_recursion::<2, T>(false)
+    }
+
+    #[test]
+    fn test_tree_recursion_with_compression() -> anyhow::Result<()> {
+        // total number of proofs to aggregate
+        const T:usize = 4;
+        run_tree_recursion::<2, T>(true)
+    }
+
+}

proof-input/src/recursion/uniform.rs

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

workflow/src/bin/aggregate.rs

@@ -1,26 +1,38 @@
 use std::env;
 use anyhow::Result;
 use plonky2::plonk::proof::ProofWithPublicInputs;
-use codex_plonky2_circuits::recursion::uniform::tree::TreeRecursion;
+use codex_plonky2_circuits::recursion::tree::TreeRecursion;
 use proof_input::params::{D, C, F, HF};
 use proof_input::serialization::file_paths::{PROOF_JSON, TREE_PROOF_JSON, VERIFIER_CIRC_DATA_JSON};
 use proof_input::serialization::json::{export_tree_proof_with_pi, import_proof_with_pi, import_verifier_circuit_data};
 
 fn main() -> Result<()> {
     // load the parameters from environment variables
-    const N: usize = 1;
-    const M: usize = 2;
+    const N: usize = 2;
 
     // take k = "number of proofs" from env arguments; default to 4 if not there
     let args: Vec<String> = env::args().collect();
-    let k: usize = if args.len() > 1 {
-        args[1]
-            .parse()
-            .expect("k not valid")
-    } else {
-        4
-    };
+    let t: usize = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(4);
 
+    match t {
+        2 => run_tree::<2>()?,
+        4 => run_tree::<4>()?,
+        8 => run_tree::<8>()?,
+        16 => run_tree::<2>()?,
+        32 => run_tree::<4>()?,
+        64 => run_tree::<8>()?,
+        128 => run_tree::<2>()?,
+        256 => run_tree::<4>()?,
+        512 => run_tree::<8>()?,
+        1024 => run_tree::<8>()?,
+        other => panic!("unsupported proof count: {}", other),
+    }
+
+    Ok(())
+}
+
+fn run_tree<const T: usize>() -> Result<()> {
+    const N: usize = 2;
     // Read the proof
     let proof_with_pi = import_proof_with_pi::<F,C,D>()?;
     println!("Proof with public input imported from: {}", PROOF_JSON);
@@ -31,11 +43,11 @@ fn main() -> Result<()> {
 
     // duplicate the proof to get k proofs
     // this is just for testing - in real scenario we would need to load k proofs
-    let proofs: Vec<ProofWithPublicInputs<F, C, D>> = (0..k).map(|_i| proof_with_pi.clone()).collect();
+    let proofs: Vec<ProofWithPublicInputs<F, C, D>> = (0..T).map(|_i| proof_with_pi.clone()).collect();
 
-    let mut tree = TreeRecursion::<F,D,C,HF, N, M>::build_with_standard_config(verifier_data.common.clone(), verifier_data.verifier_only.clone()).unwrap();
+    let mut tree = TreeRecursion::<F,D,C,HF, N, T>::build_with_standard_config(verifier_data.clone()).unwrap();
 
-    let tree_proof = tree.prove_tree(&proofs).unwrap();
+    let tree_proof = tree.prove_tree_and_compress(&proofs).unwrap();
     //export the proof to json file
     export_tree_proof_with_pi(&tree_proof)?;
     println!("Tree proof written to: {}", TREE_PROOF_JSON);
@@ -45,4 +57,4 @@ fn main() -> Result<()> {
     assert!(tree.verify_proof_and_public_input(tree_proof,inner_pi.clone(),false).is_ok());
 
     Ok(())
-}
+}