From c150525f158f930ff232fe7c9f10e05b3940da18 Mon Sep 17 00:00:00 2001
From: M Alghazwi <m.ghazwi@gmail.com>
Date: Wed, 21 May 2025 10:06:35 +0200
Subject: [PATCH] fix compression and refactor

---
 .../src/recursion/compress.rs                 | 42 ++++---------------
 codex-plonky2-circuits/src/recursion/tree.rs  |  6 +--
 2 files changed, 11 insertions(+), 37 deletions(-)

diff --git a/codex-plonky2-circuits/src/recursion/compress.rs b/codex-plonky2-circuits/src/recursion/compress.rs
index 62f1fb6..9ab3374 100644
--- a/codex-plonky2-circuits/src/recursion/compress.rs
+++ b/codex-plonky2-circuits/src/recursion/compress.rs
@@ -2,7 +2,7 @@ use std::marker::PhantomData;
 use plonky2::hash::hash_types::{HashOutTarget, 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::{CommonCircuitData, VerifierCircuitData, VerifierOnlyCircuitData};
 use plonky2::plonk::config::{AlgebraicHasher, GenericConfig};
 use plonky2::plonk::proof::{ProofWithPublicInputs, ProofWithPublicInputsTarget};
 use plonky2_field::extension::Extendable;
@@ -18,13 +18,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)]
@@ -47,18 +44,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(),
         }
     }
 
@@ -68,44 +61,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!(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);
diff --git a/codex-plonky2-circuits/src/recursion/tree.rs b/codex-plonky2-circuits/src/recursion/tree.rs
index 1a8eb57..6654d01 100644
--- a/codex-plonky2-circuits/src/recursion/tree.rs
+++ b/codex-plonky2-circuits/src/recursion/tree.rs
@@ -30,7 +30,7 @@ pub struct TreeRecursion<
 {
     leaf: LeafCircuit<F, D, C, H, N, T>,
     node: NodeCircuit<F, D, C, H, M, T>,
-    compression: CompressionCircuit<F, D, C, H>,
+    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>,
@@ -80,8 +80,8 @@ impl<
         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 node_common = node_circ_data.common.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());