impl compression and public input verification circuits.

codex-plonky2-circuits/src/circuits/sample_cells.rs

@@ -19,7 +19,8 @@ use plonky2::{
     },
     plonk::circuit_builder::CircuitBuilder,
 };
-use plonky2::plonk::config::AlgebraicHasher;
+use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData};
+use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
 use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
 
 use crate::{
@@ -55,6 +56,7 @@ impl<
             phantom_data: Default::default(),
         }
     }
+    
 }
 
 /// struct of input to the circuit as targets
@@ -350,10 +352,9 @@ impl<
         // circuit params
         let CircuitParams {
             max_depth,
-            max_log2_n_slots,
-            block_tree_depth,
             n_field_elems_per_cell,
             n_samples,
+            ..
         } = self.params;
 
         // assign n_cells_per_slot

codex-plonky2-circuits/src/error.rs

@@ -68,4 +68,10 @@ pub enum CircuitError {
 
     #[error("Expected Option {0} to contain value")]
     OptionError(String),
+
+    #[error("Public input length Error: Expected {0}, got {1}")]
+    PublicInputLengthError(usize, usize),
+
+    #[error("{0}")]
+    InvalidArgument(String),
 }

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

@@ -0,0 +1,129 @@
+use std::marker::PhantomData;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::witness::{PartialWitness, WitnessWrite};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, CommonCircuitData, 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};
+
+/// recursion compression circuit - verifies 1 inner proof
+#[derive(Clone, Debug)]
+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>,
+    phantom_data: PhantomData<(C,H)>
+}
+
+#[derive(Clone, Debug)]
+pub struct CompressionTargets<
+    const D: usize,
+>{
+    pub inner_proof: ProofWithPublicInputsTarget<D>,
+    pub verifier_data: VerifierCircuitTarget,
+}
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    C: GenericConfig<D, F = F>,
+    H: AlgebraicHasher<F>,
+> CompressionCircuit<F,D,C,H> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    pub fn new(inner_common_data: CommonCircuitData<F,D>) -> Self {
+        Self{
+            inner_common_data,
+            phantom_data:PhantomData::default(),
+        }
+    }
+
+    /// build the compression circuit
+    pub fn build(&self, builder: &mut CircuitBuilder<F, D>) -> Result<CompressionTargets<D>> {
+
+        let inner_common = self.inner_common_data.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);
+        builder.register_public_inputs(&inner_pub_input[0..4]);
+
+        // virtual target for the verifier data
+        let inner_verifier_data = builder.add_virtual_verifier_data(inner_common.config.fri_config.cap_height);
+
+        // register verifier data hash as public input.
+        let mut vd_pub_input = vec![];
+        vd_pub_input.extend_from_slice(&inner_verifier_data.circuit_digest.elements);
+        for i in 0..builder.config.fri_config.num_cap_elements() {
+            vd_pub_input.extend_from_slice(&inner_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);
+        let mut vd_to_hash = vec![];
+        vd_to_hash.extend_from_slice(&inner_pub_input[4..8]);
+        vd_to_hash.extend_from_slice(&hash_inner_vd_pub_input.elements);
+        let vd_hash = builder.hash_n_to_hash_no_pad::<H>(vd_to_hash);
+        builder.register_public_inputs(&vd_hash.elements);
+
+        // verify the proofs in-circuit
+        builder.verify_proof::<C>(&vir_proof, &inner_verifier_data, &inner_common);
+
+        // return targets
+        let t = CompressionTargets {
+            inner_proof: vir_proof,
+            verifier_data: inner_verifier_data,
+        };
+        Ok(t)
+
+    }
+
+    /// assign the compression targets with given input
+    pub fn assign_targets(
+        &self, pw: &mut PartialWitness<F>,
+        targets: &CompressionTargets<D>,
+        inner_proof: ProofWithPublicInputs<F, C, D>,
+        verifier_only_data: &VerifierOnlyCircuitData<C, D>,
+    ) -> Result<()> {
+        // assign the proof
+        pw.set_proof_with_pis_target(&targets.inner_proof, &inner_proof)
+            .map_err(|e| {
+                CircuitError::ProofTargetAssignmentError("inner-proof".to_string(), e.to_string())
+            })?;
+
+        // assign the verifier data
+        pw.set_verifier_data_target(&targets.verifier_data, verifier_only_data)
+            .map_err(|e| {
+                CircuitError::VerifierDataTargetAssignmentError(e.to_string())
+            })?;
+
+        Ok(())
+    }
+
+    /// returns the compression circuit data
+    pub fn get_circuit_data (&self) -> Result<CircuitData<F, C, D>>
+        where
+            <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+    {
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config.clone());
+
+        self.build(&mut builder)?;
+
+        let circ_data = builder.build::<C>();
+
+        Ok(circ_data)
+    }
+
+}
+
+

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

@@ -7,7 +7,6 @@ use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
 use plonky2::plonk::proof::{ProofWithPublicInputs, ProofWithPublicInputsTarget};
 use plonky2_field::extension::Extendable;
 use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
-// use crate::recursion::circuits::inner_circuit::InnerCircuit;
 use crate::{error::CircuitError,Result};
 
 /// recursion leaf circuit - verifies N inner proof

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

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

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

@@ -7,11 +7,10 @@ use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
 use plonky2::plonk::proof::{ProofWithPublicInputs, ProofWithPublicInputsTarget};
 use plonky2_field::extension::Extendable;
 use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
-// use crate::recursion::circuits::inner_circuit::InnerCircuit;
 use crate::{error::CircuitError,Result};
 use crate::circuits::utils::vec_to_array;
 
-/// recursion node circuit - verifies 2 leaf proofs
+/// recursion node circuit - verifies M leaf proofs
 #[derive(Clone, Debug)]
 pub struct NodeCircuit<
     F: RichField + Extendable<D> + Poseidon2,

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

@@ -0,0 +1,289 @@
+use std::marker::PhantomData;
+use plonky2::hash::hash_types::{HashOut, HashOutTarget, RichField};
+use plonky2::iop::target::Target;
+use plonky2::iop::witness::{PartialWitness, WitnessWrite};
+use plonky2::plonk::circuit_builder::{CircuitBuilder};
+use plonky2::plonk::circuit_data::{CommonCircuitData, VerifierCircuitData, VerifierCircuitTarget};
+use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
+use plonky2::plonk::proof::{ProofWithPublicInputs, ProofWithPublicInputsTarget};
+use plonky2_field::extension::Extendable;
+use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
+use crate::{error::CircuitError, Result};
+
+/// 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.
+/// - `T`: Total Number of inner-proofs.
+/// - `K`: Number of public input field elements per inner-proof (sampling proof).
+pub struct PublicInputVerificationCircuit<
+    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,
+    const K: usize,
+> where
+    <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>,
+{
+    pub node_common_data: CommonCircuitData<F, D>,
+    phantom: PhantomData<(C, H)>,
+}
+
+/// Holds the virtual targets for the circuit.
+/// - `inner_proof`: the proof to be verified and contains the public input to be verified.
+/// - `inner_pub_inputs`: A nested vector of targets with dimensions T×K.
+/// - `node_verifier_data`: Verifier data for the node circuit.
+/// - `leaf_verifier_data`: Verifier data for the leaf circuit.
+/// - `inner_verifier_data`: Verifier data for the inner circuit.
+pub struct PublicInputVerificationTargets<const D: usize> {
+    pub inner_proof: ProofWithPublicInputsTarget<D>,
+    pub node_verifier_data: VerifierCircuitTarget,
+    pub leaf_verifier_data: HashOutTarget,
+    pub inner_verifier_data: HashOutTarget,
+    pub inner_pub_inputs: Vec<Vec<Target>>,
+}
+
+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>
+    where
+        F: RichField + Extendable<D> + Poseidon2,
+        C: GenericConfig<D, F = F>,
+        H: AlgebraicHasher<F>,
+        <C as GenericConfig<D>>::Hasher: AlgebraicHasher<F>,
+{
+    /// Create a new instance of the circuit.
+    pub fn new(node_common_data: CommonCircuitData<F, D>) -> Self {
+        // we expect exactly 8 public inputs from the tree root proof
+        // 4 for the final aggregated public-input hash, 4 for the final aggregated verifier-data hash
+        assert_eq!(node_common_data.num_public_inputs, 8);
+
+        Self {
+            node_common_data,
+            phantom: PhantomData,
+        }
+    }
+
+    /// Builds the circuit by:
+    /// 1. Verifies a proof target with 8 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 tree hashing of node_verifier_date leaf_verifier_data and inner_verifier_data (each 4 field elements).
+    pub fn build(&self, builder: &mut CircuitBuilder<F, D>) -> 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);
+
+        // Create a VerifierCircuitTarget for the node's verifier data (unhashed).
+        let node_verifier_data = builder.add_virtual_verifier_data(
+            self.node_common_data.config.fri_config.cap_height
+        );
+
+        // verify the proof
+        builder.verify_proof::<C>(&inner_proof, &node_verifier_data, &self.node_common_data);
+
+        // create T×K targets for all inner public inputs from the base level.
+        let mut inner_pub_inputs = Vec::with_capacity(T);
+        for _ in 0..T {
+            let mut row = Vec::with_capacity(K);
+            for _ in 0..K {
+                row.push(builder.add_virtual_public_input()); // public input
+            }
+            inner_pub_inputs.push(row);
+        }
+
+        // ------------------------------------------------------------------
+        //   Summary of the logic:
+        //
+        //   let final_pi = proof.public_inputs[0..4];
+        //   let final_vd = proof.public_inputs[4..8];
+        //   ...
+        //   leaf-level pub inputs tree hashing: chunks of N -> hash -> combine with inner_verifier_data
+        //   node-level pub inputs tree hashing: chunks of M -> hash -> combine with either leaf_hash (only level 0) or node_hash
+        //   ...
+        //   check final result matches final_pi, final_vd
+        // ------------------------------------------------------------------
+
+        // 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 final_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(&node_verifier_data.circuit_digest.elements);
+        for cap_elem in node_verifier_data.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);
+        builder.register_public_inputs(&node_hash_t.elements); // public input
+
+
+        let mut pub_in_hashes_t = Vec::new();
+        let mut vd_hashes_t = Vec::new();
+
+        // hash targets for the leaf and inner circuit's verifier data.
+        let leaf_hash_t = builder.add_virtual_hash_public_input(); // public input
+        let inner_hash_t = builder.add_virtual_hash_public_input(); // public input
+
+        // 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]);
+            }
+            // hash
+            let pi_hash_chunk = builder.hash_n_to_hash_no_pad::<H>(chunk_targets);
+
+            // track these in vectors
+            pub_in_hashes_t.push(pi_hash_chunk);
+            vd_hashes_t.push(inner_hash_t);
+        }
+
+        // Now at the node level:
+
+        let mut level = 0;
+        let mut current_len = base_chunks;
+
+        while current_len > 1 {
+
+            let next_len = (current_len + (M - 1)) / M;
+
+            let mut next_pub_in_hashes_t = Vec::with_capacity(next_len);
+            let mut next_vd_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);
+
+                // flatten all pub_in_hashes in [start_idx..end_idx]
+                let mut pi_flat = Vec::with_capacity((end_idx - start_idx) * 4);
+                for j in start_idx..end_idx {
+                    pi_flat.extend_from_slice(&pub_in_hashes_t[j].elements);
+                }
+                let pi_hash = builder.hash_n_to_hash_no_pad::<H>(pi_flat);
+
+                // flatten all vd_hashes in [start_idx..end_idx]
+                let mut vd_flat = Vec::with_capacity((end_idx - start_idx) * 4);
+                for j in start_idx..end_idx {
+                    vd_flat.extend_from_slice(&vd_hashes_t[j].elements);
+                }
+                // use leaf_hash if level == 0, else node_hash
+                let hash_n_t = if level == 0 { leaf_hash_t } else { node_hash_t };
+                vd_flat.extend_from_slice(&hash_n_t.elements);
+
+                let vd_hash = builder.hash_n_to_hash_no_pad::<H>(vd_flat);
+
+                next_pub_in_hashes_t.push(pi_hash);
+                next_vd_hashes_t.push(vd_hash);
+            }
+
+            pub_in_hashes_t = next_pub_in_hashes_t;
+            vd_hashes_t = next_vd_hashes_t;
+            current_len = next_len;
+            level += 1;
+        }
+
+        // now have exactly one pub_in_hash and one vd_hash
+        let final_computed_pi_t = &pub_in_hashes_t[0];
+        let final_computed_vd_t = &vd_hashes_t[0];
+
+        // connect them to the final 8 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_vd_hash_t[i], final_computed_vd_t.elements[i]);
+        }
+
+        // return all the targets
+        Ok(PublicInputVerificationTargets {
+            inner_proof,
+            node_verifier_data,
+            leaf_verifier_data: leaf_hash_t,
+            inner_verifier_data: inner_hash_t,
+            inner_pub_inputs,
+        })
+    }
+
+    /// Assigns witness values to the targets.
+    /// - `inner_proof`: The tree root proof with 8 public inputs [pi_hash, vd_hash].
+    /// - `inner_pub_inputs_vals`: T×K public input values from inner proofs.
+    /// - `node_verifier_data`: node verifier data
+    /// - `leaf_verifier_data`: leaf circuit’s verifier data.
+    /// - `inner_verifier_data`:inner-circuit’s verifier data.
+    pub fn assign_targets(
+        &self,
+        pw: &mut PartialWitness<F>,
+        targets: &PublicInputVerificationTargets<D>,
+        inner_proof: ProofWithPublicInputs<F, C, D>,
+        inner_pub_inputs_vals: Vec<Vec<F>>,
+        node_verifier_data: &VerifierCircuitData<F, C, D>,
+        leaf_verifier_data: &VerifierCircuitData<F, C, D>,
+        inner_verifier_data: &VerifierCircuitData<F, C, D>,
+    ) -> Result<()> {
+        // Assign the final proof - it should have 8 public inputs
+        pw.set_proof_with_pis_target(&targets.inner_proof, &inner_proof)
+            .map_err(|e| {
+                CircuitError::ProofTargetAssignmentError("final-proof".to_string(), e.to_string())
+            })?;
+
+        // Assign T×K inner public inputs
+        if inner_pub_inputs_vals.len() != T {
+            return Err(CircuitError::InvalidArgument(format!(
+                "Expected T={} rows of inner_pub_inputs_vals, got {}",
+                T,
+                inner_pub_inputs_vals.len()
+            )));
+        }
+        for (i, row_vals) in inner_pub_inputs_vals.into_iter().enumerate() {
+            if row_vals.len() != K {
+                return Err(CircuitError::InvalidArgument(format!(
+                    "Expected K={} values in row {}, got {}",
+                    K,
+                    i,
+                    row_vals.len()
+                )));
+            }
+            for (j, val) in row_vals.into_iter().enumerate() {
+                pw.set_target(targets.inner_pub_inputs[i][j], val).map_err(|e| {
+                    CircuitError::TargetAssignmentError(format!("inner public input index [{}][{}]", i,j), e.to_string())
+                })?;
+            }
+        }
+
+        // Assign the node verifier data
+        pw.set_verifier_data_target(&targets.node_verifier_data, &node_verifier_data.verifier_only)
+            .map_err(|e| {
+                CircuitError::VerifierDataTargetAssignmentError(e.to_string())
+            })?;
+
+        // Assign the leaf circuit’s verifier data
+        let leaf_hash = Self::get_hash_of_verifier_data(leaf_verifier_data);
+        pw.set_hash_target(targets.leaf_verifier_data, leaf_hash).map_err(|e| {
+            CircuitError::HashTargetAssignmentError("leaf verifier data hash".to_string(), e.to_string())
+        })?;
+
+        // Assign the inner circuit’s verifier data
+        let inner_hash = Self::get_hash_of_verifier_data(inner_verifier_data);
+        pw.set_hash_target(targets.inner_verifier_data, inner_hash).map_err(|e| {
+            CircuitError::HashTargetAssignmentError("inner verifier data hash".to_string(), e.to_string())
+        })?;
+
+        Ok(())
+    }
+
+    /// helper fn to generate hash of verifier data
+    fn get_hash_of_verifier_data(verifier_data: &VerifierCircuitData<F, C, D>) -> HashOut<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)
+    }
+}

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

@@ -6,10 +6,10 @@ use plonky2::plonk::circuit_data::{CircuitConfig, CircuitData, CommonCircuitData
 use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
 use plonky2::plonk::proof::ProofWithPublicInputs;
 use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
-// use crate::recursion::circuits::inner_circuit::InnerCircuit;
 use plonky2_field::extension::Extendable;
 use crate::{error::CircuitError, Result};
 use crate::recursion::uniform::{leaf::{LeafTargets,LeafCircuit},node::{NodeTargets,NodeCircuit}};
+use crate::recursion::uniform::compress::{CompressionCircuit, CompressionTargets};
 
 /// tree recursion
 pub struct TreeRecursion<
@@ -24,10 +24,13 @@ pub struct TreeRecursion<
 {
     leaf: LeafCircuit<F, D, C, H, N>,
     node: NodeCircuit<F, D, C, H, M>,
+    compression: CompressionCircuit<F, D, C, H>,
     leaf_circ_data: CircuitData<F, C, D>,
     node_circ_data: CircuitData<F, C, D>,
+    compression_circ_data: CircuitData<F, C, D>,
     leaf_targets: LeafTargets<D>,
     node_targets: NodeTargets<D>,
+    compression_targets: CompressionTargets<D>,
     phantom_data: PhantomData<(H)>
 }
 
@@ -63,13 +66,25 @@ impl<
         let node_circ_data = builder.build::<C>();
         // println!("node circuit size = {:?}", node_circ_data.common.degree_bits());
 
+        // compression build
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+        let node_common = node_circ_data.common.clone();
+        let compression_circ = CompressionCircuit::new(node_common);
+        let compression_targets = compression_circ.build(&mut builder)?;
+        let compression_circ_data = builder.build::<C>();
+        // println!("compress circuit size = {:?}", compression_circ_data.common.degree_bits());
+
         Ok(Self{
             leaf,
             node,
+            compression: compression_circ,
             leaf_circ_data,
             node_circ_data,
+            compression_circ_data,
             leaf_targets,
             node_targets,
+            compression_targets,
             phantom_data: Default::default(),
         })
     }
@@ -78,10 +93,34 @@ impl<
         self.leaf_circ_data.verifier_data()
     }
 
+    pub fn get_node_common_data(&self) -> CommonCircuitData<F, D>{
+        self.node_circ_data.common.clone()
+    }
+
+    pub fn get_leaf_common_data(&self) -> CommonCircuitData<F, D>{
+        self.leaf_circ_data.common.clone()
+    }
+
     pub fn get_node_verifier_data(&self) -> VerifierCircuitData<F, C, D>{
         self.node_circ_data.verifier_data()
     }
 
+    pub fn prove_tree_and_compress(
+        &mut self,
+        proofs_with_pi: &[ProofWithPublicInputs<F, C, D>],
+        inner_verifier_only_data: &VerifierOnlyCircuitData<C, D>,
+    ) -> Result<(ProofWithPublicInputs<F, C, D>)>
+    {
+        let proof =
+            self.prove_tree(proofs_with_pi, inner_verifier_only_data)?;
+        let mut pw = PartialWitness::<F>::new();
+        self.compression.assign_targets(&mut pw, &self.compression_targets, proof, &self.node_circ_data.verifier_only)?;
+
+        self.compression_circ_data.prove(pw).map_err(
+            |e| CircuitError::InvalidProofError(e.to_string())
+        )
+    }
+
     pub fn prove_tree
     (
         &mut self,
@@ -162,19 +201,46 @@ impl<
         self.prove(&new_proofs, &self.node_circ_data.verifier_only)
     }
 
+    pub fn verify_proof(
+        &self,
+        proof: ProofWithPublicInputs<F, C, D>,
+        is_compressed: bool,
+    ) -> Result<()>{
+        if is_compressed{
+            self.compression_circ_data.verify(proof)
+                .map_err(|e| CircuitError::InvalidProofError(e.to_string()))
+        }else {
+            self.node_circ_data.verify(proof)
+                .map_err(|e| CircuitError::InvalidProofError(e.to_string()))
+        }
+    }
+
     pub fn verify_proof_and_public_input(
         &self,
         proof: ProofWithPublicInputs<F, C, D>,
         inner_public_input: Vec<Vec<F>>,
-        inner_verifier_data: &VerifierCircuitData<F, C, D>) -> Result<()>
-    {
+        inner_verifier_data: &VerifierCircuitData<F, C, D>,
+        is_compressed: bool,
+    ) -> Result<()>{
         let public_input = proof.public_inputs.clone();
-        self.node_circ_data.verify(proof)
-            .map_err(|e| CircuitError::InvalidProofError(e.to_string()))?;
-        self.verify_public_input(public_input, inner_public_input, inner_verifier_data)
+        if is_compressed{
+            self.compression_circ_data.verify(proof)
+                .map_err(|e| CircuitError::InvalidProofError(e.to_string()))?;
+            self.verify_public_input(public_input, inner_public_input, inner_verifier_data, is_compressed)
+        }else {
+            self.node_circ_data.verify(proof)
+                .map_err(|e| CircuitError::InvalidProofError(e.to_string()))?;
+            self.verify_public_input(public_input, inner_public_input, inner_verifier_data, is_compressed)
+        }
     }
 
-    pub fn verify_public_input(&self, public_input: Vec<F>, inner_public_input: Vec<Vec<F>>, inner_verifier_data: &VerifierCircuitData<F, C, D>) -> Result<()>{
+    pub fn verify_public_input(
+        &self,
+        public_input: Vec<F>,
+        inner_public_input: Vec<Vec<F>>,
+        inner_verifier_data: &VerifierCircuitData<F, C, D>,
+        is_compressed: bool,
+    ) -> Result<()>{
         assert_eq!(public_input.len(), 8);
 
         let given_input_hash = &public_input[0..4];
@@ -227,7 +293,14 @@ impl<
 
         //check expected hash
         let expected_pi_hash = pub_in_hashes[0];
-        let expected_vd_hash = inner_vd_hashes[0];
+        let mut expected_vd_hash = inner_vd_hashes[0];
+
+        if is_compressed {
+            let mut vd_to_hash = vec![];
+            vd_to_hash.extend_from_slice(&expected_vd_hash.elements);
+            vd_to_hash.extend_from_slice(&node_hash.elements);
+            expected_vd_hash = H::hash_no_pad(&vd_to_hash);
+        }
 
         assert_eq!(given_input_hash, expected_pi_hash.elements);
         assert_eq!(given_vd_hash, expected_vd_hash.elements);

proof-input/Cargo.toml

@@ -14,4 +14,8 @@ plonky2  = { workspace = true }
 plonky2_field = { workspace = true }
 # --- local ---
 plonky2_poseidon2 = { path = "../plonky2_poseidon2" }
-codex-plonky2-circuits = { path = "../codex-plonky2-circuits" }
+codex-plonky2-circuits = { path = "../codex-plonky2-circuits" }
+
+[features]
+default = []
+parallel = ["plonky2/parallel"]

proof-input/src/recursion/uniform.rs

@@ -1,8 +1,7 @@
-// some tests for approach 2 of the tree recursion
+// some tests for the tree recursion
 
 #[cfg(test)]
 mod tests {
-    use std::fs;
     use plonky2::iop::witness::{PartialWitness};
     use plonky2::plonk::circuit_builder::CircuitBuilder;
     use plonky2::plonk::circuit_data::{CircuitConfig};
@@ -13,6 +12,7 @@ mod tests {
     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;
 
     #[test]
     fn test_uniform_recursion() -> anyhow::Result<()> {
@@ -35,7 +35,8 @@ mod tests {
         println!("sampling circuit degree bits = {:?}", inner_data.common.degree_bits());
         let inner_proof = inner_data.prove(pw)?;
 
-        let proofs: Vec<ProofWithPublicInputs<F, C, D>> = (0..4).map(|i| inner_proof.clone()).collect();
+        let num_of_proofs = 4;
+        let proofs: Vec<ProofWithPublicInputs<F, C, D>> = (0..num_of_proofs).map(|i| inner_proof.clone()).collect();
 
         // ------------------- tree --------------------
         const N: usize = 1;
@@ -45,14 +46,106 @@ mod tests {
 
         let root = tree.prove_tree(&proofs, &inner_data.verifier_only)?;
         println!("pub input size = {}", root.public_inputs.len());
+        println!("proof size = {:?} bytes", root.to_bytes().len());
+
+        let root_compressed = tree.prove_tree_and_compress(&proofs, &inner_data.verifier_only)?;
+        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,&inner_data.verifier_data()).is_ok(),
+            tree.verify_proof_and_public_input(root,inner_pi.clone(),&inner_data.verifier_data(), false).is_ok(),
             "proof verification failed"
         );
 
+        assert!(
+            tree.verify_proof_and_public_input(root_compressed,inner_pi,&inner_data.verifier_data(), true).is_ok(),
+            "compressed proof verification failed"
+        );
+
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_pi_verifier() -> anyhow::Result<()> {
+
+        let config = CircuitConfig::standard_recursion_config();
+        let mut sampling_builder = CircuitBuilder::<F, D>::new(config);
+
+        //------------ 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 = samp_circ.sample_slot_circuit_with_public_input(&mut sampling_builder)?;
+        // get generate a sampling proof
+        let mut pw = PartialWitness::<F>::new();
+        samp_circ.sample_slot_assign_witness(&mut pw,&inner_tar,&one_circ_input)?;
+        let inner_data = sampling_builder.build::<C>();
+        println!("sampling circuit degree bits = {:?}", inner_data.common.degree_bits());
+        let inner_proof = inner_data.prove(pw)?;
+
+        // 9 field elems as public inputs in the sampling circuit
+        const K:usize = 9;
+        // change the following as needed.
+        const T: usize = 4;
+        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(inner_data.common.clone())?;
+
+        let root = tree.prove_tree(&proofs, &inner_data.verifier_only)?;
+        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(),&inner_data.verifier_data(), 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());
+
+        let config = CircuitConfig::standard_recursion_config();
+        let mut builder = CircuitBuilder::<F, D>::new(config);
+        let pi_tarq = pi_verifier_circ.build(&mut builder)?;
+
+        let pi_circ_data = builder.build::<C>();
+        println!("PI verifier circuit degree bits = {:?}", pi_circ_data.common.degree_bits());
+
+        let mut pw = PartialWitness::<F>::new();
+
+        pi_verifier_circ.assign_targets(&mut pw, &pi_tarq, root, inner_pi.clone(), &tree.get_node_verifier_data(), &tree.get_leaf_verifier_data(), &inner_data.verifier_data())?;
+
+        let proof = pi_circ_data.prove(pw)?;
+        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();
+        let num_pi = proof.public_inputs.len();
+
+        // sanity check
+        for (i, e) in proof.public_inputs.iter().enumerate(){
+            if i < pub_input_flat.len() {
+                assert_eq!(*e, pub_input_flat[i])
+            }
+        }
+
+        assert!(
+            pi_circ_data.verify(proof).is_ok(),
+            "pi-verifier proof verification failed"
+        );
+
         Ok(())
     }
 }

workflow/benches/compression.rs

@@ -0,0 +1,111 @@
+use anyhow::{anyhow, Result};
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use plonky2::gates::noop::NoopGate;
+use plonky2::hash::hash_types::HashOut;
+use plonky2::iop::witness::{PartialWitness, WitnessWrite};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::CircuitConfig;
+use plonky2::plonk::config::GenericConfig;
+
+use codex_plonky2_circuits::recursion::uniform::compress::CompressionCircuit;
+use proof_input::params::{D, C, F, HF};
+
+/// Benchmark for building, proving, and verifying the Plonky2 circuit.
+fn bench_compression_runtime(c: &mut Criterion, circuit_size: usize) -> Result<()>{
+
+    #[cfg(feature = "parallel")]
+    println!("Parallel feature is ENABLED");
+
+    // Create the circuit configuration
+    let config = CircuitConfig::standard_recursion_config();
+    let mut builder = CircuitBuilder::<F, D>::new(config);
+
+    let num_dummy_gates = match circuit_size {
+        0 => return Err(anyhow!("size must be at least 1")),
+        1 => 0,
+        2 => 1,
+        n => (1 << (n - 1)) + 1,
+    };
+
+    for _ in 0..num_dummy_gates {
+        builder.add_gate(NoopGate, vec![]);
+    }
+
+    // 2 virtual hashes (8 field elems) as public input - same as in the recursion tree
+    let mut pi = vec![];
+    for i in 0..2{
+        pi.push(builder.add_virtual_hash_public_input());
+    }
+
+    let inner_data = builder.build::<C>();
+    println!("inner circuit size = {:?}", inner_data.common.degree_bits());
+
+    // prove with dummy public input
+    let mut pw = PartialWitness::<F>::new();
+    pw.set_hash_target(pi[0], HashOut::<F>::ZERO)?;
+    pw.set_hash_target(pi[1], HashOut::<F>::ZERO)?;
+    let inner_proof = inner_data.prove(pw)?;
+
+    // Compression circuit
+    let config = CircuitConfig::standard_recursion_config();
+    let mut builder = CircuitBuilder::<F, D>::new(config);
+    let compression_circ = CompressionCircuit::<F,D,C,HF>::new(inner_data.common.clone());
+    let compression_targets = compression_circ.build(&mut builder)?;
+
+    // Benchmark Group
+    let mut group = c.benchmark_group(format!("Compression Circuit Benchmark for inner-proof size = {}", circuit_size));
+
+    // Benchmark the Circuit Building Phase
+    group.bench_function("Build Circuit", |b| {
+        b.iter(|| {
+            let config = CircuitConfig::standard_recursion_config();
+            let mut local_builder = CircuitBuilder::<F, D>::new(config);
+            let _compression_targets = compression_circ.build(&mut local_builder);
+            let _data = local_builder.build::<C>();
+        })
+    });
+
+    // Build the circuit once for proving and verifying benchmarks
+    let compression_circ_data = builder.build::<C>();
+    println!("compress circuit size = {:?}", compression_circ_data.common.degree_bits());
+
+    let mut pw = PartialWitness::<F>::new();
+    compression_circ.assign_targets(&mut pw, &compression_targets, inner_proof, &inner_data.verifier_only)?;
+
+    group.bench_function("Prove Circuit", |b| {
+        b.iter_batched(
+            || pw.clone(),
+            |local_pw| compression_circ_data.prove(local_pw).expect("Failed to prove circuit"),
+            BatchSize::SmallInput,
+        )
+    });
+
+    let proof = compression_circ_data.prove(pw)?;
+    println!("Proof size: {} bytes", proof.to_bytes().len());
+
+    let verifier_data = compression_circ_data.verifier_data();
+
+    // Benchmark the Verifying Phase
+    group.bench_function("Verify Proof", |b| {
+        b.iter(|| {
+            verifier_data.verify(proof.clone()).expect("Failed to verify proof");
+        })
+    });
+
+    group.finish();
+    Ok(())
+}
+
+fn bench_compression(c: &mut Criterion) -> Result<()>{
+    bench_compression_runtime(c, 13)?;
+    bench_compression_runtime(c, 14)?;
+    Ok(())
+}
+
+/// Criterion benchmark group
+criterion_group!{
+    name = benches;
+    config = Criterion::default().sample_size(10);
+    targets = bench_compression
+}
+criterion_main!(benches);