add bn254 wrapper

codex-plonky2-circuits/Cargo.toml

@@ -18,6 +18,9 @@ plonky2_poseidon2 = { path = "../plonky2_poseidon2" }
 itertools = { workspace = true }
 plonky2_maybe_rayon = { workspace = true }
 hashbrown = "0.14.5"
+ff = { package = "ff", version = "0.13", features = ["derive"] }
+num = "0.4.3"
+lazy_static = "1.5.0"
 
 [dev-dependencies]
 criterion = { version = "0.5.1", default-features = false }

codex-plonky2-circuits/src/bn254_wrapper/bn254_fr.rs

@@ -0,0 +1,7 @@
+use ff::PrimeField;
+
+#[derive(PrimeField)]
+#[PrimeFieldModulus = "21888242871839275222246405745257275088548364400416034343698204186575808495617"]
+#[PrimeFieldGenerator = "7"]
+#[PrimeFieldReprEndianness = "little"]
+pub struct Fr([u64; 4]);

codex-plonky2-circuits/src/bn254_wrapper/config.rs

@@ -0,0 +1,238 @@
+use core::fmt;
+use std::error::Error;
+use std::marker::PhantomData;
+
+use ff::{Field as ff_Field, PrimeField};
+use num::BigUint;
+use plonky2::field::extension::quadratic::QuadraticExtension;
+use plonky2::field::goldilocks_field::GoldilocksField;
+use plonky2::field::types::Field;
+use plonky2::hash::hash_types::RichField;
+use plonky2::hash::poseidon::{PoseidonHash, PoseidonPermutation};
+use plonky2::plonk::config::{GenericConfig, GenericHashOut, Hasher};
+use serde::de::Visitor;
+use serde::{Deserialize, Deserializer, Serialize, Serializer};
+
+use crate::bn254_wrapper::poseidon_bn254::{permution, GOLDILOCKS_ELEMENTS, RATE};
+use crate::bn254_wrapper::bn254_fr::{Fr, FrRepr};
+
+/// Configuration using Poseidon BN254 over the Goldilocks field.
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Serialize)]
+pub struct PoseidonBN254GoldilocksConfig;
+impl GenericConfig<2> for PoseidonBN254GoldilocksConfig {
+    type F = GoldilocksField;
+    type FE = QuadraticExtension<Self::F>;
+    type Hasher = PoseidonBN254Hash;
+    type InnerHasher = PoseidonHash;
+}
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub struct PoseidonBN254HashOut<F: Field> {
+    pub value: Fr,
+    _phantom: PhantomData<F>,
+}
+
+fn hash_out_to_bytes<F: Field>(hash: PoseidonBN254HashOut<F>) -> Vec<u8> {
+    let binding = hash.value.to_repr();
+    let limbs = binding.as_ref();
+    limbs.to_vec()
+}
+
+impl<F: RichField> GenericHashOut<F> for PoseidonBN254HashOut<F> {
+    fn to_bytes(&self) -> Vec<u8> {
+        hash_out_to_bytes(*self)
+    }
+
+    fn from_bytes(bytes: &[u8]) -> Self {
+        let sized_bytes: [u8; 32] = bytes.try_into().unwrap();
+        let fr_repr = FrRepr(sized_bytes);
+        let fr = Fr::from_repr(fr_repr).unwrap();
+
+        Self {
+            value: fr,
+            _phantom: PhantomData,
+        }
+    }
+
+    fn to_vec(&self) -> Vec<F> {
+        let bytes = hash_out_to_bytes(*self);
+        bytes
+            // Chunks of 7 bytes since 8 bytes would allow collisions.
+            .chunks(7)
+            .map(|bytes| {
+                let mut arr = [0; 8];
+                arr[..bytes.len()].copy_from_slice(bytes);
+                F::from_canonical_u64(u64::from_le_bytes(arr))
+            })
+            .collect()
+    }
+}
+
+impl<F: RichField> Serialize for PoseidonBN254HashOut<F> {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        // Output the hash as a bigint string.
+        let binding = self.value.to_repr();
+        let limbs = binding.as_ref();
+
+        let big_int = BigUint::from_bytes_le(limbs);
+        serializer.serialize_str(big_int.to_str_radix(10).as_str())
+    }
+}
+
+impl<'de, F: RichField> Deserialize<'de> for PoseidonBN254HashOut<F> {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: Deserializer<'de>,
+    {
+        struct PoseidonBN254HashOutVisitor;
+
+        impl<'a> Visitor<'a> for PoseidonBN254HashOutVisitor {
+            type Value = String;
+
+            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
+                formatter.write_str("a string with integer value within BN254 scalar field")
+            }
+
+            fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
+            where
+                E: Error,
+            {
+                Ok(v.to_string())
+            }
+        }
+
+        let deserialized_str = deserializer
+            .deserialize_str(PoseidonBN254HashOutVisitor)
+            .unwrap();
+        let big_int = BigUint::parse_bytes(deserialized_str.as_bytes(), 10).unwrap();
+
+        let mut bytes = big_int.to_bytes_le();
+        for _i in bytes.len()..32 {
+            bytes.push(0);
+        }
+
+        let sized_bytes: [u8; 32] = bytes.try_into().unwrap();
+        let fr_repr = FrRepr(sized_bytes);
+        let fr = Fr::from_repr(fr_repr).unwrap();
+
+        Ok(Self {
+            value: fr,
+            _phantom: PhantomData,
+        })
+    }
+}
+
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub struct PoseidonBN254Hash;
+impl<F: RichField> Hasher<F> for PoseidonBN254Hash {
+    const HASH_SIZE: usize = 32; // Hash output is 4 limbs of u64
+    type Hash = PoseidonBN254HashOut<F>;
+    type Permutation = PoseidonPermutation<F>;
+
+    fn hash_no_pad(input: &[F]) -> Self::Hash {
+        let mut state = [Fr::ZERO; 4];
+
+        state[0] = Fr::ZERO;
+        for rate_chunk in input.chunks(RATE * 3) {
+            for (j, bn254_chunk) in rate_chunk.chunks(3).enumerate() {
+                let mut bytes = bn254_chunk[0].to_canonical_u64().to_le_bytes().to_vec();
+
+                for gl_element in bn254_chunk.iter().skip(1) {
+                    let chunk_bytes = gl_element.to_canonical_u64().to_le_bytes();
+                    bytes.extend_from_slice(&chunk_bytes);
+                }
+
+                for _i in bytes.len()..32 {
+                    bytes.push(0);
+                }
+
+                let sized_bytes: [u8; 32] = bytes.try_into().unwrap();
+                let fr_repr = FrRepr(sized_bytes);
+                state[j + 1] = Fr::from_repr(fr_repr).unwrap();
+            }
+            permution(&mut state);
+        }
+
+        PoseidonBN254HashOut {
+            value: state[0],
+            _phantom: PhantomData,
+        }
+    }
+
+    fn hash_pad(input: &[F]) -> Self::Hash {
+        let mut padded_input = input.to_vec();
+        padded_input.push(F::ONE);
+        while (padded_input.len() + 1) % (RATE * GOLDILOCKS_ELEMENTS) != 0 {
+            padded_input.push(F::ZERO);
+        }
+        padded_input.push(F::ONE);
+        Self::hash_no_pad(&padded_input)
+    }
+
+    fn hash_or_noop(inputs: &[F]) -> Self::Hash {
+        if inputs.len() * 8 <= GOLDILOCKS_ELEMENTS * 8 {
+            let mut inputs_bytes = vec![0u8; 32];
+            for i in 0..inputs.len() {
+                inputs_bytes[i * 8..(i + 1) * 8]
+                    .copy_from_slice(&inputs[i].to_canonical_u64().to_le_bytes());
+            }
+            Self::Hash::from_bytes(&inputs_bytes)
+        } else {
+            Self::hash_no_pad(inputs)
+        }
+    }
+
+    fn two_to_one(left: Self::Hash, right: Self::Hash) -> Self::Hash {
+        let mut state = [Fr::ZERO, Fr::ZERO, left.value, right.value];
+        permution(&mut state);
+
+        PoseidonBN254HashOut {
+            value: state[0],
+            _phantom: PhantomData,
+        }
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    use super::*;
+
+    #[test]
+    fn test_byte_methods() {
+        type F = GoldilocksField;
+
+        let fr = Fr::from_str_vartime(
+            "11575173631114898451293296430061690731976535592475236587664058405912382527658",
+        )
+        .unwrap();
+        let hash = PoseidonBN254HashOut::<F> {
+            value: fr,
+            _phantom: PhantomData,
+        };
+
+        let bytes = hash.to_bytes();
+
+        let hash_from_bytes = PoseidonBN254HashOut::<F>::from_bytes(&bytes);
+        assert_eq!(hash, hash_from_bytes);
+    }
+
+    #[test]
+    fn test_serialization() {
+        let fr = Fr::from_str_vartime(
+            "11575173631114898451293296430061690731976535592475236587664058405912382527658",
+        )
+        .unwrap();
+        let hash = PoseidonBN254HashOut::<GoldilocksField> {
+            value: fr,
+            _phantom: PhantomData,
+        };
+
+        let serialized = serde_json::to_string(&hash).unwrap();
+        let deserialized: PoseidonBN254HashOut<GoldilocksField> =
+            serde_json::from_str(&serialized).unwrap();
+        assert_eq!(hash, deserialized);
+    }
+}

codex-plonky2-circuits/src/bn254_wrapper/mod.rs

@@ -0,0 +1,5 @@
+pub mod config;
+pub mod bn254_fr;
+pub mod poseidon_bn254_constants;
+pub mod poseidon_bn254;
+pub mod wrap;

codex-plonky2-circuits/src/bn254_wrapper/poseidon_bn254.rs

@@ -0,0 +1,268 @@
+use std::ops::{AddAssign, MulAssign};
+
+use ff::Field;
+
+use crate::bn254_wrapper::poseidon_bn254_constants::{
+    C_CONSTANTS, M_MATRIX, P_MATRIX, S_CONSTANTS,
+};
+use crate::bn254_wrapper::bn254_fr::Fr;
+
+pub const RATE: usize = 3;
+pub const WIDTH: usize = 4;
+pub const FULL_ROUNDS: usize = 8;
+pub const PARTIAL_ROUNDS: usize = 56;
+pub const GOLDILOCKS_ELEMENTS: usize = 3;
+
+pub type PoseidonState = [Fr; WIDTH];
+
+// This poseidon BN254 implementation is based on the following implementation:
+// https://github.com/iden3/go-iden3-crypto/blob/e5cf066b8be3da9a3df9544c65818df189fdbebe/poseidon/poseidon.go
+pub fn permution(state: &mut PoseidonState) {
+    ark(state, 0);
+    full_rounds(state, true);
+    partial_rounds(state);
+    full_rounds(state, false);
+}
+
+fn ark(state: &mut PoseidonState, it: usize) {
+    for i in 0..WIDTH {
+        state[i].add_assign(&C_CONSTANTS[it + i]);
+    }
+}
+
+fn exp5(mut x: Fr) -> Fr {
+    let aux = x;
+    x = x.square();
+    x = x.square();
+    x.mul_assign(&aux);
+
+    x
+}
+
+fn exp5_state(state: &mut PoseidonState) {
+    for state_element in state.iter_mut().take(WIDTH) {
+        *state_element = exp5(*state_element);
+    }
+}
+
+fn full_rounds(state: &mut PoseidonState, first: bool) {
+    for i in 0..FULL_ROUNDS / 2 - 1 {
+        exp5_state(state);
+        if first {
+            ark(state, (i + 1) * WIDTH);
+        } else {
+            ark(
+                state,
+                (FULL_ROUNDS / 2 + 1) * WIDTH + PARTIAL_ROUNDS + i * WIDTH,
+            );
+        }
+        mix(state, &M_MATRIX);
+    }
+
+    exp5_state(state);
+    if first {
+        ark(state, (FULL_ROUNDS / 2) * WIDTH);
+        mix(state, &P_MATRIX);
+    } else {
+        mix(state, &M_MATRIX);
+    }
+}
+
+fn partial_rounds(state: &mut PoseidonState) {
+    for i in 0..PARTIAL_ROUNDS {
+        state[0] = exp5(state[0]);
+        state[0].add_assign(&C_CONSTANTS[(FULL_ROUNDS / 2 + 1) * WIDTH + i]);
+
+        let mut mul;
+        let mut new_state0 = Fr::ZERO;
+        for j in 0..WIDTH {
+            mul = Fr::ZERO;
+            mul.add_assign(&S_CONSTANTS[(WIDTH * 2 - 1) * i + j]);
+            mul.mul_assign(&state[j]);
+            new_state0.add_assign(&mul);
+        }
+
+        for k in 1..WIDTH {
+            mul = Fr::ZERO;
+            mul.add_assign(&state[0]);
+            mul.mul_assign(&S_CONSTANTS[(WIDTH * 2 - 1) * i + WIDTH + k - 1]);
+            state[k].add_assign(&mul);
+        }
+
+        state[0] = new_state0;
+    }
+}
+
+fn mix(state: &mut PoseidonState, constant_matrix: &[Vec<Fr>]) {
+    let mut result: PoseidonState = [Fr::ZERO; WIDTH];
+
+    let mut mul;
+    for (i, result_element) in result.iter_mut().enumerate().take(WIDTH) {
+        for j in 0..WIDTH {
+            mul = Fr::ZERO;
+            mul.add_assign(&constant_matrix[j][i]);
+            mul.mul_assign(&state[j]);
+            result_element.add_assign(&mul);
+        }
+    }
+
+    state[..WIDTH].copy_from_slice(&result[..WIDTH]);
+}
+
+#[cfg(test)]
+mod permutation_tests {
+    use anyhow::Ok;
+    use ff::{Field, PrimeField};
+
+    use super::{permution, WIDTH};
+    use crate::bn254_wrapper::bn254_fr::Fr;
+
+    #[test]
+    fn test_permuation() -> Result<(), anyhow::Error> {
+        // Test inputs are:
+        // 1. all zeros
+        // 2. range 0..WIDTH
+        // 3. all max BN254 values
+        // 4. random elements of BN254.
+        // Expected output calculated from this poseidon implementation:  https://github.com/iden3/go-iden3-crypto/blob/master/poseidon/poseidon.go#L65
+
+        let max_value: Fr = Fr::from_str_vartime(
+            "21888242871839275222246405745257275088548364400416034343698204186575808495616",
+        )
+        .unwrap();
+
+        let test_vectors: Vec<([Fr; 4], [Fr; 4])> = vec![
+            (
+                [Fr::ZERO; 4],
+                [
+                    Fr::from_str_vartime("5317387130258456662214331362918410991734007599705406860481038345552731150762").unwrap(),
+                    Fr::from_str_vartime("17768273200467269691696191901389126520069745877826494955630904743826040320364").unwrap(),
+                    Fr::from_str_vartime("19413739268543925182080121099097652227979760828059217876810647045303340666757").unwrap(),
+                    Fr::from_str_vartime("3717738800218482999400886888123026296874264026760636028937972004600663725187").unwrap(),
+                ]
+            ),
+            (
+                [
+                    Fr::from_str_vartime("0").unwrap(),
+                    Fr::from_str_vartime("1").unwrap(),
+                    Fr::from_str_vartime("2").unwrap(),
+                    Fr::from_str_vartime("3").unwrap(),
+                ],
+                [
+                    Fr::from_str_vartime("6542985608222806190361240322586112750744169038454362455181422643027100751666").unwrap(),
+                    Fr::from_str_vartime("3478427836468552423396868478117894008061261013954248157992395910462939736589").unwrap(),
+                    Fr::from_str_vartime("1904980799580062506738911865015687096398867595589699208837816975692422464009").unwrap(),
+                    Fr::from_str_vartime("11971464497515232077059236682405357499403220967704831154657374522418385384151").unwrap(),
+                ]
+            ),
+            (
+                [max_value; 4],
+                [
+                    Fr::from_str_vartime("13055670547682322550638362580666986963569035646873545133474324633020685301274").unwrap(),
+                    Fr::from_str_vartime("19087936485076376314486368416882351797015004625427655501762827988254486144933").unwrap(),
+                    Fr::from_str_vartime("10391468779200270580383536396630001155994223659670674913170907401637624483385").unwrap(),
+                    Fr::from_str_vartime("17202557688472898583549180366140168198092766974201433936205272956998081177816").unwrap(),
+                ]
+            ),
+            (
+                [
+                    Fr::from_str_vartime("6542985608222806190361240322586112750744169038454362455181422643027100751666").unwrap(),
+                    Fr::from_str_vartime("3478427836468552423396868478117894008061261013954248157992395910462939736589").unwrap(),
+                    Fr::from_str_vartime("1904980799580062506738911865015687096398867595589699208837816975692422464009").unwrap(),
+                    Fr::from_str_vartime("11971464497515232077059236682405357499403220967704831154657374522418385384151").unwrap(),
+                ],
+                [
+                    Fr::from_str_vartime("21792249080447013894140672594027696524030291802493510986509431008224624594361").unwrap(),
+                    Fr::from_str_vartime("3536096706123550619294332177231935214243656967137545251021848527424156573335").unwrap(),
+                    Fr::from_str_vartime("14869351042206255711434675256184369368509719143073814271302931417334356905217").unwrap(),
+                    Fr::from_str_vartime("5027523131326906886284185656868809493297314443444919363729302983434650240523").unwrap(),
+                ]
+            ),
+        ];
+
+        for (mut input, expected_output) in test_vectors.into_iter() {
+            permution(&mut input);
+            for i in 0..WIDTH {
+                assert_eq!(input[i], expected_output[i]);
+            }
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod merkle_tree_tests {
+    use anyhow::Result;
+    use plonky2::field::extension::Extendable;
+    use plonky2::hash::hash_types::RichField;
+    use plonky2::hash::merkle_proofs::verify_merkle_proof_to_cap;
+    use plonky2::hash::merkle_tree::MerkleTree;
+    use plonky2::plonk::config::GenericConfig;
+
+    use crate::bn254_wrapper::config::PoseidonBN254GoldilocksConfig;
+
+    fn random_data<F: RichField>(n: usize, k: usize) -> Vec<Vec<F>> {
+        (0..n).map(|_| F::rand_vec(k)).collect()
+    }
+
+    fn verify_all_leaves<
+        F: RichField + Extendable<D>,
+        C: GenericConfig<D, F = F>,
+        const D: usize,
+    >(
+        leaves: Vec<Vec<F>>,
+        cap_height: usize,
+    ) -> Result<()> {
+        let tree = MerkleTree::<F, C::Hasher>::new(leaves.clone(), cap_height);
+        for (i, leaf) in leaves.into_iter().enumerate() {
+            let proof = tree.prove(i);
+            verify_merkle_proof_to_cap(leaf, i, &tree.cap, &proof)?;
+        }
+        Ok(())
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_cap_height_too_big() {
+        const D: usize = 2;
+        type C = PoseidonBN254GoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+
+        let log_n = 8;
+        let cap_height = log_n + 1; // Should panic if `cap_height > len_n`.
+
+        let leaves = random_data::<F>(1 << log_n, 7);
+        let _ = MerkleTree::<F, <C as GenericConfig<D>>::Hasher>::new(leaves, cap_height);
+    }
+
+    #[test]
+    fn test_cap_height_eq_log2_len() -> Result<()> {
+        const D: usize = 2;
+        type C = PoseidonBN254GoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+
+        let log_n = 8;
+        let n = 1 << log_n;
+        let leaves = random_data::<F>(n, 7);
+
+        verify_all_leaves::<F, C, D>(leaves, log_n)?;
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_merkle_trees() -> Result<()> {
+        const D: usize = 2;
+        type C = PoseidonBN254GoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+
+        let log_n = 8;
+        let n = 1 << log_n;
+        let leaves = random_data::<F>(n, 7);
+
+        verify_all_leaves::<F, C, D>(leaves, 1)?;
+
+        Ok(())
+    }
+}

codex-plonky2-circuits/src/bn254_wrapper/wrap.rs

@@ -0,0 +1,223 @@
+use std::fs::{self, File};
+use std::marker::PhantomData;
+use std::path::Path;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::witness::{PartialWitness, WitnessWrite};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+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;
+use serde::Serialize;
+use plonky2_poseidon2::Poseidon2;
+use crate::circuit_helper::Plonky2Circuit;
+use crate::error::CircuitError;
+
+/// Wrap circuit - wraps the plonky2 proof with
+/// InnerParameters: Config params for the inner proof - this is the default config
+/// OuterParameters: Config params for the outer proof - this is the bn254 config
+#[derive(Debug)]
+pub struct WrapCircuit<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    InnerParameters: GenericConfig<D, F = F>,
+    OuterParameters: GenericConfig<D, F = F>,
+> {
+    inner_verifier_data: VerifierCircuitData<F, InnerParameters, D>,
+    phantom_data: PhantomData<OuterParameters>
+}
+
+#[derive(Clone, Debug)]
+pub struct WrapTargets<
+    const D: usize,
+>{
+    pub inner_proof: ProofWithPublicInputsTarget<D>,
+}
+
+#[derive(Clone, Debug)]
+pub struct WrapInput<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    InnerParameters: GenericConfig<D, F = F>,
+>{
+    pub inner_proof: ProofWithPublicInputs<F, InnerParameters, D>
+}
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    InnerParameters: GenericConfig<D, F = F>,
+    OuterParameters: GenericConfig<D, F = F>,
+    > Plonky2Circuit<F,OuterParameters,D> for WrapCircuit<F, D, InnerParameters, OuterParameters>
+    where
+        <InnerParameters as GenericConfig<D>>::Hasher: AlgebraicHasher<F>
+{
+    type Targets = WrapTargets<D>;
+    type Input = WrapInput<F, D, InnerParameters>;
+
+    fn add_targets(&self, builder: &mut CircuitBuilder<F, D>, register_pi: bool) -> crate::Result<Self::Targets> {
+        let inner_common = self.inner_verifier_data.common.clone();
+
+        // the proof virtual targets
+        let vir_proof = builder.add_virtual_proof_with_pis(&inner_common);
+        // make inner public input into outer public input
+        if register_pi {
+            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.verifier_only);
+
+        // verify the proofs in-circuit
+        builder.verify_proof::<InnerParameters>(&vir_proof, &const_verifier_data, &inner_common);
+
+        Ok(
+            WrapTargets{
+            inner_proof:vir_proof,
+            }
+        )
+    }
+
+    fn assign_targets(&self, pw: &mut PartialWitness<F>, targets: &Self::Targets, input: &Self::Input) -> crate::Result<()> where
+        <InnerParameters as GenericConfig<D>>::Hasher: AlgebraicHasher<F> {
+        // assign the proof
+        pw.set_proof_with_pis_target(&targets.inner_proof, &input.inner_proof)
+            .map_err(|e| {
+                CircuitError::ProofTargetAssignmentError("inner-proof".to_string(), e.to_string())
+            })?;
+
+        Ok(())
+    }
+}
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    const D: usize,
+    InnerParameters: GenericConfig<D, F = F>,
+    OuterParameters: GenericConfig<D, F = F>,
+>  WrapCircuit<F, D, InnerParameters, OuterParameters>
+{
+    pub fn new(
+        inner_verifier_data: VerifierCircuitData<F, InnerParameters, D>,
+    ) -> Self {
+        Self{
+            inner_verifier_data,
+            phantom_data: Default::default(),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct WrappedOutput<
+    F: RichField + Extendable<D> + Poseidon2,
+    C: GenericConfig<D, F = F>,
+    const D: usize
+> {
+    pub proof: ProofWithPublicInputs<F, C, D>,
+    pub common_data: CommonCircuitData<F, D>,
+    pub verifier_data: VerifierOnlyCircuitData<C, D>,
+}
+
+impl<
+    F: RichField + Extendable<D> + Poseidon2,
+    C: GenericConfig<D, F = F>,
+    const D: usize
+> WrappedOutput<F,C,D> {
+    pub fn save<P: AsRef<Path>>(&self, path: P) -> anyhow::Result<()>
+    where
+        C: Serialize,
+    {
+        if !path.as_ref().exists() {
+            fs::create_dir_all(&path)?;
+        }
+        let common_data_file = File::create(path.as_ref().join("common_circuit_data.json"))?;
+        serde_json::to_writer(&common_data_file, &self.common_data)?;
+        println!("Succesfully wrote common circuit data to common_circuit_data.json");
+
+        let verifier_data_file =
+            File::create(path.as_ref().join("verifier_only_circuit_data.json"))?;
+        serde_json::to_writer(&verifier_data_file, &self.verifier_data)?;
+        println!("Succesfully wrote verifier data to verifier_only_circuit_data.json");
+
+        let proof_file = File::create(path.as_ref().join("proof_with_public_inputs.json"))?;
+        serde_json::to_writer(&proof_file, &self.proof)?;
+        println!("Succesfully wrote proof to proof_with_public_inputs.json");
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use plonky2::field::goldilocks_field::GoldilocksField;
+    use plonky2::field::types::Field;
+    use plonky2::gates::noop::NoopGate;
+    use plonky2::plonk::circuit_builder::CircuitBuilder;
+    use plonky2::plonk::circuit_data::CircuitConfig;
+    use plonky2::plonk::config::PoseidonGoldilocksConfig;
+    use crate::bn254_wrapper::config::PoseidonBN254GoldilocksConfig;
+    use super::*;
+
+    #[test]
+    fn test_full_wrap() -> anyhow::Result<()>{
+        const D: usize = 2;
+
+        type F = GoldilocksField;
+        type InnerParameters = PoseidonGoldilocksConfig;
+        type OuterParameters = PoseidonBN254GoldilocksConfig;
+
+        let build_path = "./verifier_data".to_string();
+        let test_path = format!("{}/test_small/", build_path);
+
+        let conf = CircuitConfig::standard_recursion_config();
+        let mut builder =  CircuitBuilder::<F, D>::new(conf);
+
+        // let a = builder.add_virtual_public_input();
+        // let b = builder.add_virtual_public_input();
+        // let c = builder.add(a,b);
+        // builder.register_public_input(c);
+        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();
+
+        // Set up the dummy circuit and wrapper.
+        let dummy_circuit = builder.build::<InnerParameters>();
+        let mut pw = PartialWitness::new();
+        // pw.set_target(a, GoldilocksField::from_canonical_u64(1))?;
+        // pw.set_target(b, GoldilocksField::from_canonical_u64(2))?;
+        pw.set_target(t, F::ZERO).expect("faulty assign");
+        println!(
+            "dummy circuit degree: {}",
+            dummy_circuit.common.degree_bits()
+        );
+        let dummy_inner_proof = dummy_circuit.prove(pw).unwrap();
+        assert!(dummy_circuit.verify(dummy_inner_proof.clone()).is_ok());
+        println!("Verified dummy_circuit");
+
+        // wrap this in the outer circuit.
+        let wrapper = WrapCircuit::<F,D,InnerParameters,OuterParameters>::new(dummy_circuit.verifier_data());
+        let (targ, data) = wrapper.build_with_standard_config().unwrap();
+        println!(
+            "wrapper circuit degree: {}",
+            data.common.degree_bits()
+        );
+        let verifier_data = data.verifier_data();
+        let prover_data = data.prover_data();
+        let wrap_input = WrapInput{
+            inner_proof: dummy_inner_proof,
+        };
+        let proof = wrapper.prove(&targ, &wrap_input,&prover_data).unwrap();
+
+        let wrap_circ = WrappedOutput::<F, OuterParameters,D>{
+            proof,
+            common_data: prover_data.common,
+            verifier_data: verifier_data.verifier_only,
+        };
+
+        wrap_circ.save(test_path).unwrap();
+        println!("Saved test wrapped circuit");
+        Ok(())
+    }
+}

codex-plonky2-circuits/src/lib.rs

@@ -3,5 +3,6 @@ pub mod recursion;
 pub mod error;
 pub mod circuit_helper;
 mod bundle;
+pub mod bn254_wrapper;
 
 pub type Result<T> = core::result::Result<T, error::CircuitError>;

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

@@ -180,7 +180,7 @@ mod tests {
     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};
+    use plonky2_poseidon2::poseidon2_hash::poseidon2::Poseidon2Hash;
 
     // For our tests, we define:
     const D: usize = 2;

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

@@ -260,7 +260,7 @@ mod tests {
     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 plonky2_poseidon2::poseidon2_hash::poseidon2:: Poseidon2Hash;
     use crate::recursion::leaf::BUCKET_SIZE;
 
 
@@ -276,7 +276,7 @@ mod tests {
         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 {
+        for _i in 0..9+B {
             pub_input.push(builder.add_virtual_public_input());
         }
         let data = builder.build::<C>();

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

@@ -307,7 +307,7 @@ mod tests {
     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_poseidon2::poseidon2_hash::poseidon2::Poseidon2Hash;
     use plonky2::iop::witness::WitnessWrite;
 
     const D: usize = 2;

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

@@ -166,14 +166,13 @@ mod tests {
     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_poseidon2::poseidon2_hash::poseidon2::Poseidon2;
     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> {