@@ -101,8 +101,8 @@ }
     }
 
     /// Creates an new virtual polynomial from a MLE and its coefficient.
-    pub fn new_from_mle(mle: Rc<DenseMultilinearExtension<F>>, coefficient: F) -> Self {
-        let mle_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(&mle);
+    pub fn new_from_mle(mle: &Rc<DenseMultilinearExtension<F>>, coefficient: F) -> Self {
+        let mle_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(mle);
         let mut hm = HashMap::new();
         hm.insert(mle_ptr, 0);
 
@@ -115,7 +115,7 @@ phantom: PhantomData::default(),
             },
             // here `0` points to the first polynomial of `flattened_ml_extensions`
             products: vec![(coefficient, vec![0])],
-            flattened_ml_extensions: vec![mle],
+            flattened_ml_extensions: vec![mle.clone()],
             raw_pointers_lookup_table: hm,
         }
     }
@@ -475,7 +475,7 @@ VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
                 let (b, _b_sum) = random_mle_list(nv, 1, &mut rng);
                 let b_mle = b[0].clone();
                 let coeff = Fr::rand(&mut rng);
-                let b_vp = VirtualPolynomial::new_from_mle(b_mle.clone(), coeff);
+                let b_vp = VirtualPolynomial::new_from_mle(&b_mle, coeff);
 
                 let mut c = a.clone();
 

@@ -8,8 +8,7 @@ use ark_std::{end_timer, log2, start_timer, One, Zero};
 use errors::HyperPlonkErrors;
 use pcs::prelude::{compute_qx_degree, merge_polynomials, PCSErrors, PolynomialCommitmentScheme};
 use poly_iop::{
-    identity_permutation_mle,
-    prelude::{PermutationCheck, SumCheck, ZeroCheck},
+    prelude::{identity_permutation_mle, PermutationCheck, ZeroCheck},
     PolyIOP,
 };
 use selectors::SelectorColumn;
@@ -25,10 +24,9 @@ mod structs;
 mod utils;
 mod witness;
 
-/// A trait for HyperPlonk Poly-IOPs.
+/// A trait for HyperPlonk SNARKs.
 /// A HyperPlonk is derived from SumChecks, ZeroChecks and PermutationChecks.
-pub trait HyperPlonkSNARK<E, PCS>:
-    SumCheck<E::Fr> + ZeroCheck<E::Fr> + PermutationCheck<E::Fr>
+pub trait HyperPlonkSNARK<E, PCS>: PermutationCheck<E, PCS>
 where
     E: PairingEngine,
     PCS: PolynomialCommitmentScheme<E>,
@@ -99,7 +97,7 @@ {
     type Parameters = HyperPlonkParams;
     type ProvingKey = HyperPlonkProvingKey<E, PCS>;
     type VerifyingKey = HyperPlonkVerifyingKey<E, PCS>;
-    type Proof = HyperPlonkProof<E, PCS, Self, Self>;
+    type Proof = HyperPlonkProof<E, Self, PCS>;
 
     /// Generate the preprocessed polynomials output by the indexer.
     ///
@@ -277,7 +275,7 @@ "merged witness poly has a different num_var ({}) from expected ({})",
                 w_merged.num_vars, merged_nv
             )));
         }
-        let w_merged_com = PCS::commit(&pk.pcs_param, &Rc::new(w_merged.clone()))?;
+        let w_merged_com = PCS::commit(&pk.pcs_param, &w_merged)?;
 
         transcript.append_serializable_element(b"w", &w_merged_com)?;
         end_timer!(step);
@@ -309,51 +307,24 @@
         // =======================================================================
         // 3. Run permutation check on `\{w_i(x)\}` and `permutation_oracles`, and
         // obtain a PermCheckSubClaim.
-        //
-        // 3.1. `generate_challenge` from current transcript (generate beta, gamma)
-        // 3.2. `compute_product` to build `prod(x)` etc. from f, g and s_perm
-        // 3.3. push a commitment of `prod(x)` to the transcript
-        // 3.4. `update_challenge` with the updated transcript
-        // 3.5. `prove` to generate the proof
-        // 3.6. open `prod(0,x)`, `prod(1, x)`, `prod(x, 0)`, `prod(x, 1)` at
-        //      zero_check.point
         // =======================================================================
         let step = start_timer!(|| "Permutation check on w_i(x)");
 
-        // 3.1 `generate_challenge` from current transcript (generate beta, gamma)
-        let mut permutation_challenge = Self::generate_challenge(&mut transcript)?;
-
-        // 3.2. `compute_product` to build `prod(x)` etc. from f, g and s_perm
-
-        // This function returns 3 MLEs:
-        // - prod(x)
-        // - numerator
-        // - denominator
-        // See function signature for details.
-        let prod_x_and_aux_info = Self::compute_prod_evals(
-            &permutation_challenge,
+        let (perm_check_proof, prod_x) = <Self as PermutationCheck<E, PCS>>::prove(
+            &pk.pcs_param,
             &w_merged,
             &w_merged,
             &pk.permutation_oracles,
-        )?;
-        let prod_x = Rc::new(prod_x_and_aux_info[0].clone());
-
-        // 3.3 push a commitment of `prod(x)` to the transcript
-        let prod_com = PCS::commit(&pk.pcs_param, &prod_x)?;
-
-        // 3.4. `update_challenge` with the updated transcript
-        Self::update_challenge(&mut permutation_challenge, &mut transcript, &prod_com)?;
-
-        // 3.5. `prove` to generate the proof
-        let perm_check_proof = <Self as PermutationCheck<E::Fr>>::prove(
-            &prod_x_and_aux_info,
-            &permutation_challenge,
             &mut transcript,
         )?;
 
-        // 3.6 open prod(0,x), prod(1, x), prod(x, 0), prod(x, 1) at zero_check.point
+        // open prod(0,x), prod(1, x), prod(x, 0), prod(x, 1) at zero_check.point
         // prod(0, x)
-        let tmp_point = [perm_check_proof.point.as_slice(), &[E::Fr::zero()]].concat();
+        let tmp_point = [
+            perm_check_proof.zero_check_proof.point.as_slice(),
+            &[E::Fr::zero()],
+        ]
+        .concat();
         let (prod_0_x_opening, prod_0_x_eval) = PCS::open(&pk.pcs_param, &prod_x, &tmp_point)?;
         #[cfg(feature = "extensive_sanity_checks")]
         {
@@ -370,7 +341,11 @@ ));
             }
         }
         // prod(1, x)
-        let tmp_point = [perm_check_proof.point.as_slice(), &[E::Fr::one()]].concat();
+        let tmp_point = [
+            perm_check_proof.zero_check_proof.point.as_slice(),
+            &[E::Fr::one()],
+        ]
+        .concat();
         let (prod_1_x_opening, prod_1_x_eval) = PCS::open(&pk.pcs_param, &prod_x, &tmp_point)?;
         #[cfg(feature = "extensive_sanity_checks")]
         {
@@ -387,7 +362,11 @@ ));
             }
         }
         // prod(x, 0)
-        let tmp_point = [&[E::Fr::zero()], perm_check_proof.point.as_slice()].concat();
+        let tmp_point = [
+            &[E::Fr::zero()],
+            perm_check_proof.zero_check_proof.point.as_slice(),
+        ]
+        .concat();
         let (prod_x_0_opening, prod_x_0_eval) = PCS::open(&pk.pcs_param, &prod_x, &tmp_point)?;
         #[cfg(feature = "extensive_sanity_checks")]
         {
@@ -405,7 +384,11 @@ ));
             }
         }
         // prod(x, 1)
-        let tmp_point = [&[E::Fr::one()], perm_check_proof.point.as_slice()].concat();
+        let tmp_point = [
+            &[E::Fr::one()],
+            perm_check_proof.zero_check_proof.point.as_slice(),
+        ]
+        .concat();
         let (prod_x_1_opening, prod_x_1_eval) = PCS::open(&pk.pcs_param, &prod_x, &tmp_point)?;
         #[cfg(feature = "extensive_sanity_checks")]
         {
@@ -447,18 +430,20 @@
         // open permutation check proof
         let (witness_perm_check_opening, witness_perm_check_eval) = PCS::open(
             &pk.pcs_param,
-            &Rc::new(w_merged.clone()),
-            &perm_check_proof.point,
+            &w_merged,
+            &perm_check_proof.zero_check_proof.point,
         )?;
 
         #[cfg(feature = "extensive_sanity_checks")]
         {
             // sanity checks
-            let eval = w_merged.evaluate(&perm_check_proof.point).ok_or_else(|| {
-                HyperPlonkErrors::InvalidParameters(
-                    "evaluation dimension does not match".to_string(),
-                )
-            })?;
+            let eval = w_merged
+                .evaluate(&perm_check_proof.zero_check_proof.point)
+                .ok_or_else(|| {
+                    HyperPlonkErrors::InvalidParameters(
+                        "evaluation dimension does not match".to_string(),
+                    )
+                })?;
             if eval != witness_perm_check_eval {
                 return Err(HyperPlonkErrors::InvalidProver(
                     "Evaluation is different from PCS opening".to_string(),
@@ -492,7 +477,7 @@ // Open permutation polynomial at perm_check_point
         let (s_perm_opening, s_perm_eval) = PCS::open(
             &pk.pcs_param,
             &pk.permutation_oracles,
-            &perm_check_proof.point,
+            &perm_check_proof.zero_check_proof.point,
         )?;
 
         #[cfg(feature = "extensive_sanity_checks")]
@@ -500,7 +485,7 @@ {
             // sanity check
             let eval = pk
                 .permutation_oracles
-                .evaluate(&perm_check_proof.point)
+                .evaluate(&perm_check_proof.zero_check_proof.point)
                 .ok_or_else(|| {
                     HyperPlonkErrors::InvalidParameters(
                         "evaluation dimension does not match".to_string(),
@@ -576,7 +561,6 @@ // =======================================================================
             // PCS components: permutation check
             // =======================================================================
             // We do not validate prod(x), this is checked by subclaim
-            prod_commit: prod_com,
             prod_evals: vec![prod_0_x_eval, prod_1_x_eval, prod_x_0_eval, prod_x_1_eval],
             prod_openings: vec![
                 prod_0_x_opening,
@@ -715,6 +699,7 @@ // =======================================================================
         // 2. Verify perm_check_proof on `\{w_i(x)\}` and `permutation_oracles`
         // =======================================================================
         let step = start_timer!(|| "verify permutation check");
+
         // Zero check and sum check have different AuxInfo because `w_merged` and
         // `Prod(x)` have degree and num_vars
         let perm_check_aux_info = VPAuxInfo::<E::Fr> {
@@ -724,23 +709,21 @@ // degree of merged poly
             num_variables: merged_nv,
             phantom: PhantomData::default(),
         };
-        let mut challenge = <Self as PermutationCheck<E::Fr>>::generate_challenge(&mut transcript)?;
-        <Self as PermutationCheck<E::Fr>>::update_challenge(
-            &mut challenge,
-            &mut transcript,
-            &proof.prod_commit,
-        )?;
-
-        let perm_check_sub_claim = <Self as PermutationCheck<E::Fr>>::verify(
+        let perm_check_sub_claim = <Self as PermutationCheck<E, PCS>>::verify(
             &proof.perm_check_proof,
             &perm_check_aux_info,
             &mut transcript,
         )?;
+
         let perm_check_point = &perm_check_sub_claim
+            .product_check_sub_claim
             .zero_check_sub_claim
             .sum_check_sub_claim
             .point;
 
+        let alpha = perm_check_sub_claim.product_check_sub_claim.challenge;
+        let (beta, gamma) = perm_check_sub_claim.challenges;
+
         // check perm check subclaim:
         // proof.witness_perm_check_eval ?= perm_check_sub_claim.expected_eval
         //
@@ -754,9 +737,6 @@ // - prod(1, x) ... from prod(x) evaluated over (1, zero_point)
         // - g(x), f(x) are both w_merged over (zero_point)
         // - s_perm(x) and s_id(x) from vk_param.perm_oracle
         // - alpha, beta, gamma from challenge
-        let alpha = challenge
-            .alpha
-            .ok_or_else(|| HyperPlonkErrors::InvalidVerifier("alpha is not set".to_string()))?;
 
         let s_id = identity_permutation_mle::<E::Fr>(perm_check_point.len());
         let s_id_eval = s_id.evaluate(perm_check_point).ok_or_else(|| {
@@ -765,16 +745,13 @@ })?;
 
         let q_x_rec = proof.prod_evals[1] - proof.prod_evals[2] * proof.prod_evals[3]
             + alpha
-                * ((proof.witness_perm_check_eval
-                    + challenge.beta * proof.perm_oracle_eval
-                    + challenge.gamma)
+                * ((proof.witness_perm_check_eval + beta * proof.perm_oracle_eval + gamma)
                     * proof.prod_evals[0]
-                    - (proof.witness_perm_check_eval
-                        + challenge.beta * s_id_eval
-                        + challenge.gamma));
+                    - (proof.witness_perm_check_eval + beta * s_id_eval + gamma));
 
         if q_x_rec
             != perm_check_sub_claim
+                .product_check_sub_claim
                 .zero_check_sub_claim
                 .expected_evaluation
         {
@@ -823,7 +800,7 @@
         // prod(0, x)
         if !PCS::verify(
             &vk.pcs_param,
-            &proof.prod_commit,
+            &proof.perm_check_proof.prod_x_comm,
             &[perm_check_point.as_slice(), &[E::Fr::zero()]].concat(),
             &proof.prod_evals[0],
             &proof.prod_openings[0],
@@ -835,7 +812,7 @@ }
         // prod(1, x)
         if !PCS::verify(
             &vk.pcs_param,
-            &proof.prod_commit,
+            &proof.perm_check_proof.prod_x_comm,
             &[perm_check_point.as_slice(), &[E::Fr::one()]].concat(),
             &proof.prod_evals[1],
             &proof.prod_openings[1],
@@ -847,7 +824,7 @@ }
         // prod(x, 0)
         if !PCS::verify(
             &vk.pcs_param,
-            &proof.prod_commit,
+            &proof.perm_check_proof.prod_x_comm,
             &[&[E::Fr::zero()], perm_check_point.as_slice()].concat(),
             &proof.prod_evals[2],
             &proof.prod_openings[2],
@@ -859,7 +836,7 @@ }
         // prod(x, 1)
         if !PCS::verify(
             &vk.pcs_param,
-            &proof.prod_commit,
+            &proof.perm_check_proof.prod_x_comm,
             &[&[E::Fr::one()], perm_check_point.as_slice()].concat(),
             &proof.prod_evals[3],
             &proof.prod_openings[3],
@@ -940,7 +917,7 @@ use crate::{selectors::SelectorColumn, structs::CustomizedGates, witness::WitnessColumn};
     use ark_bls12_381::Bls12_381;
     use ark_std::test_rng;
     use pcs::prelude::KZGMultilinearPCS;
-    use poly_iop::random_permutation_mle;
+    use poly_iop::prelude::random_permutation_mle;
 
     #[test]
     fn test_hyperplonk_e2e() -> Result<(), HyperPlonkErrors> {

@@ -1,25 +1,10 @@
 //! Main module for the HyperPlonk PolyIOP.
 
 use ark_ec::PairingEngine;
-use ark_ff::PrimeField;
 use ark_poly::DenseMultilinearExtension;
 use pcs::PolynomialCommitmentScheme;
 use poly_iop::prelude::{PermutationCheck, ZeroCheck};
 use std::rc::Rc;
-
-/// The sub-claim for the HyperPlonk PolyIOP, consists of the following:
-///   - the SubClaim for the zero-check PIOP
-///   - the SubClaim for the permutation-check PIOP
-///   - the SubClaim for public input consistency
-#[derive(Clone, Debug, Default, PartialEq)]
-pub struct HyperPlonkSubClaim<F: PrimeField, ZC: ZeroCheck<F>, PC: PermutationCheck<F>> {
-    /// the SubClaim for the custom gate zerocheck
-    pub zero_check_sub_claim: ZC::ZeroCheckSubClaim,
-    /// the SubClaim for the permutation check
-    pub perm_check_sub_claim: PC::PermutationCheckSubClaim,
-    /// the public input consistency check
-    pub pub_input_sub_claim: (Vec<F>, F), // (point, expected_eval)
-}
 
 /// The proof for the HyperPlonk PolyIOP, consists of the following:
 ///   - a batch commitment to all the witness MLEs
@@ -27,12 +12,12 @@ ///   - a batch opening to all the MLEs at certain index
 ///   - the zero-check proof for checking custom gate-satisfiability
 ///   - the permutation-check proof for checking the copy constraints
 #[derive(Clone, Debug, Default, PartialEq)]
-pub struct HyperPlonkProof<
+pub struct HyperPlonkProof<E, PC, PCS>
+where
     E: PairingEngine,
+    PC: PermutationCheck<E, PCS>,
     PCS: PolynomialCommitmentScheme<E>,
-    ZC: ZeroCheck<E::Fr>,
-    PC: PermutationCheck<E::Fr>,
-> {
+{
     // =======================================================================
     // PCS components: common
     // =======================================================================
@@ -43,9 +28,6 @@ pub w_merged_com: PCS::Commitment,
     // =======================================================================
     // PCS components: permutation check
     // =======================================================================
-    /// PCS commit for prod(x)
-    // TODO: replace me with a batch commitment
-    pub prod_commit: PCS::Commitment,
     /// prod(x)'s evaluations
     /// sequence: prod(0,x), prod(1, x), prod(x, 0), prod(x, 1)
     pub prod_evals: Vec<E::Fr>,
@@ -86,7 +68,7 @@ // =======================================================================
     // IOP components
     // =======================================================================
     /// the custom gate zerocheck proof
-    pub zero_check_proof: ZC::ZeroCheckProof,
+    pub zero_check_proof: <PC as ZeroCheck<E::Fr>>::ZeroCheckProof,
     /// the permutation check proof for copy constraints
     pub perm_check_proof: PC::PermutationProof,
 }
@@ -97,7 +79,7 @@ ///   - binary log of the number of public input variables
 ///   - binary log of the number of selectors
 ///   - binary log of the number of witness wires
 ///   - the customized gate function
-#[derive(Clone, Debug, Default, PartialEq)]
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub struct HyperPlonkParams {
     /// the number of variables in polys
     pub nv: usize,

@@ -107,7 +107,7 @@ /// merge a set of polynomials. Returns an error if the
 /// polynomials do not share a same number of nvs.
 pub fn merge_polynomials<F: PrimeField>(
     polynomials: &[Rc<DenseMultilinearExtension<F>>],
-) -> Result<DenseMultilinearExtension<F>, PCSErrors> {
+) -> Result<Rc<DenseMultilinearExtension<F>>, PCSErrors> {
     let nv = polynomials[0].num_vars();
     for poly in polynomials.iter() {
         if nv != poly.num_vars() {
@@ -123,9 +123,9 @@ for poly in polynomials.iter() {
         scalars.extend_from_slice(poly.to_evaluations().as_slice());
     }
     scalars.extend_from_slice(vec![F::zero(); (1 << merged_nv) - scalars.len()].as_ref());
-    Ok(DenseMultilinearExtension::from_evaluations_vec(
+    Ok(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
         merged_nv, scalars,
-    ))
+    )))
 }
 
 /// Given a list of points, build `l(points)` which is a list of univariate
@@ -337,7 +337,7 @@ F::zero(),
                 F::from(1u64),
                 F::from(2u64),
             ];
-            let w_rec = DenseMultilinearExtension::from_evaluations_vec(3, w_eval);
+            let w_rec = Rc::new(DenseMultilinearExtension::from_evaluations_vec(3, w_eval));
 
             assert_eq!(w, w_rec);
         }
@@ -387,7 +387,7 @@ F::zero(),
                 F::zero(),
                 F::zero(),
             ];
-            let w_rec = DenseMultilinearExtension::from_evaluations_vec(4, w_eval);
+            let w_rec = Rc::new(DenseMultilinearExtension::from_evaluations_vec(4, w_eval));
 
             assert_eq!(w, w_rec);
         }

@@ -6,14 +6,13 @@
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-pcs = { path = "../pcs" }
 
 ark-ff = { version = "^0.3.0", default-features = false }
 ark-std = { version = "^0.3.0", default-features = false }
+ark-ec = { version = "^0.3.0", default-features = false }
 ark-poly = { version = "^0.3.0", default-features = false }
 ark-serialize =  { version = "^0.3.0", default-features = false }
 ark-bls12-381 = { version = "0.3.0", default-features = false, features = [ "curve" ] }
-ark-ec = { version = "^0.3.0", default-features = false }
 
 rand_chacha = { version = "0.3.0", default-features = false }
 displaydoc = { version = "0.2.3", default-features = false }
@@ -21,6 +20,7 @@ rayon = { version = "1.5.2", default-features = false, optional = true }
 
 transcript = { path = "../transcript" }
 arithmetic = { path = "../arithmetic" }
+pcs = { path = "../pcs" }
 
 [dev-dependencies]
 ark-ec = { version = "^0.3.0", default-features = false }
@@ -39,10 +39,11 @@ "rayon",
     "arithmetic/parallel",
     "ark-std/parallel", 
     "ark-ff/parallel",  
-    "ark-poly/parallel",
+    "ark-poly/parallel", 
     "pcs/parallel",
     ]
 print-trace = [ 
     "arithmetic/print-trace",
-    "ark-std/print-trace" 
+    "ark-std/print-trace",
+    "pcs/print-trace",
     ] 

@@ -1,11 +1,14 @@
 use arithmetic::{VPAuxInfo, VirtualPolynomial};
-use ark_bls12_381::Fr;
+use ark_bls12_381::{Bls12_381, Fr};
 use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
-use ark_std::{test_rng, UniformRand};
+use ark_std::test_rng;
+use pcs::{prelude::KZGMultilinearPCS, PolynomialCommitmentScheme};
 use poly_iop::prelude::{
     identity_permutation_mle, PermutationCheck, PolyIOP, PolyIOPErrors, SumCheck, ZeroCheck,
 };
-use std::{marker::PhantomData, time::Instant};
+use std::{marker::PhantomData, rc::Rc, time::Instant};
+
+type KZG = KZGMultilinearPCS<Bls12_381>;
 
 fn main() -> Result<(), PolyIOPErrors> {
     bench_permutation_check()?;
@@ -135,6 +138,9 @@ fn bench_permutation_check() -> Result<(), PolyIOPErrors> {
     let mut rng = test_rng();
 
     for nv in 4..20 {
+        let srs = KZG::gen_srs_for_testing(&mut rng, nv + 1)?;
+        let (pcs_param, _) = KZG::trim(&srs, nv + 1, Some(nv + 1))?;
+
         let repetition = if nv < 10 {
             100
         } else if nv < 20 {
@@ -143,34 +149,22 @@ } else {
             10
         };
 
-        let w = DenseMultilinearExtension::rand(nv, &mut rng);
+        let w = Rc::new(DenseMultilinearExtension::rand(nv, &mut rng));
 
         // s_perm is the identity map
         let s_perm = identity_permutation_mle(nv);
 
         let proof = {
             let start = Instant::now();
-            let mut transcript = <PolyIOP<Fr> as PermutationCheck<Fr>>::init_transcript();
+            let mut transcript =
+                <PolyIOP<Fr> as PermutationCheck<Bls12_381, KZG>>::init_transcript();
             transcript.append_message(b"testing", b"initializing transcript for testing")?;
 
-            let mut challenge =
-                <PolyIOP<Fr> as PermutationCheck<Fr>>::generate_challenge(&mut transcript)?;
-
-            let prod_x_and_aux = <PolyIOP<Fr> as PermutationCheck<Fr>>::compute_prod_evals(
-                &challenge, &w, &w, &s_perm,
-            )?;
-
-            let prod_x_binding = mock_commit(&prod_x_and_aux[0]);
-
-            <PolyIOP<Fr> as PermutationCheck<Fr>>::update_challenge(
-                &mut challenge,
-                &mut transcript,
-                &prod_x_binding,
-            )?;
-
-            let proof = <PolyIOP<Fr> as PermutationCheck<Fr>>::prove(
-                &prod_x_and_aux,
-                &challenge,
+            let (proof, _q_x) = <PolyIOP<Fr> as PermutationCheck<Bls12_381, KZG>>::prove(
+                &pcs_param,
+                &w,
+                &w,
+                &s_perm,
                 &mut transcript,
             )?;
 
@@ -190,10 +184,14 @@ phantom: PhantomData::default(),
             };
 
             let start = Instant::now();
-            let mut transcript = <PolyIOP<Fr> as PermutationCheck<Fr>>::init_transcript();
+            let mut transcript =
+                <PolyIOP<Fr> as PermutationCheck<Bls12_381, KZG>>::init_transcript();
             transcript.append_message(b"testing", b"initializing transcript for testing")?;
-            let _subclaim =
-                <PolyIOP<Fr> as PermutationCheck<Fr>>::verify(&proof, &poly_info, &mut transcript)?;
+            let _perm_check_sum_claim = <PolyIOP<Fr> as PermutationCheck<Bls12_381, KZG>>::verify(
+                &proof,
+                &poly_info,
+                &mut transcript,
+            )?;
             println!(
                 "permutation check verification time for {} variables: {} ns",
                 nv,
@@ -206,8 +204,3 @@ }
 
     Ok(())
 }
-
-fn mock_commit(_f: &DenseMultilinearExtension<Fr>) -> Fr {
-    let mut rng = test_rng();
-    Fr::rand(&mut rng)
-}

@@ -10,17 +10,7 @@ mod sum_check;
 mod utils;
 mod zero_check;
 
-pub use errors::PolyIOPErrors;
-pub use perm_check::{
-    util::{identity_permutation_mle, random_permutation_mle},
-    PermutationCheck,
-};
-pub use prod_check::ProductCheck;
-pub use sum_check::SumCheck;
-pub use utils::*;
-pub use zero_check::ZeroCheck;
-
-#[derive(Clone, Debug, Default, Copy)]
+#[derive(Clone, Debug, Default, Copy, PartialEq, Eq)]
 /// Struct for PolyIOP protocol.
 /// It has an associated type `F` that defines the prime field the multi-variate
 /// polynomial operates on.

@@ -1,20 +1,30 @@
 //! Main module for the Permutation Check protocol
 
-use crate::{
-    errors::PolyIOPErrors,
-    perm_check::util::{build_prod_partial_eval, compute_prod_0},
-    structs::IOPProof,
-    utils::get_index,
-    PolyIOP, ZeroCheck,
-};
-use arithmetic::VirtualPolynomial;
-use ark_ff::PrimeField;
+use self::util::computer_num_and_denom;
+use crate::{errors::PolyIOPErrors, prelude::ProductCheck, PolyIOP};
+use ark_ec::PairingEngine;
 use ark_poly::DenseMultilinearExtension;
-use ark_serialize::CanonicalSerialize;
 use ark_std::{end_timer, start_timer};
+use pcs::PolynomialCommitmentScheme;
 use std::rc::Rc;
 use transcript::IOPTranscript;
 
+/// A permutation subclaim consists of
+/// - the SubClaim from the ProductCheck
+/// - Challenges beta and gamma
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct PermutationCheckSubClaim<E, PCS, PC>
+where
+    E: PairingEngine,
+    PC: ProductCheck<E, PCS>,
+    PCS: PolynomialCommitmentScheme<E>,
+{
+    /// the SubClaim from the ProductCheck
+    pub product_check_sub_claim: PC::ProductCheckSubClaim,
+    /// Challenges beta and gamma
+    pub challenges: (E::Fr, E::Fr),
+}
+
 pub mod util;
 
 /// A PermutationCheck is derived from ZeroCheck.
@@ -25,31 +35,14 @@ /// Inputs:
 /// - f(x)
 /// - g(x)
 /// - permutation s_perm(x)
-///
-/// Steps:
-/// 1. `generate_challenge` from current transcript (generate beta, gamma)
-/// 2. `compute_product` to build `prod(x)` etc. from f, g and s_perm
-/// 3. push a commitment of `prod(x)` to the transcript (done by the snark
-/// caller)
-/// 4. `update_challenge` with the updated transcript (generate alpha)
-/// 5. `prove` to generate the proof
-pub trait PermutationCheck<F: PrimeField>: ZeroCheck<F> {
+pub trait PermutationCheck<E, PCS>: ProductCheck<E, PCS>
+where
+    E: PairingEngine,
+    PCS: PolynomialCommitmentScheme<E>,
+{
     type PermutationCheckSubClaim;
-    type PermutationChallenge;
     type PermutationProof;
 
-    /// Generate the preprocessed polynomial for the permutation check.
-    ///
-    /// The algorithm takes as input a permutation and outputs a merged
-    /// multilinear polynomial s(X0, X1, ..., Xn) such that
-    /// - s(0, X1, ..., Xn) = s_id(X1, ..., Xn) (identity permutation
-    ///   polynomial)
-    /// - s(1, X1, ..., Xn) = s_perm(X1, ..., Xn) (permutation polynomial)
-    fn preprocess(
-        permutation: &[F],
-        aux_info: &Self::VPAuxInfo,
-    ) -> Result<DenseMultilinearExtension<F>, PolyIOPErrors>;
-
     /// Initialize the system with a transcript
     ///
     /// This function is optional -- in the case where a PermutationCheck is
@@ -58,77 +51,23 @@ /// may be initialized by this complex protocol, and passed to the
     /// PermutationCheck prover/verifier.
     fn init_transcript() -> Self::Transcript;
 
-    /// Step 1 of the IOP.
-    /// Generate challenge beta and gamma from a transcript.
-    fn generate_challenge(
-        transcript: &mut Self::Transcript,
-    ) -> Result<Self::PermutationChallenge, PolyIOPErrors>;
-
-    /// Step 2 of the IOP.
-    ///
-    /// Input:
-    /// - f(x), g(x), s_perm(x) are mle-s
-    /// - challenges, consists of beta and gamma
-    ///
-    /// Output: the evaluations for the following 3 polynomials
-    /// - prod(x)
-    /// - numerator
-    /// - denominator
-    ///
-    /// where
-    /// - `prod(0,x) := prod(0, x0, x1, …, xn)` which is the MLE over the
-    /// evaluations of the following polynomial on the boolean hypercube
-    /// {0,1}^n:
-    ///
-    ///  (f(x) + \beta s_id(x) + \gamma)/(g(x) + \beta s_perm(x) + \gamma)
-    ///
-    ///   where s_id(x) is an identity permutation
-    ///
-    /// - numerator is the MLE for `f(x) + \beta s_id(x) + \gamma`
-    /// - denominator is the MLE for `g(x) + \beta s_perm(x) + \gamma`
-    ///
-    /// The caller needs to check num_vars matches in f/g/s_id/s_perm
-    /// Cost: linear in N.
-    ///
-    /// TODO: replace argument `s_perm` with the merged polynomial `s`.
-    fn compute_prod_evals(
-        challenge: &Self::PermutationChallenge,
-        fx: &DenseMultilinearExtension<F>,
-        gx: &DenseMultilinearExtension<F>,
-        s_perm: &DenseMultilinearExtension<F>,
-    ) -> Result<[DenseMultilinearExtension<F>; 3], PolyIOPErrors>;
-
-    /// Step 3 of the IOP.
-    /// push a commitment of `prod(x)` to the transcript
-    /// IMPORTANT: this step is done by the snark caller
-    fn commit_prod_x() {
-        unimplemented!()
-    }
-
-    /// Step 4 of the IOP.
-    /// Update the challenge with alpha; returns an error if
-    /// alpha already exists.
-    fn update_challenge(
-        challenge: &mut Self::PermutationChallenge,
-        transcript: &mut Self::Transcript,
-        prod_x_binding: &impl CanonicalSerialize,
-    ) -> Result<(), PolyIOPErrors>;
-
-    /// Step 5 of the IOP.
+    /// Inputs:
+    /// - f(x)
+    /// - g(x)
+    /// - permutation s_perm(x)
+    /// Outputs:
+    /// - a permutation check proof proving that g is a permutation of f under
+    ///   s_perm
+    /// - the Q(x) polynomial build during product check
     ///
-    /// Initialize the prover to argue that an MLE g(x) is a permutation of
-    /// MLE f(x) over a permutation given by s_perm.
-    /// Inputs:
-    /// - 3 MLEs from the second step
-    /// - challenge: `Self::Challenge` that has been updated
-    /// - transcript: a transcript that is used to generate the challenges beta
-    ///   and gamma
     /// Cost: O(N)
     fn prove(
-        prod_x_and_aux_info: &[DenseMultilinearExtension<F>; 3],
-        challenge: &Self::PermutationChallenge,
-        transcript: &mut IOPTranscript<F>,
-    ) -> Result<Self::PermutationProof, PolyIOPErrors>;
+        pcs_param: &PCS::ProverParam,
+        fx: &Self::MultilinearExtension,
+        gx: &Self::MultilinearExtension,
+        s_perm: &Self::MultilinearExtension,
+        transcript: &mut IOPTranscript<E::Fr>,
+    ) -> Result<(Self::PermutationProof, Self::MultilinearExtension), PolyIOPErrors>;
 
     /// Verify that an MLE g(x) is a permutation of
     /// MLE f(x) over a permutation given by s_perm.
@@ -139,32 +78,6 @@ transcript: &mut Self::Transcript,
     ) -> Result<Self::PermutationCheckSubClaim, PolyIOPErrors>;
 }
 
-/// A permutation subclaim consists of
-/// - A zero check IOP subclaim for Q(x) is 0, consists of the following:
-///  (See `build_qx` for definition of Q(x).)
-///   - the SubClaim from the SumCheck
-///   - the initial challenge r which is used to build eq(x, r) in ZeroCheck
-/// - A final query for `prod(1, ..., 1, 0) = 1`.
-// Note that this final query is in fact a constant that
-// is independent from the proof. So we should avoid
-// (de)serialize it.
-#[derive(Clone, Debug, Default, PartialEq)]
-pub struct PermutationCheckSubClaim<F: PrimeField, ZC: ZeroCheck<F>> {
-    // the SubClaim from the ZeroCheck
-    pub zero_check_sub_claim: ZC::ZeroCheckSubClaim,
-    // final query which consists of
-    // - the vector `(1, ..., 1, 0)`
-    // - the evaluation `1`
-    final_query: (Vec<F>, F),
-}
-
-#[derive(Debug, Clone)]
-pub struct PermutationChallenge<F: PrimeField> {
-    pub alpha: Option<F>,
-    pub beta: F,
-    pub gamma: F,
-}
-
 /// A PermutationCheck is derived from ZeroCheck.
 ///
 /// A Permutation Check IOP takes the following steps:
@@ -181,28 +94,16 @@ /// 3. push a commitment of `prod(x)` to the transcript (done by the snark
 /// caller)
 /// 4. `update_challenge` with the updated transcript (generate alpha)
 /// 5. `prove` to generate the proof
-impl<F: PrimeField> PermutationCheck<F> for PolyIOP<F> {
+impl<E, PCS> PermutationCheck<E, PCS> for PolyIOP<E::Fr>
+where
+    E: PairingEngine,
+    PCS: PolynomialCommitmentScheme<E, Polynomial = Rc<DenseMultilinearExtension<E::Fr>>>,
+{
     /// A Permutation SubClaim is indeed a ZeroCheck SubClaim that consists of
     /// - the SubClaim from the SumCheck
     /// - the initial challenge r which is used to build eq(x, r)
-    type PermutationCheckSubClaim = PermutationCheckSubClaim<F, Self>;
-
-    type PermutationProof = Self::SumCheckProof;
-    type PermutationChallenge = PermutationChallenge<F>;
-
-    /// Generate the preprocessed polynomial for the permutation check.
-    ///
-    /// The algorithm takes as input a permutation and outputs a merged
-    /// multilinear polynomial s(X0, X1, ..., Xn) such that
-    /// - s(0, X1, ..., Xn) = s_id(X1, ..., Xn) (identity permutation
-    ///   polynomial)
-    /// - s(1, X1, ..., Xn) = s_perm(X1, ..., Xn) (permutation polynomial)
-    fn preprocess(
-        _permutation: &[F],
-        _aux_info: &Self::VPAuxInfo,
-    ) -> Result<DenseMultilinearExtension<F>, PolyIOPErrors> {
-        unimplemented!();
-    }
+    type PermutationCheckSubClaim = PermutationCheckSubClaim<E, PCS, Self>;
+    type PermutationProof = Self::ProductCheckProof;
 
     /// Initialize the system with a transcript
     ///
@@ -211,162 +112,50 @@ /// an building block for a more complex protocol, the transcript
     /// may be initialized by this complex protocol, and passed to the
     /// PermutationCheck prover/verifier.
     fn init_transcript() -> Self::Transcript {
-        IOPTranscript::<F>::new(b"Initializing PermutationCheck transcript")
+        IOPTranscript::<E::Fr>::new(b"Initializing PermutationCheck transcript")
     }
 
-    /// Step 1 of the IOP.
-    /// Generate challenge beta and gamma from a transcript.
-    fn generate_challenge(
-        transcript: &mut Self::Transcript,
-    ) -> Result<Self::PermutationChallenge, PolyIOPErrors> {
-        Ok(Self::PermutationChallenge {
-            beta: transcript.get_and_append_challenge(b"beta")?,
-            gamma: transcript.get_and_append_challenge(b"gamma")?,
-            alpha: None,
-        })
-    }
-
-    /// Step 2 of the IOP.
+    /// Inputs:
+    /// - f(x)
+    /// - g(x)
+    /// - permutation s_perm(x)
+    /// Outputs:
+    /// - a permutation check proof proving that g is a permutation of f under
+    ///   s_perm
+    /// - the Q(x) polynomial build during product check
     ///
-    /// Input:
-    /// - f(x), g(x), s_perm(x) are mle-s
-    /// - challenges, consists of beta and gamma
-    ///
-    /// Output: the evaluations for the following 3 polynomials
-    /// - prod(x)
-    /// - numerator
-    /// - denominator
-    ///
-    /// where
-    /// - `prod(0,x) := prod(0, x0, x1, …, xn)` which is the MLE over the
-    /// evaluations of the following polynomial on the boolean hypercube
-    /// {0,1}^n:
-    ///
-    ///  (f(x) + \beta s_id(x) + \gamma)/(g(x) + \beta s_perm(x) + \gamma)
-    ///
-    ///   where s_id(x) is an identity permutation
-    ///
-    /// - numerator is the MLE for `f(x) + \beta s_id(x) + \gamma`
-    /// - denominator is the MLE for `g(x) + \beta s_perm(x) + \gamma`
-    ///
-    /// The caller needs to check num_vars matches in f/g/s_id/s_perm
-    /// Cost: linear in N.
-    ///
-    /// TODO: replace argument `s_perm` with the merged polynomial `s`.
-    fn compute_prod_evals(
-        challenge: &Self::PermutationChallenge,
-        fx: &DenseMultilinearExtension<F>,
-        gx: &DenseMultilinearExtension<F>,
-        s_perm: &DenseMultilinearExtension<F>,
-    ) -> Result<[DenseMultilinearExtension<F>; 3], PolyIOPErrors> {
-        let start = start_timer!(|| "compute evaluations of prod polynomial");
-
-        if challenge.alpha.is_some() {
-            return Err(PolyIOPErrors::InvalidChallenge(
-                "alpha is already sampled".to_string(),
+    /// Cost: O(N)
+    fn prove(
+        pcs_param: &PCS::ProverParam,
+        fx: &Self::MultilinearExtension,
+        gx: &Self::MultilinearExtension,
+        s_perm: &Self::MultilinearExtension,
+        transcript: &mut IOPTranscript<E::Fr>,
+    ) -> Result<(Self::PermutationProof, Self::MultilinearExtension), PolyIOPErrors> {
+        let start = start_timer!(|| "Permutation check prove");
+        if fx.num_vars != gx.num_vars {
+            return Err(PolyIOPErrors::InvalidParameters(
+                "fx and gx have different number of variables".to_string(),
             ));
         }
 
-        let num_vars = fx.num_vars;
-
-        // ===================================
-        // prod(0, x)
-        // ===================================
-        let (prod_0x_eval, numerator_eval, denominator_eval) =
-            compute_prod_0(&challenge.beta, &challenge.gamma, fx, gx, s_perm)?;
-
-        // ===================================
-        // prod(1, x)
-        // ===================================
-        //
-        // `prod(1, x)` can be computed via recursing the following formula for 2^n-1
-        // times
-        //
-        // `prod(1, x_1, ..., x_n) :=
-        //      prod(x_1, x_2, ..., x_n, 0) * prod(x_1, x_2, ..., x_n, 1)`
-        //
-        // At any given step, the right hand side of the equation
-        // is available via either eval_0x or the current view of eval_1x
-        let mut prod_1x_eval = vec![];
-        for x in 0..(1 << num_vars) - 1 {
-            // sign will decide if the evaluation should be looked up from eval_0x or
-            // eval_1x; x_zero_index is the index for the evaluation (x_2, ..., x_n,
-            // 0); x_one_index is the index for the evaluation (x_2, ..., x_n, 1);
-            let (x_zero_index, x_one_index, sign) = get_index(x, num_vars);
-            if !sign {
-                prod_1x_eval.push(prod_0x_eval[x_zero_index] * prod_0x_eval[x_one_index]);
-            } else {
-                // sanity check: if we are trying to look up from the eval_1x table,
-                // then the target index must already exist
-                if x_zero_index >= prod_1x_eval.len() || x_one_index >= prod_1x_eval.len() {
-                    return Err(PolyIOPErrors::ShouldNotArrive);
-                }
-                prod_1x_eval.push(prod_1x_eval[x_zero_index] * prod_1x_eval[x_one_index]);
-            }
+        if fx.num_vars != s_perm.num_vars {
+            return Err(PolyIOPErrors::InvalidParameters(
+                "fx and s_perm have different number of variables".to_string(),
+            ));
         }
 
-        // prod(1, 1, ..., 1) := 0
-        prod_1x_eval.push(F::zero());
-
-        // ===================================
-        // prod(x)
-        // ===================================
-        // prod(x)'s evaluation is indeed `e := [eval_0x[..], eval_1x[..]].concat()`
-        let eval = [prod_0x_eval.as_slice(), prod_1x_eval.as_slice()].concat();
+        // generate challenge `beta` and `gamma` from current transcript
+        let beta = transcript.get_and_append_challenge(b"beta")?;
+        let gamma = transcript.get_and_append_challenge(b"gamma")?;
+        let (numerator, denominator) = computer_num_and_denom(&beta, &gamma, fx, gx, s_perm)?;
 
-        let fx = DenseMultilinearExtension::from_evaluations_vec(num_vars + 1, eval);
-        let numerator = DenseMultilinearExtension::from_evaluations_vec(num_vars, numerator_eval);
-        let denominator =
-            DenseMultilinearExtension::from_evaluations_vec(num_vars, denominator_eval);
+        // invoke product check on numerator and denominator
+        let (proof, poly) =
+            <Self as ProductCheck<E, PCS>>::prove(pcs_param, &numerator, &denominator, transcript)?;
 
         end_timer!(start);
-        Ok([fx, numerator, denominator])
-    }
-
-    /// Step 4 of the IOP.
-    /// Update the challenge with alpha; returns an error if
-    /// alpha already exists.
-    fn update_challenge(
-        challenge: &mut Self::PermutationChallenge,
-        transcript: &mut Self::Transcript,
-        prod_x_binding: &impl CanonicalSerialize,
-    ) -> Result<(), PolyIOPErrors> {
-        if challenge.alpha.is_some() {
-            return Err(PolyIOPErrors::InvalidChallenge(
-                "alpha should not be sampled at the current stage".to_string(),
-            ));
-        }
-        transcript.append_serializable_element(b"prod(x)", prod_x_binding)?;
-        challenge.alpha = Some(transcript.get_and_append_challenge(b"alpha")?);
-        Ok(())
-    }
-
-    /// Step 5 of the IOP.
-    ///
-    /// Initialize the prover to argue that an MLE g(x) is a permutation of
-    /// MLE f(x) over a permutation given by s_perm.
-    /// Inputs:
-    /// - 3 MLEs from the second step
-    /// - challenge: `Self::Challenge` that has been updated
-    /// - transcript: a transcript that is used to generate the challenges beta
-    ///   and gamma
-    /// Cost: O(N)
-    fn prove(
-        prod_x_and_aux_info: &[DenseMultilinearExtension<F>; 3],
-        challenge: &Self::PermutationChallenge,
-        transcript: &mut IOPTranscript<F>,
-    ) -> Result<Self::PermutationProof, PolyIOPErrors> {
-        let alpha = match challenge.alpha {
-            Some(p) => p,
-            None => {
-                return Err(PolyIOPErrors::InvalidChallenge(
-                    "alpha is not sampled yet".to_string(),
-                ))
-            },
-        };
-
-        let (proof, _q_x) = prove_internal(prod_x_and_aux_info, &alpha, transcript)?;
-        Ok(proof)
+        Ok((proof, poly))
     }
 
     /// Verify that an MLE g(x) is a permutation of an
@@ -378,235 +167,187 @@ transcript: &mut Self::Transcript,
     ) -> Result<Self::PermutationCheckSubClaim, PolyIOPErrors> {
         let start = start_timer!(|| "Permutation check verify");
 
+        let beta = transcript.get_and_append_challenge(b"beta")?;
+        let gamma = transcript.get_and_append_challenge(b"gamma")?;
+
         // invoke the zero check on the iop_proof
-        let zero_check_sub_claim = <Self as ZeroCheck<F>>::verify(proof, aux_info, transcript)?;
-        let mut final_query = vec![F::one(); aux_info.num_variables];
-        final_query[aux_info.num_variables - 1] = F::zero();
-        let final_eval = F::one();
+        let product_check_sub_claim =
+            <Self as ProductCheck<E, PCS>>::verify(proof, aux_info, transcript)?;
 
         end_timer!(start);
-
         Ok(PermutationCheckSubClaim {
-            zero_check_sub_claim,
-            final_query: (final_query, final_eval),
+            product_check_sub_claim,
+            challenges: (beta, gamma),
         })
     }
 }
 
-/// Step 5 of the IOP.
-///
-/// Generate a proof to argue that an MLE g(x) is a permutation of
-/// MLE f(x) over a permutation given by s_perm.
-/// Inputs:
-/// - 3 MLEs from the second step
-/// - challenge alpha
-/// - transcript: a transcript that is used to generate the challenges beta and
-///   gamma
-///
-/// Returns proof and Q(x) for testing purpose.
-///
-/// Cost: O(N)
-fn prove_internal<F: PrimeField>(
-    prod_x_and_aux_info: &[DenseMultilinearExtension<F>; 3],
-    alpha: &F,
-    transcript: &mut IOPTranscript<F>,
-) -> Result<(IOPProof<F>, VirtualPolynomial<F>), PolyIOPErrors> {
-    let start = start_timer!(|| "Permutation check prove");
-
-    let prod_partial_evals = build_prod_partial_eval(&prod_x_and_aux_info[0])?;
-    let prod_0x = Rc::new(prod_partial_evals[0].clone());
-    let prod_1x = Rc::new(prod_partial_evals[1].clone());
-    let prod_x0 = Rc::new(prod_partial_evals[2].clone());
-    let prod_x1 = Rc::new(prod_partial_evals[3].clone());
-    let numerator = Rc::new(prod_x_and_aux_info[1].clone());
-    let denominator = Rc::new(prod_x_and_aux_info[2].clone());
-
-    // compute (g(x) + beta * s_perm(x) + gamma) * prod(0, x) * alpha
-    // which is prods[6] * prod[1] * alpha
-    let mut q_x = VirtualPolynomial::new_from_mle(denominator, F::one());
-    q_x.mul_by_mle(prod_0x, *alpha)?;
-
-    //   (g(x) + beta * s_perm(x) + gamma) * prod(0, x) * alpha
-    // - (f(x) + beta * s_id(x)   + gamma) * alpha
-    q_x.add_mle_list([numerator], -*alpha)?;
-
-    // Q(x) := prod(1,x) - prod(x, 0) * prod(x, 1)
-    //       + alpha * (
-    //             (g(x) + beta * s_perm(x) + gamma) * prod(0, x)
-    //           - (f(x) + beta * s_id(x)   + gamma))
-    q_x.add_mle_list([prod_x0, prod_x1], -F::one())?;
-    q_x.add_mle_list([prod_1x], F::one())?;
-
-    let iop_proof = <PolyIOP<F> as ZeroCheck<F>>::prove(&q_x, transcript)?;
-
-    end_timer!(start);
-    Ok((iop_proof, q_x))
-}
-
 #[cfg(test)]
 mod test {
     use super::{util::build_prod_partial_eval, PermutationCheck};
     use crate::{
         errors::PolyIOPErrors,
-        perm_check::prove_internal,
+        perm_check::util::computer_num_and_denom,
         prelude::{identity_permutation_mle, random_permutation_mle},
-        structs::IOPProof,
         utils::bit_decompose,
         PolyIOP,
     };
     use arithmetic::{VPAuxInfo, VirtualPolynomial};
-    use ark_bls12_381::{Fr, G1Affine};
-    use ark_ec::{AffineCurve, ProjectiveCurve};
-    use ark_ff::{UniformRand, Zero};
+    use ark_bls12_381::Bls12_381;
+    use ark_ec::PairingEngine;
+    use ark_ff::{One, Zero};
     use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
     use ark_std::test_rng;
-    use std::marker::PhantomData;
+    use pcs::{prelude::KZGMultilinearPCS, PolynomialCommitmentScheme};
+    use std::{marker::PhantomData, rc::Rc};
 
-    /// This is a mock function to generate some commitment element for testing.
-    fn mock_commit<G: AffineCurve>(_f: &DenseMultilinearExtension<G::ScalarField>) -> G {
-        let mut rng = test_rng();
-        G::Projective::rand(&mut rng).into_affine()
-    }
+    type KZG = KZGMultilinearPCS<Bls12_381>;
 
-    fn test_permutation_check_helper(
-        f: &DenseMultilinearExtension<Fr>,
-        g: &DenseMultilinearExtension<Fr>,
-        s_perm: &DenseMultilinearExtension<Fr>,
-    ) -> Result<(IOPProof<Fr>, VirtualPolynomial<Fr>), PolyIOPErrors> {
-        let mut transcript = <PolyIOP<Fr> as PermutationCheck<Fr>>::init_transcript();
+    fn test_permutation_check_helper<E, PCS>(
+        pcs_param: &PCS::ProverParam,
+        fx: &Rc<DenseMultilinearExtension<E::Fr>>,
+        gx: &Rc<DenseMultilinearExtension<E::Fr>>,
+        s_perm: &Rc<DenseMultilinearExtension<E::Fr>>,
+    ) -> Result<(), PolyIOPErrors>
+    where
+        E: PairingEngine,
+        PCS: PolynomialCommitmentScheme<E, Polynomial = Rc<DenseMultilinearExtension<E::Fr>>>,
+    {
+        let nv = fx.num_vars;
+        let poly_info = VPAuxInfo {
+            max_degree: 2,
+            num_variables: nv,
+            phantom: PhantomData::default(),
+        };
+
+        // prover
+        let mut transcript = <PolyIOP<E::Fr> as PermutationCheck<E, PCS>>::init_transcript();
         transcript.append_message(b"testing", b"initializing transcript for testing")?;
+        let (proof, q_x) = <PolyIOP<E::Fr> as PermutationCheck<E, PCS>>::prove(
+            pcs_param,
+            fx,
+            gx,
+            s_perm,
+            &mut transcript,
+        )?;
 
-        let mut challenge =
-            <PolyIOP<Fr> as PermutationCheck<Fr>>::generate_challenge(&mut transcript)?;
+        // verifier
+        let mut transcript = <PolyIOP<E::Fr> as PermutationCheck<E, PCS>>::init_transcript();
+        transcript.append_message(b"testing", b"initializing transcript for testing")?;
+        let perm_check_sum_claim = <PolyIOP<E::Fr> as PermutationCheck<E, PCS>>::verify(
+            &proof,
+            &poly_info,
+            &mut transcript,
+        )?;
 
-        let prod_x_and_aux =
-            <PolyIOP<Fr> as PermutationCheck<Fr>>::compute_prod_evals(&challenge, f, g, s_perm)?;
+        let prod_partial_evals = build_prod_partial_eval(&q_x)?;
+        let prod_0x = prod_partial_evals[0].clone();
+        let prod_1x = prod_partial_evals[1].clone();
+        let prod_x0 = prod_partial_evals[2].clone();
+        let prod_x1 = prod_partial_evals[3].clone();
 
-        let prod_x_binding: G1Affine = mock_commit(&prod_x_and_aux[0]);
+        let (numerator, denominator) = computer_num_and_denom(
+            &perm_check_sum_claim.challenges.0,
+            &perm_check_sum_claim.challenges.1,
+            fx,
+            gx,
+            &s_perm,
+        )?;
 
-        <PolyIOP<Fr> as PermutationCheck<Fr>>::update_challenge(
-            &mut challenge,
-            &mut transcript,
-            &prod_x_binding,
+        // compute (g(x) + beta * s_perm(x) + gamma) * prod(0, x) * alpha
+        // which is prods[6] * prod[1] * alpha
+        let mut q_x = VirtualPolynomial::new_from_mle(&denominator, E::Fr::one());
+        q_x.mul_by_mle(
+            prod_0x,
+            perm_check_sum_claim.product_check_sub_claim.challenge,
         )?;
-        let alpha = challenge.alpha.unwrap();
 
-        prove_internal(&prod_x_and_aux, &alpha, &mut transcript)
+        //   (g(x) + beta * s_perm(x) + gamma) * prod(0, x) * alpha
+        // - (f(x) + beta * s_id(x)   + gamma) * alpha
+        q_x.add_mle_list(
+            [numerator],
+            -perm_check_sum_claim.product_check_sub_claim.challenge,
+        )?;
+
+        // Q(x) := prod(1,x) - prod(x, 0) * prod(x, 1)
+        //       + alpha * (
+        //             (g(x) + beta * s_perm(x) + gamma) * prod(0, x)
+        //           - (f(x) + beta * s_id(x)   + gamma))
+        q_x.add_mle_list([prod_x0, prod_x1], -E::Fr::one())?;
+        q_x.add_mle_list([prod_1x], E::Fr::one())?;
+
+        if q_x
+            .evaluate(
+                &perm_check_sum_claim
+                    .product_check_sub_claim
+                    .zero_check_sub_claim
+                    .sum_check_sub_claim
+                    .point,
+            )
+            .unwrap()
+            != perm_check_sum_claim
+                .product_check_sub_claim
+                .zero_check_sub_claim
+                .sum_check_sub_claim
+                .expected_evaluation
+        {
+            return Err(PolyIOPErrors::InvalidVerifier("wrong subclaim".to_string()));
+        };
+
+        // test q_x is a 0 over boolean hypercube
+        for i in 0..1 << nv {
+            let bit_sequence = bit_decompose(i, nv);
+            let eval: Vec<E::Fr> = bit_sequence
+                .iter()
+                .map(|x| E::Fr::from(*x as u64))
+                .collect();
+            let res = q_x.evaluate(&eval).unwrap();
+            if !res.is_zero() {}
+        }
+        Ok(())
     }
 
     fn test_permutation_check(nv: usize) -> Result<(), PolyIOPErrors> {
         let mut rng = test_rng();
 
-        let poly_info = VPAuxInfo {
-            max_degree: 2,
-            num_variables: nv,
-            phantom: PhantomData::default(),
-        };
+        let srs = KZGMultilinearPCS::<Bls12_381>::gen_srs_for_testing(&mut rng, nv + 1)?;
+        let (pcs_param, _) = KZGMultilinearPCS::<Bls12_381>::trim(&srs, nv + 1, Some(nv + 1))?;
 
         {
             // good path: w is a permutation of w itself under the identify map
-            let w = DenseMultilinearExtension::rand(nv, &mut rng);
-
+            let w = Rc::new(DenseMultilinearExtension::rand(nv, &mut rng));
             // s_perm is the identity map
             let s_perm = identity_permutation_mle(nv);
-
-            let (proof, q_x) = test_permutation_check_helper(&w, &w, &s_perm)?;
-
-            let mut transcript = <PolyIOP<Fr> as PermutationCheck<Fr>>::init_transcript();
-            transcript.append_message(b"testing", b"initializing transcript for testing")?;
-            let subclaim =
-                <PolyIOP<Fr> as PermutationCheck<Fr>>::verify(&proof, &poly_info, &mut transcript)?
-                    .zero_check_sub_claim;
-            assert_eq!(
-                q_x.evaluate(&subclaim.sum_check_sub_claim.point).unwrap(),
-                subclaim.sum_check_sub_claim.expected_evaluation,
-                "wrong subclaim"
-            );
-
-            // test q_x is a 0 over boolean hypercube
-            for i in 0..1 << nv {
-                let bit_sequence = bit_decompose(i, nv);
-                let eval: Vec<Fr> = bit_sequence.iter().map(|x| Fr::from(*x as u64)).collect();
-                let res = q_x.evaluate(&eval)?;
-                assert!(res.is_zero())
-            }
+            test_permutation_check_helper::<Bls12_381, KZG>(&pcs_param, &w, &w, &s_perm)?;
         }
 
         {
             // bad path 1: w is a not permutation of w itself under a random map
-            let w = DenseMultilinearExtension::rand(nv, &mut rng);
-
+            let w = Rc::new(DenseMultilinearExtension::rand(nv, &mut rng));
             // s_perm is a random map
             let s_perm = random_permutation_mle(nv, &mut rng);
 
-            let (proof, q_x) = test_permutation_check_helper(&w, &w, &s_perm)?;
-
-            let mut transcript = <PolyIOP<Fr> as PermutationCheck<Fr>>::init_transcript();
-            transcript.append_message(b"testing", b"initializing transcript for testing")?;
-            if nv != 1 {
-                assert!(<PolyIOP<Fr> as PermutationCheck<Fr>>::verify(
-                    &proof,
-                    &poly_info,
-                    &mut transcript
+            if nv == 1 {
+                test_permutation_check_helper::<Bls12_381, KZG>(&pcs_param, &w, &w, &s_perm)?;
+            } else {
+                assert!(test_permutation_check_helper::<Bls12_381, KZG>(
+                    &pcs_param, &w, &w, &s_perm
                 )
                 .is_err());
-            } else {
-                // a trivial poly is always a permutation of itself, so the zero check should
-                // pass
-                let subclaim = <PolyIOP<Fr> as PermutationCheck<Fr>>::verify(
-                    &proof,
-                    &poly_info,
-                    &mut transcript,
-                )?
-                .zero_check_sub_claim;
-
-                // the evaluation should fail because a different s_perm is used for proof and
-                // for w |-> w mapping
-                assert_ne!(
-                    q_x.evaluate(&subclaim.sum_check_sub_claim.point).unwrap(),
-                    subclaim.sum_check_sub_claim.expected_evaluation,
-                    "wrong subclaim"
-                );
             }
         }
 
         {
             // bad path 2: f is a not permutation of g under a identity map
-            let f = DenseMultilinearExtension::rand(nv, &mut rng);
-            let g = DenseMultilinearExtension::rand(nv, &mut rng);
-
+            let f = Rc::new(DenseMultilinearExtension::rand(nv, &mut rng));
+            let g = Rc::new(DenseMultilinearExtension::rand(nv, &mut rng));
             // s_perm is the identity map
             let s_perm = identity_permutation_mle(nv);
 
-            let (proof, q_x) = test_permutation_check_helper(&f, &g, &s_perm)?;
-
-            let mut transcript = <PolyIOP<Fr> as PermutationCheck<Fr>>::init_transcript();
-            transcript.append_message(b"testing", b"initializing transcript for testing")?;
-            if nv != 1 {
-                assert!(<PolyIOP<Fr> as PermutationCheck<Fr>>::verify(
-                    &proof,
-                    &poly_info,
-                    &mut transcript
-                )
-                .is_err());
-            } else {
-                // a trivial poly is always a permutation of itself, so the zero check should
-                // pass
-                let subclaim = <PolyIOP<Fr> as PermutationCheck<Fr>>::verify(
-                    &proof,
-                    &poly_info,
-                    &mut transcript,
-                )?
-                .zero_check_sub_claim;
-
-                // the evaluation should fail because a different s_perm is used for proof and
-                // for f |-> g mapping
-                assert_ne!(
-                    q_x.evaluate(&subclaim.sum_check_sub_claim.point).unwrap(),
-                    subclaim.sum_check_sub_claim.expected_evaluation,
-                    "wrong subclaim"
-                );
-            }
+            assert!(
+                test_permutation_check_helper::<Bls12_381, KZG>(&pcs_param, &f, &g, &s_perm)
+                    .is_err()
+            );
         }
 
         Ok(())
@@ -624,48 +365,6 @@
     #[test]
     fn zero_polynomial_should_error() -> Result<(), PolyIOPErrors> {
         assert!(test_permutation_check(0).is_err());
-        Ok(())
-    }
-
-    #[test]
-    fn test_compute_prod() -> Result<(), PolyIOPErrors> {
-        let mut rng = test_rng();
-
-        for num_vars in 2..6 {
-            let f = DenseMultilinearExtension::rand(num_vars, &mut rng);
-            let g = DenseMultilinearExtension::rand(num_vars, &mut rng);
-
-            let s_id = identity_permutation_mle::<Fr>(num_vars);
-            let s_perm = random_permutation_mle(num_vars, &mut rng);
-
-            let mut transcript = <PolyIOP<Fr> as PermutationCheck<Fr>>::init_transcript();
-            transcript.append_message(b"testing", b"initializing transcript for testing")?;
-            let challenge =
-                <PolyIOP<Fr> as PermutationCheck<Fr>>::generate_challenge(&mut transcript)?;
-
-            let prod_and_aux = <PolyIOP<Fr> as PermutationCheck<Fr>>::compute_prod_evals(
-                &challenge, &f, &g, &s_perm,
-            )?;
-            let prod_partial_eval = build_prod_partial_eval(&prod_and_aux[0])?;
-
-            for i in 0..1 << num_vars {
-                let r: Vec<Fr> = bit_decompose(i, num_vars)
-                    .iter()
-                    .map(|&x| Fr::from(x))
-                    .collect();
-
-                let eval = prod_partial_eval[0].evaluate(&r).unwrap();
-
-                let f_eval = f.evaluate(&r).unwrap();
-                let g_eval = g.evaluate(&r).unwrap();
-                let s_id_eval = s_id.evaluate(&r).unwrap();
-                let s_perm_eval = s_perm.evaluate(&r).unwrap();
-                let eval_rec = (f_eval + challenge.beta * s_id_eval + challenge.gamma)
-                    / (g_eval + challenge.beta * s_perm_eval + challenge.gamma);
-
-                assert_eq!(eval, eval_rec);
-            }
-        }
         Ok(())
     }
 }

@@ -1,9 +1,10 @@
 //! This module implements useful functions for the permutation check protocol.
 
-use crate::PolyIOPErrors;
+use crate::errors::PolyIOPErrors;
 use ark_ff::PrimeField;
 use ark_poly::DenseMultilinearExtension;
 use ark_std::{end_timer, rand::RngCore, start_timer};
+use std::rc::Rc;
 
 /// Returns the evaluations of three MLEs:
 /// - prod(0,x)
@@ -25,8 +26,7 @@ /// - denominator is the MLE for `g(x) + \beta s_perm(x) + \gamma`
 ///
 /// The caller needs to check num_vars matches in f/g/s_id/s_perm
 /// Cost: linear in N.
-///
-/// TODO: replace `s_perm` with the merged poly `s`.
+#[cfg(test)]
 #[allow(clippy::type_complexity)]
 pub(super) fn compute_prod_0<F: PrimeField>(
     beta: &F,
@@ -60,17 +60,74 @@ end_timer!(start);
     Ok((prod_0x_evals, numerator_evals, denominator_evals))
 }
 
+/// Returns the evaluations of two MLEs:
+/// - numerator
+/// - denominator
+///
+///  where
+///  - beta and gamma are challenges
+///  - f(x), g(x), s_id(x), s_perm(x) are mle-s
+///
+/// - numerator is the MLE for `f(x) + \beta s_id(x) + \gamma`
+/// - denominator is the MLE for `g(x) + \beta s_perm(x) + \gamma`
+#[allow(clippy::type_complexity)]
+pub(super) fn computer_num_and_denom<F: PrimeField>(
+    beta: &F,
+    gamma: &F,
+    fx: &DenseMultilinearExtension<F>,
+    gx: &DenseMultilinearExtension<F>,
+    s_perm: &DenseMultilinearExtension<F>,
+) -> Result<
+    (
+        Rc<DenseMultilinearExtension<F>>,
+        Rc<DenseMultilinearExtension<F>>,
+    ),
+    PolyIOPErrors,
+> {
+    let start = start_timer!(|| "compute numerator and denominator");
+
+    let num_vars = fx.num_vars;
+    let mut numerator_evals = vec![];
+    let mut denominator_evals = vec![];
+    let s_id = identity_permutation_mle::<F>(num_vars);
+
+    for (&fi, (&gi, (&s_id_i, &s_perm_i))) in
+        fx.iter().zip(gx.iter().zip(s_id.iter().zip(s_perm.iter())))
+    {
+        let numerator = fi + *beta * s_id_i + gamma;
+        let denominator = gi + *beta * s_perm_i + gamma;
+
+        numerator_evals.push(numerator);
+        denominator_evals.push(denominator);
+    }
+    let numerator = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars,
+        numerator_evals,
+    ));
+    let denominator = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars,
+        denominator_evals,
+    ));
+
+    end_timer!(start);
+    Ok((numerator, denominator))
+}
+
 /// An MLE that represent an identity permutation: `f(index) \mapto index`
-pub fn identity_permutation_mle<F: PrimeField>(num_vars: usize) -> DenseMultilinearExtension<F> {
+pub fn identity_permutation_mle<F: PrimeField>(
+    num_vars: usize,
+) -> Rc<DenseMultilinearExtension<F>> {
     let s_id_vec = (0..1u64 << num_vars).map(F::from).collect();
-    DenseMultilinearExtension::from_evaluations_vec(num_vars, s_id_vec)
+    Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars, s_id_vec,
+    ))
 }
 
 /// An MLE that represent a random permutation
 pub fn random_permutation_mle<F: PrimeField, R: RngCore>(
     num_vars: usize,
     rng: &mut R,
-) -> DenseMultilinearExtension<F> {
+) -> Rc<DenseMultilinearExtension<F>> {
     let len = 1u64 << num_vars;
     let mut s_id_vec: Vec<F> = (0..len).map(F::from).collect();
     let mut s_perm_vec = vec![];
@@ -78,7 +135,9 @@ for _ in 0..len {
         let index = rng.next_u64() as usize % s_id_vec.len();
         s_perm_vec.push(s_id_vec.remove(index));
     }
-    DenseMultilinearExtension::from_evaluations_vec(num_vars, s_perm_vec)
+    Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars, s_perm_vec,
+    ))
 }
 
 /// Helper function of the IOP.
@@ -91,22 +150,25 @@ /// - prod(0, x)
 /// - prod(1, x)
 /// - prod(x, 0)
 /// - prod(x, 1)
+#[cfg(test)]
 pub(super) fn build_prod_partial_eval<F: PrimeField>(
-    prod_x: &DenseMultilinearExtension<F>,
-) -> Result<[DenseMultilinearExtension<F>; 4], PolyIOPErrors> {
+    prod_x: &Rc<DenseMultilinearExtension<F>>,
+) -> Result<[Rc<DenseMultilinearExtension<F>>; 4], PolyIOPErrors> {
     let start = start_timer!(|| "build prod polynomial");
 
     let prod_x_eval = &prod_x.evaluations;
     let num_vars = prod_x.num_vars - 1;
 
     // prod(0, x)
-    let prod_0_x =
-        DenseMultilinearExtension::from_evaluations_slice(num_vars, &prod_x_eval[0..1 << num_vars]);
+    let prod_0_x = Rc::new(DenseMultilinearExtension::from_evaluations_slice(
+        num_vars,
+        &prod_x_eval[0..1 << num_vars],
+    ));
     // prod(1, x)
-    let prod_1_x = DenseMultilinearExtension::from_evaluations_slice(
+    let prod_1_x = Rc::new(DenseMultilinearExtension::from_evaluations_slice(
         num_vars,
         &prod_x_eval[1 << num_vars..1 << (num_vars + 1)],
-    );
+    ));
 
     // ===================================
     // prod(x, 0) and prod(x, 1)
@@ -124,8 +186,12 @@ } else {
             eval_x1.push(prod_x);
         }
     }
-    let prod_x_0 = DenseMultilinearExtension::from_evaluations_vec(num_vars, eval_x0);
-    let prod_x_1 = DenseMultilinearExtension::from_evaluations_vec(num_vars, eval_x1);
+    let prod_x_0 = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars, eval_x0,
+    ));
+    let prod_x_1 = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars, eval_x1,
+    ));
 
     end_timer!(start);
 

@@ -4,6 +4,7 @@ perm_check::{
         util::{identity_permutation_mle, random_permutation_mle},
         PermutationCheck,
     },
+    prod_check::ProductCheck,
     sum_check::SumCheck,
     utils::*,
     zero_check::ZeroCheck,

@@ -3,7 +3,8 @@
 use crate::{
     errors::PolyIOPErrors,
     prod_check::util::{compute_product_poly, prove_zero_check},
-    PolyIOP, ZeroCheck,
+    zero_check::ZeroCheck,
+    PolyIOP,
 };
 use arithmetic::VPAuxInfo;
 use ark_ec::PairingEngine;
@@ -11,7 +12,7 @@ use ark_ff::{One, PrimeField, Zero};
 use ark_poly::DenseMultilinearExtension;
 use ark_std::{end_timer, start_timer};
 use pcs::prelude::PolynomialCommitmentScheme;
-use std::{marker::PhantomData, rc::Rc};
+use std::rc::Rc;
 use transcript::IOPTranscript;
 
 mod util;
@@ -42,8 +43,7 @@ E: PairingEngine,
     PCS: PolynomialCommitmentScheme<E>,
 {
     type ProductCheckSubClaim;
-    type ProductProof;
-    type Polynomial;
+    type ProductCheckProof;
 
     /// Initialize the system with a transcript
     ///
@@ -69,17 +69,17 @@ /// - the product polynomial (used for testing)
     ///
     /// Cost: O(N)
     fn prove(
-        fx: &Self::Polynomial,
-        gx: &Self::Polynomial,
+        pcs_param: &PCS::ProverParam,
+        fx: &Self::MultilinearExtension,
+        gx: &Self::MultilinearExtension,
         transcript: &mut IOPTranscript<E::Fr>,
-        pk: &PCS::ProverParam,
-    ) -> Result<(Self::ProductProof, Self::Polynomial), PolyIOPErrors>;
+    ) -> Result<(Self::ProductCheckProof, Self::MultilinearExtension), PolyIOPErrors>;
 
     /// Verify that for witness multilinear polynomials f(x), g(x)
     /// it holds that `\prod_{x \in {0,1}^n} f(x) = \prod_{x \in {0,1}^n} g(x)`
     fn verify(
-        proof: &Self::ProductProof,
-        num_vars: usize,
+        proof: &Self::ProductCheckProof,
+        aux_info: &VPAuxInfo<E::Fr>,
         transcript: &mut Self::Transcript,
     ) -> Result<Self::ProductCheckSubClaim, PolyIOPErrors>;
 }
@@ -98,23 +98,26 @@ // (de)serialize it.
 #[derive(Clone, Debug, Default, PartialEq)]
 pub struct ProductCheckSubClaim<F: PrimeField, ZC: ZeroCheck<F>> {
     // the SubClaim from the ZeroCheck
-    zero_check_sub_claim: ZC::ZeroCheckSubClaim,
+    pub zero_check_sub_claim: ZC::ZeroCheckSubClaim,
     // final query which consists of
     // - the vector `(1, ..., 1, 0)` (needs to be reversed because Arkwork's MLE uses big-endian
     //   format for points)
     // The expected final query evaluation is 1
     final_query: (Vec<F>, F),
-    challenge: F,
+    pub challenge: F,
 }
 
 /// A product check proof consists of
 /// - a zerocheck proof
 /// - a product polynomial commitment
 #[derive(Clone, Debug, Default, PartialEq)]
-pub struct ProductProof<E: PairingEngine, PCS: PolynomialCommitmentScheme<E>, ZC: ZeroCheck<E::Fr>>
-{
-    zero_check_proof: ZC::ZeroCheckProof,
-    prod_x_comm: PCS::Commitment,
+pub struct ProductCheckProof<
+    E: PairingEngine,
+    PCS: PolynomialCommitmentScheme<E>,
+    ZC: ZeroCheck<E::Fr>,
+> {
+    pub zero_check_proof: ZC::ZeroCheckProof,
+    pub prod_x_comm: PCS::Commitment,
 }
 
 impl<E, PCS> ProductCheck<E, PCS> for PolyIOP<E::Fr>
@@ -123,19 +126,18 @@ E: PairingEngine,
     PCS: PolynomialCommitmentScheme<E, Polynomial = Rc<DenseMultilinearExtension<E::Fr>>>,
 {
     type ProductCheckSubClaim = ProductCheckSubClaim<E::Fr, Self>;
-    type ProductProof = ProductProof<E, PCS, Self>;
-    type Polynomial = Rc<DenseMultilinearExtension<E::Fr>>;
+    type ProductCheckProof = ProductCheckProof<E, PCS, Self>;
 
     fn init_transcript() -> Self::Transcript {
         IOPTranscript::<E::Fr>::new(b"Initializing ProductCheck transcript")
     }
 
     fn prove(
-        fx: &Self::Polynomial,
-        gx: &Self::Polynomial,
+        pcs_param: &PCS::ProverParam,
+        fx: &Self::MultilinearExtension,
+        gx: &Self::MultilinearExtension,
         transcript: &mut IOPTranscript<E::Fr>,
-        pk: &PCS::ProverParam,
-    ) -> Result<(Self::ProductProof, Self::Polynomial), PolyIOPErrors> {
+    ) -> Result<(Self::ProductCheckProof, Self::MultilinearExtension), PolyIOPErrors> {
         let start = start_timer!(|| "prod_check prove");
 
         if fx.num_vars != gx.num_vars {
@@ -148,7 +150,7 @@ // compute the product polynomial
         let prod_x = compute_product_poly(fx, gx)?;
 
         // generate challenge
-        let prod_x_comm = PCS::commit(pk, &Rc::new(prod_x.clone()))?;
+        let prod_x_comm = PCS::commit(pcs_param, &prod_x)?;
         transcript.append_serializable_element(b"prod(x)", &prod_x_comm)?;
         let alpha = transcript.get_and_append_challenge(b"alpha")?;
 
@@ -158,17 +160,17 @@
         end_timer!(start);
 
         Ok((
-            ProductProof {
+            ProductCheckProof {
                 zero_check_proof,
                 prod_x_comm,
             },
-            Rc::new(prod_x.clone()),
+            prod_x,
         ))
     }
 
     fn verify(
-        proof: &Self::ProductProof,
-        num_vars: usize,
+        proof: &Self::ProductCheckProof,
+        aux_info: &VPAuxInfo<E::Fr>,
         transcript: &mut Self::Transcript,
     ) -> Result<Self::ProductCheckSubClaim, PolyIOPErrors> {
         let start = start_timer!(|| "prod_check verify");
@@ -179,13 +181,8 @@ let alpha = transcript.get_and_append_challenge(b"alpha")?;
 
         // invoke the zero check on the iop_proof
         // the virtual poly info for Q(x)
-        let aux_info = VPAuxInfo {
-            max_degree: 2,
-            num_variables: num_vars,
-            phantom: PhantomData::default(),
-        };
         let zero_check_sub_claim =
-            <Self as ZeroCheck<E::Fr>>::verify(&proof.zero_check_proof, &aux_info, transcript)?;
+            <Self as ZeroCheck<E::Fr>>::verify(&proof.zero_check_proof, aux_info, transcript)?;
 
         // the final query is on prod_x, hence has length `num_vars` + 1
         let mut final_query = vec![E::Fr::one(); aux_info.num_variables + 1];
@@ -207,18 +204,19 @@ #[cfg(test)]
 mod test {
     use super::ProductCheck;
     use crate::{errors::PolyIOPErrors, PolyIOP};
+    use arithmetic::VPAuxInfo;
     use ark_bls12_381::{Bls12_381, Fr};
     use ark_ec::PairingEngine;
     use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
     use ark_std::test_rng;
     use pcs::{prelude::KZGMultilinearPCS, PolynomialCommitmentScheme};
-    use std::rc::Rc;
+    use std::{marker::PhantomData, rc::Rc};
 
     // f and g are guaranteed to have the same product
     fn test_product_check_helper<E, PCS>(
         f: &DenseMultilinearExtension<E::Fr>,
         g: &DenseMultilinearExtension<E::Fr>,
-        pk: &PCS::ProverParam,
+        pcs_param: &PCS::ProverParam,
     ) -> Result<(), PolyIOPErrors>
     where
         E: PairingEngine,
@@ -228,17 +226,22 @@ let mut transcript = <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::init_transcript();
         transcript.append_message(b"testing", b"initializing transcript for testing")?;
 
         let (proof, prod_x) = <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::prove(
+            pcs_param,
             &Rc::new(f.clone()),
             &Rc::new(g.clone()),
             &mut transcript,
-            pk,
         )?;
 
         let mut transcript = <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::init_transcript();
         transcript.append_message(b"testing", b"initializing transcript for testing")?;
 
+        let aux_info = VPAuxInfo {
+            max_degree: 2,
+            num_variables: f.num_vars,
+            phantom: PhantomData::default(),
+        };
         let subclaim =
-            <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::verify(&proof, f.num_vars, &mut transcript)?;
+            <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::verify(&proof, &aux_info, &mut transcript)?;
         assert_eq!(
             prod_x.evaluate(&subclaim.final_query.0).unwrap(),
             subclaim.final_query.1,
@@ -251,17 +254,17 @@ transcript.append_message(b"testing", b"initializing transcript for testing")?;
 
         let h = f + g;
         let (bad_proof, prod_x_bad) = <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::prove(
+            pcs_param,
             &Rc::new(f.clone()),
             &Rc::new(h.clone()),
             &mut transcript,
-            pk,
         )?;
 
         let mut transcript = <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::init_transcript();
         transcript.append_message(b"testing", b"initializing transcript for testing")?;
         let bad_subclaim = <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::verify(
             &bad_proof,
-            f.num_vars,
+            &aux_info,
             &mut transcript,
         )?;
         assert_ne!(
@@ -281,9 +284,9 @@ let mut g = f.clone();
         g.evaluations.reverse();
 
         let srs = KZGMultilinearPCS::<Bls12_381>::gen_srs_for_testing(&mut rng, nv + 1)?;
-        let (pk, _) = KZGMultilinearPCS::<Bls12_381>::trim(&srs, nv + 1, Some(nv + 1))?;
+        let (pcs_param, _) = KZGMultilinearPCS::<Bls12_381>::trim(&srs, nv + 1, Some(nv + 1))?;
 
-        test_product_check_helper::<Bls12_381, KZGMultilinearPCS<Bls12_381>>(&f, &g, &pk)?;
+        test_product_check_helper::<Bls12_381, KZGMultilinearPCS<Bls12_381>>(&f, &g, &pcs_param)?;
 
         Ok(())
     }

@@ -1,6 +1,8 @@
 //! This module implements useful functions for the product check protocol.
 
-use crate::{errors::PolyIOPErrors, structs::IOPProof, utils::get_index, PolyIOP, ZeroCheck};
+use crate::{
+    errors::PolyIOPErrors, structs::IOPProof, utils::get_index, zero_check::ZeroCheck, PolyIOP,
+};
 use arithmetic::VirtualPolynomial;
 use ark_ff::PrimeField;
 use ark_poly::DenseMultilinearExtension;
@@ -19,9 +21,9 @@ ///
 /// The caller needs to check num_vars matches in f and g
 /// Cost: linear in N.
 pub(super) fn compute_product_poly<F: PrimeField>(
-    fx: &DenseMultilinearExtension<F>,
-    gx: &DenseMultilinearExtension<F>,
-) -> Result<DenseMultilinearExtension<F>, PolyIOPErrors> {
+    fx: &Rc<DenseMultilinearExtension<F>>,
+    gx: &Rc<DenseMultilinearExtension<F>>,
+) -> Result<Rc<DenseMultilinearExtension<F>>, PolyIOPErrors> {
     let start = start_timer!(|| "compute evaluations of prod polynomial");
     let num_vars = fx.num_vars;
 
@@ -69,7 +71,10 @@ // ===================================
     // prod(x)'s evaluation is indeed `e := [eval_0x[..], eval_1x[..]].concat()`
     let eval = [prod_0x_eval.as_slice(), prod_1x_eval.as_slice()].concat();
 
-    let prod_x = DenseMultilinearExtension::from_evaluations_vec(num_vars + 1, eval);
+    let prod_x = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars + 1,
+        eval,
+    ));
 
     end_timer!(start);
     Ok(prod_x)
@@ -82,21 +87,19 @@ /// Returns proof and Q(x) for testing purpose.
 ///
 /// Cost: O(N)
 pub(super) fn prove_zero_check<F: PrimeField>(
-    fx: &DenseMultilinearExtension<F>,
-    gx: &DenseMultilinearExtension<F>,
-    prod_x: &DenseMultilinearExtension<F>,
+    fx: &Rc<DenseMultilinearExtension<F>>,
+    gx: &Rc<DenseMultilinearExtension<F>>,
+    prod_x: &Rc<DenseMultilinearExtension<F>>,
     alpha: &F,
     transcript: &mut IOPTranscript<F>,
 ) -> Result<(IOPProof<F>, VirtualPolynomial<F>), PolyIOPErrors> {
     let start = start_timer!(|| "zerocheck in product check");
 
     let prod_partial_evals = build_prod_partial_eval(prod_x)?;
-    let prod_0x = Rc::new(prod_partial_evals[0].clone());
-    let prod_1x = Rc::new(prod_partial_evals[1].clone());
-    let prod_x0 = Rc::new(prod_partial_evals[2].clone());
-    let prod_x1 = Rc::new(prod_partial_evals[3].clone());
-    let fx = Rc::new(fx.clone());
-    let gx = Rc::new(gx.clone());
+    let prod_0x = prod_partial_evals[0].clone();
+    let prod_1x = prod_partial_evals[1].clone();
+    let prod_x0 = prod_partial_evals[2].clone();
+    let prod_x1 = prod_partial_evals[3].clone();
 
     // compute g(x) * prod(0, x) * alpha
     let mut q_x = VirtualPolynomial::new_from_mle(gx, F::one());
@@ -104,7 +107,7 @@ q_x.mul_by_mle(prod_0x, *alpha)?;
 
     //   g(x) * prod(0, x) * alpha
     // - f(x) * alpha
-    q_x.add_mle_list([fx], -*alpha)?;
+    q_x.add_mle_list([fx.clone()], -*alpha)?;
 
     // Q(x) := prod(1,x) - prod(x, 0) * prod(x, 1)
     //       + alpha * (
@@ -130,21 +133,23 @@ /// - prod(1, x)
 /// - prod(x, 0)
 /// - prod(x, 1)
 fn build_prod_partial_eval<F: PrimeField>(
-    prod_x: &DenseMultilinearExtension<F>,
-) -> Result<[DenseMultilinearExtension<F>; 4], PolyIOPErrors> {
+    prod_x: &Rc<DenseMultilinearExtension<F>>,
+) -> Result<[Rc<DenseMultilinearExtension<F>>; 4], PolyIOPErrors> {
     let start = start_timer!(|| "build partial prod polynomial");
 
     let prod_x_eval = &prod_x.evaluations;
     let num_vars = prod_x.num_vars - 1;
 
     // prod(0, x)
-    let prod_0_x =
-        DenseMultilinearExtension::from_evaluations_slice(num_vars, &prod_x_eval[0..1 << num_vars]);
+    let prod_0_x = Rc::new(DenseMultilinearExtension::from_evaluations_slice(
+        num_vars,
+        &prod_x_eval[0..1 << num_vars],
+    ));
     // prod(1, x)
-    let prod_1_x = DenseMultilinearExtension::from_evaluations_slice(
+    let prod_1_x = Rc::new(DenseMultilinearExtension::from_evaluations_slice(
         num_vars,
         &prod_x_eval[1 << num_vars..1 << (num_vars + 1)],
-    );
+    ));
 
     // ===================================
     // prod(x, 0) and prod(x, 1)
@@ -162,8 +167,12 @@ } else {
             eval_x1.push(prod_x);
         }
     }
-    let prod_x_0 = DenseMultilinearExtension::from_evaluations_vec(num_vars, eval_x0);
-    let prod_x_1 = DenseMultilinearExtension::from_evaluations_vec(num_vars, eval_x1);
+    let prod_x_0 = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars, eval_x0,
+    ));
+    let prod_x_1 = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars, eval_x1,
+    ));
 
     end_timer!(start);
 

@@ -8,7 +8,7 @@ /// An IOP proof is a collections of
 /// - messages from prover to verifier at each round through the interactive
 ///   protocol.
 /// - a point that is generated by the transcript for evaluation
-#[derive(Clone, Debug, Default, PartialEq)]
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub struct IOPProof<F: PrimeField> {
     pub point: Vec<F>,
     pub proofs: Vec<IOPProverMessage<F>>,

@@ -9,6 +9,7 @@ use arithmetic::{VPAuxInfo, VirtualPolynomial};
 use ark_ff::PrimeField;
 use ark_poly::DenseMultilinearExtension;
 use ark_std::{end_timer, start_timer};
+use std::{fmt::Debug, rc::Rc};
 use transcript::IOPTranscript;
 
 mod prover;
@@ -20,9 +21,9 @@ type VirtualPolynomial;
     type VPAuxInfo;
     type MultilinearExtension;
 
-    type SumCheckProof;
+    type SumCheckProof: Clone + Debug + Default + PartialEq;
     type Transcript;
-    type SumCheckSubClaim;
+    type SumCheckSubClaim: Clone + Debug + Default + PartialEq;
 
     /// Extract sum from the proof
     fn extract_sum(proof: &Self::SumCheckProof) -> F;
@@ -126,7 +127,7 @@ impl<F: PrimeField> SumCheck<F> for PolyIOP<F> {
     type SumCheckProof = IOPProof<F>;
     type VirtualPolynomial = VirtualPolynomial<F>;
     type VPAuxInfo = VPAuxInfo<F>;
-    type MultilinearExtension = DenseMultilinearExtension<F>;
+    type MultilinearExtension = Rc<DenseMultilinearExtension<F>>;
     type SumCheckSubClaim = SumCheckSubClaim<F>;
     type Transcript = IOPTranscript<F>;
 

@@ -1,5 +1,7 @@
 //! Main module for the ZeroCheck protocol.
 
+use std::fmt::Debug;
+
 use crate::{errors::PolyIOPErrors, sum_check::SumCheck, PolyIOP};
 use arithmetic::build_eq_x_r;
 use ark_ff::PrimeField;
@@ -22,8 +24,8 @@ }
 
 /// A ZeroCheck is derived from SumCheck.
 pub trait ZeroCheck<F: PrimeField>: SumCheck<F> {
-    type ZeroCheckSubClaim;
-    type ZeroCheckProof;
+    type ZeroCheckSubClaim: Clone + Debug + Default + PartialEq;
+    type ZeroCheckProof: Clone + Debug + Default + PartialEq;
 
     /// Initialize the system with a transcript
     ///