refactor prodcheck

subroutines/src/poly_iop/perm_check/mod.rs

@@ -121,8 +121,12 @@ where
         let (numerator, denominator) = computer_num_and_denom(&beta, &gamma, fx, gx, s_perm)?;
 
         // invoke product check on numerator and denominator
-        let (proof, prod_poly) =
+        let (proof, prod_poly, _frac_poly) = <Self as ProductCheck<E, PCS>>::prove(
-            <Self as ProductCheck<E, PCS>>::prove(pcs_param, &numerator, &denominator, transcript)?;
+            pcs_param,
+            &[numerator],
+            &[denominator],
+            transcript,
+        )?;
 
         end_timer!(start);
         Ok((proof, prod_poly))

subroutines/src/poly_iop/prod_check/mod.rs

@@ -4,7 +4,7 @@ use crate::{
     pcs::PolynomialCommitmentScheme,
     poly_iop::{
         errors::PolyIOPErrors,
-        prod_check::util::{compute_product_poly, prove_zero_check},
+        prod_check::util::{compute_frac_poly, compute_product_poly, prove_zero_check},
         zero_check::ZeroCheck,
         PolyIOP,
     },
@@ -19,22 +19,29 @@ use transcript::IOPTranscript;
 
 mod util;
 
-/// A product-check proves that two n-variate multilinear polynomials `f(x),
+/// A product-check proves that two lists of n-variate multilinear polynomials
-/// g(x)` satisfy:
+/// `(f1, f2, ..., fk)` and `(g1, ..., gk)` satisfy:
-/// \prod_{x \in {0,1}^n} f(x) = \prod_{x \in {0,1}^n} g(x)
+/// \prod_{x \in {0,1}^n} f1(x) * ... * fk(x) = \prod_{x \in {0,1}^n} g1(x) *
+/// ... * gk(x)
 ///
 /// A ProductCheck is derived from ZeroCheck.
 ///
 /// Prover steps:
-/// 1. build `prod(x0, ..., x_n)` from f and g,
+/// 1. build MLE `frac(x)` s.t. `frac(x) = f1(x) * ... * fk(x) / (g1(x) * ... *
-///    such that `prod(0, x1, ..., xn)` equals `f/g` over domain {0,1}^n
+/// gk(x))` for all x \in {0,1}^n 2. build `prod(x)` from `frac(x)`, where
-/// 2. push commitments of `prod(x)` to the transcript,
+/// `prod(x)` equals to `v(1,x)` in the paper 2. push commitments of `frac(x)`
-///    and `generate_challenge` from current transcript (generate alpha)
+/// and `prod(x)` to the transcript,    and `generate_challenge` from current
-/// 3. generate the zerocheck proof for the virtual polynomial
+/// transcript (generate alpha) 3. generate the zerocheck proof for the virtual
-///    prod(1, x) - prod(x, 0) * prod(x, 1) + alpha * (f(x) - prod(0, x) * g(x))
+/// polynomial Q(x):       prod(x) - p1(x) * p2(x)
+///     + alpha * frac(x) * g1(x) * ... * gk(x)
+///     - alpha * f1(x) * ... * fk(x)
+/// where p1(x) = (1-x1) * frac(x2, ..., xn, 0)
+///             + x1 * prod(x2, ..., xn, 0),
+/// and p2(x) = (1-x1) * frac(x2, ..., xn, 1)
+///           + x1 * prod(x2, ..., xn, 1)
 ///
 /// Verifier steps:
-/// 1. Extract commitments of `prod(x)` from the proof, push
+/// 1. Extract commitments of `frac(x)` and `prod(x)` from the proof, push
 /// them to the transcript
 /// 2. `generate_challenge` from current transcript (generate alpha)
 /// 3. `verify` to verify the zerocheck proof and generate the subclaim for
@@ -55,30 +62,42 @@ where
     /// ProductCheck prover/verifier.
     fn init_transcript() -> Self::Transcript;
 
-    /// Generate a proof for product check, showing that witness multilinear
+    /// Proves that two lists of n-variate multilinear polynomials `(f1, f2,
-    /// polynomials f(x), g(x) satisfy `\prod_{x \in {0,1}^n} f(x) =
+    /// ..., fk)` and `(g1, ..., gk)` satisfy:
-    /// \prod_{x \in {0,1}^n} g(x)`
+    ///   \prod_{x \in {0,1}^n} f1(x) * ... * fk(x)
+    /// = \prod_{x \in {0,1}^n} g1(x) * ... * gk(x)
     ///
     /// Inputs:
-    /// - fx: the numerator multilinear polynomial
+    /// - fxs: the list of numerator multilinear polynomial
-    /// - gx: the denominator multilinear polynomial
+    /// - gxs: the list of denominator multilinear polynomial
     /// - transcript: the IOP transcript
     /// - pk: PCS committing key
     ///
     /// Outputs
     /// - the product check proof
     /// - the product polynomial (used for testing)
+    /// - the fractional polynomial (used for testing)
     ///
     /// Cost: O(N)
+    #[allow(clippy::type_complexity)]
     fn prove(
         pcs_param: &PCS::ProverParam,
-        fx: &Self::MultilinearExtension,
+        fxs: &[Self::MultilinearExtension],
-        gx: &Self::MultilinearExtension,
+        gxs: &[Self::MultilinearExtension],
         transcript: &mut IOPTranscript<E::Fr>,
-    ) -> Result<(Self::ProductCheckProof, Self::MultilinearExtension), PolyIOPErrors>;
+    ) -> Result<
+        (
+            Self::ProductCheckProof,
+            Self::MultilinearExtension,
+            Self::MultilinearExtension,
+        ),
+        PolyIOPErrors,
+    >;
 
-    /// Verify that for witness multilinear polynomials f(x), g(x)
+    /// Verify that for witness multilinear polynomials (f1, ..., fk, g1, ...,
-    /// it holds that `\prod_{x \in {0,1}^n} f(x) = \prod_{x \in {0,1}^n} g(x)`
+    /// gk) it holds that
+    ///      `\prod_{x \in {0,1}^n} f1(x) * ... * fk(x)
+    ///     = \prod_{x \in {0,1}^n} g1(x) * ... * gk(x)`
     fn verify(
         proof: &Self::ProductCheckProof,
         aux_info: &VPAuxInfo<E::Fr>,
@@ -87,9 +106,7 @@ where
 }
 
 /// A product check subclaim consists of
-/// - A zero check IOP subclaim for
+/// - A zero check IOP subclaim for the virtual polynomial
-/// `Q(x) = prod(1, x) - prod(x, 0) * prod(x, 1) + alpha * (f(x) - prod(0,
-/// x) * g(x)) = 0`
 /// - The random challenge `alpha`
 /// - A final query for `prod(1, ..., 1, 0) = 1`.
 // Note that this final query is in fact a constant that
@@ -110,6 +127,7 @@ pub struct ProductCheckSubClaim<F: PrimeField, ZC: ZeroCheck<F>> {
 /// A product check proof consists of
 /// - a zerocheck proof
 /// - a product polynomial commitment
+/// - a polynomial commitment for the fractional polynomial
 #[derive(Clone, Debug, Default, PartialEq)]
 pub struct ProductCheckProof<
     E: PairingEngine,
@@ -118,6 +136,7 @@ pub struct ProductCheckProof<
 > {
     pub zero_check_proof: ZC::ZeroCheckProof,
     pub prod_x_comm: PCS::Commitment,
+    pub frac_comm: PCS::Commitment,
 }
 
 impl<E, PCS> ProductCheck<E, PCS> for PolyIOP<E::Fr>
@@ -134,28 +153,51 @@ where
 
     fn prove(
         pcs_param: &PCS::ProverParam,
-        fx: &Self::MultilinearExtension,
+        fxs: &[Self::MultilinearExtension],
-        gx: &Self::MultilinearExtension,
+        gxs: &[Self::MultilinearExtension],
         transcript: &mut IOPTranscript<E::Fr>,
-    ) -> Result<(Self::ProductCheckProof, Self::MultilinearExtension), PolyIOPErrors> {
+    ) -> Result<
+        (
+            Self::ProductCheckProof,
+            Self::MultilinearExtension,
+            Self::MultilinearExtension,
+        ),
+        PolyIOPErrors,
+    > {
         let start = start_timer!(|| "prod_check prove");
 
-        if fx.num_vars != gx.num_vars {
+        if fxs.is_empty() {
+            return Err(PolyIOPErrors::InvalidParameters("fxs is empty".to_string()));
+        }
+        if fxs.len() != gxs.len() {
             return Err(PolyIOPErrors::InvalidParameters(
-                "fx and gx have different number of variables".to_string(),
+                "fxs and gxs have different number of polynomials".to_string(),
             ));
         }
+        for poly in fxs.iter().chain(gxs.iter()) {
+            if poly.num_vars != fxs[0].num_vars {
+                return Err(PolyIOPErrors::InvalidParameters(
+                    "fx and gx have different number of variables".to_string(),
+                ));
+            }
+        }
 
+        // compute the fractional polynomial frac_p s.t.
+        // frac_p(x) = f1(x) * ... * fk(x) / (g1(x) * ... * gk(x))
+        let frac_poly = compute_frac_poly(fxs, gxs)?;
         // compute the product polynomial
-        let prod_x = compute_product_poly(fx, gx)?;
+        let prod_x = compute_product_poly(&frac_poly)?;
 
         // generate challenge
+        let frac_comm = PCS::commit(pcs_param, &frac_poly)?;
         let prod_x_comm = PCS::commit(pcs_param, &prod_x)?;
+        transcript.append_serializable_element(b"frac(x)", &frac_comm)?;
         transcript.append_serializable_element(b"prod(x)", &prod_x_comm)?;
         let alpha = transcript.get_and_append_challenge(b"alpha")?;
 
         // build the zero-check proof
-        let (zero_check_proof, _) = prove_zero_check(fx, gx, &prod_x, &alpha, transcript)?;
+        let (zero_check_proof, _) =
+            prove_zero_check(fxs, gxs, &frac_poly, &prod_x, &alpha, transcript)?;
 
         end_timer!(start);
 
@@ -163,8 +205,10 @@ where
             ProductCheckProof {
                 zero_check_proof,
                 prod_x_comm,
+                frac_comm,
             },
             prod_x,
+            frac_poly,
         ))
     }
 
@@ -176,6 +220,7 @@ where
         let start = start_timer!(|| "prod_check verify");
 
         // update transcript and generate challenge
+        transcript.append_serializable_element(b"frac(x)", &proof.frac_comm)?;
         transcript.append_serializable_element(b"prod(x)", &proof.prod_x_comm)?;
         let alpha = transcript.get_and_append_challenge(b"alpha")?;
 
@@ -184,8 +229,8 @@ where
         let zero_check_sub_claim =
             <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
+        // the final query is on prod_x
-        let mut final_query = vec![E::Fr::one(); aux_info.num_variables + 1];
+        let mut final_query = vec![E::Fr::one(); aux_info.num_variables];
         // the point has to be reversed because Arkworks uses big-endian.
         final_query[0] = E::Fr::zero();
         let final_eval = E::Fr::one();
@@ -214,10 +259,35 @@ mod test {
     use ark_std::test_rng;
     use std::{marker::PhantomData, rc::Rc};
 
-    // f and g are guaranteed to have the same product
+    fn check_frac_poly<E>(
+        frac_poly: &Rc<DenseMultilinearExtension<E::Fr>>,
+        fs: &[Rc<DenseMultilinearExtension<E::Fr>>],
+        gs: &[Rc<DenseMultilinearExtension<E::Fr>>],
+    ) where
+        E: PairingEngine,
+    {
+        let mut flag = true;
+        let num_vars = frac_poly.num_vars;
+        for i in 0..1 << num_vars {
+            let nom = fs
+                .iter()
+                .fold(E::Fr::from(1u8), |acc, f| acc * f.evaluations[i]);
+            let denom = gs
+                .iter()
+                .fold(E::Fr::from(1u8), |acc, g| acc * g.evaluations[i]);
+            if denom * frac_poly.evaluations[i] != nom {
+                flag = false;
+                break;
+            }
+        }
+        assert_eq!(flag, true);
+    }
+    // fs and gs are guaranteed to have the same product
+    // fs and hs doesn't have the same product
     fn test_product_check_helper<E, PCS>(
-        f: &DenseMultilinearExtension<E::Fr>,
+        fs: &[Rc<DenseMultilinearExtension<E::Fr>>],
-        g: &DenseMultilinearExtension<E::Fr>,
+        gs: &[Rc<DenseMultilinearExtension<E::Fr>>],
+        hs: &[Rc<DenseMultilinearExtension<E::Fr>>],
         pcs_param: &PCS::ProverParam,
     ) -> Result<(), PolyIOPErrors>
     where
@@ -227,19 +297,16 @@ mod test {
         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(
+        let (proof, prod_x, frac_poly) =
-            pcs_param,
+            <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::prove(pcs_param, fs, gs, &mut transcript)?;
-            &Rc::new(f.clone()),
-            &Rc::new(g.clone()),
-            &mut transcript,
-        )?;
 
         let mut transcript = <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::init_transcript();
         transcript.append_message(b"testing", b"initializing transcript for testing")?;
 
+        // what's aux_info for?
         let aux_info = VPAuxInfo {
-            max_degree: 2,
+            max_degree: fs.len() + 1,
-            num_variables: f.num_vars,
+            num_variables: fs[0].num_vars,
             phantom: PhantomData::default(),
         };
         let prod_subclaim =
@@ -249,18 +316,14 @@ mod test {
             prod_subclaim.final_query.1,
             "different product"
         );
+        check_frac_poly::<E>(&frac_poly, fs, gs);
 
         // bad path
         let mut transcript = <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::init_transcript();
         transcript.append_message(b"testing", b"initializing transcript for testing")?;
 
-        let h = f + g;
+        let (bad_proof, prod_x_bad, frac_poly) =
-        let (bad_proof, prod_x_bad) = <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::prove(
+            <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::prove(pcs_param, fs, hs, &mut transcript)?;
-            pcs_param,
-            &Rc::new(f.clone()),
-            &Rc::new(h),
-            &mut transcript,
-        )?;
 
         let mut transcript = <PolyIOP<E::Fr> as ProductCheck<E, PCS>>::init_transcript();
         transcript.append_message(b"testing", b"initializing transcript for testing")?;
@@ -274,6 +337,8 @@ mod test {
             bad_subclaim.final_query.1,
             "can't detect wrong proof"
         );
+        // the frac_poly should still be computed correctly
+        check_frac_poly::<E>(&frac_poly, fs, hs);
 
         Ok(())
     }
@@ -281,14 +346,28 @@ mod test {
     fn test_product_check(nv: usize) -> Result<(), PolyIOPErrors> {
         let mut rng = test_rng();
 
-        let f: DenseMultilinearExtension<Fr> = DenseMultilinearExtension::rand(nv, &mut rng);
+        let f1: DenseMultilinearExtension<Fr> = DenseMultilinearExtension::rand(nv, &mut rng);
-        let mut g = f.clone();
+        let mut g1 = f1.clone();
-        g.evaluations.reverse();
+        g1.evaluations.reverse();
+        let f2: DenseMultilinearExtension<Fr> = DenseMultilinearExtension::rand(nv, &mut rng);
+        let mut g2 = f2.clone();
+        g2.evaluations.reverse();
+        let fs = vec![Rc::new(f1), Rc::new(f2)];
+        let gs = vec![Rc::new(g2), Rc::new(g1)];
+        let mut hs = vec![];
+        for _ in 0..fs.len() {
+            hs.push(Rc::new(DenseMultilinearExtension::rand(
+                fs[0].num_vars,
+                &mut rng,
+            )));
+        }
 
-        let srs = MultilinearKzgPCS::<Bls12_381>::gen_srs_for_testing(&mut rng, nv + 1)?;
+        let srs = MultilinearKzgPCS::<Bls12_381>::gen_srs_for_testing(&mut rng, nv)?;
-        let (pcs_param, _) = MultilinearKzgPCS::<Bls12_381>::trim(&srs, None, Some(nv + 1))?;
+        let (pcs_param, _) = MultilinearKzgPCS::<Bls12_381>::trim(&srs, None, Some(nv))?;
 
-        test_product_check_helper::<Bls12_381, MultilinearKzgPCS<Bls12_381>>(&f, &g, &pcs_param)?;
+        test_product_check_helper::<Bls12_381, MultilinearKzgPCS<Bls12_381>>(
+            &fs, &gs, &hs, &pcs_param,
+        )?;
 
         Ok(())
     }

subroutines/src/poly_iop/prod_check/util.rs

@@ -5,198 +5,166 @@ use arithmetic::{get_index, VirtualPolynomial};
 use ark_ff::PrimeField;
 use ark_poly::DenseMultilinearExtension;
 use ark_std::{end_timer, start_timer};
-use rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
 use std::rc::Rc;
 use transcript::IOPTranscript;
 
-/// Compute the product polynomial `prod(x)` where
+/// Compute multilinear fractional polynomial s.t. frac(x) = f1(x) * ... * fk(x)
+/// / (g1(x) * ... * gk(x)) for all x \in {0,1}^n
 ///
-///  - `prod(0,x) := prod(0, x1, …, xn)` is the MLE over the
+/// The caller needs to sanity-check that the number of polynomials and
-/// evaluations of `f(x)/g(x)` on the boolean hypercube {0,1}^n
+/// variables match in fxs and gxs; and gi(x) has no zero entries.
-///
+pub(super) fn compute_frac_poly<F: PrimeField>(
-/// - `prod(1,x)` is a MLE over the evaluations of `prod(x, 0) * prod(x, 1)`
+    fxs: &[Rc<DenseMultilinearExtension<F>>],
-/// on the boolean hypercube {0,1}^n
+    gxs: &[Rc<DenseMultilinearExtension<F>>],
+) -> Result<Rc<DenseMultilinearExtension<F>>, PolyIOPErrors> {
+    let start = start_timer!(|| "compute frac(x)");
+
+    let mut f_evals = vec![F::one(); 1 << fxs[0].num_vars];
+    for fx in fxs.iter() {
+        for (f_eval, fi) in f_evals.iter_mut().zip(fx.iter()) {
+            *f_eval *= fi;
+        }
+    }
+    let mut g_evals = vec![F::one(); 1 << gxs[0].num_vars];
+    for gx in gxs.iter() {
+        for (g_eval, gi) in g_evals.iter_mut().zip(gx.iter()) {
+            *g_eval *= gi;
+        }
+    }
+    for (f_eval, g_eval) in f_evals.iter_mut().zip(g_evals.iter()) {
+        if *g_eval == F::zero() {
+            return Err(PolyIOPErrors::InvalidParameters(
+                "gxs has zero entries in the boolean hypercube".to_string(),
+            ));
+        }
+        *f_eval /= g_eval;
+    }
+
+    end_timer!(start);
+    Ok(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        fxs[0].num_vars,
+        f_evals,
+    )))
+}
+
+/// Compute the product polynomial `prod(x)` such that
+/// `prod(x) = [(1-x1)*frac(x2, ..., xn, 0) + x1*prod(x2, ..., xn, 0)] *
+/// [(1-x1)*frac(x2, ..., xn, 1) + x1*prod(x2, ..., xn, 1)]` on the boolean
+/// hypercube {0,1}^n
 ///
 /// 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: &Rc<DenseMultilinearExtension<F>>,
+    frac_poly: &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;
+    let num_vars = frac_poly.num_vars;
-
+    let frac_evals = &frac_poly.evaluations;
-    // ===================================
-    // prod(0, x)
-    // ===================================
-    let prod_0x_eval = compute_prod_0(fx, gx)?;
-
-    // ===================================
-    // 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]);
-        }
-    }
-
-    // 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();
+    // `prod(x)` can be computed via recursing the following formula for 2^n-1
-
+    // times
-    let prod_x = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+    //
-        num_vars + 1,
+    // `prod(x_1, ..., x_n) :=
-        eval,
+    //      [(1-x1)*frac(x2, ..., xn, 0) + x1*prod(x2, ..., xn, 0)] *
-    ));
+    //      [(1-x1)*frac(x2, ..., xn, 1) + x1*prod(x2, ..., xn, 1)]`
+    //
+    // At any given step, the right hand side of the equation
+    // is available via either frac_x or the current view of prod_x
+    let mut prod_x_evals = vec![];
+    for x in 0..(1 << num_vars) - 1 {
+        // sign will decide if the evaluation should be looked up from frac_x or
+        // prod_x; 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_x_evals.push(frac_evals[x_zero_index] * frac_evals[x_one_index]);
+        } else {
+            // sanity check: if we are trying to look up from the prod_x_evals table,
+            // then the target index must already exist
+            if x_zero_index >= prod_x_evals.len() || x_one_index >= prod_x_evals.len() {
+                return Err(PolyIOPErrors::ShouldNotArrive);
+            }
+            prod_x_evals.push(prod_x_evals[x_zero_index] * prod_x_evals[x_one_index]);
+        }
+    }
 
+    // prod(1, 1, ..., 1) := 0
+    prod_x_evals.push(F::zero());
     end_timer!(start);
-    Ok(prod_x)
+
+    Ok(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars,
+        prod_x_evals,
+    )))
 }
 
 /// generate the zerocheck proof for the virtual polynomial
-///    prod(1, x) - prod(x, 0) * prod(x, 1) + alpha * (f(x) - prod(0, x) * g(x))
+///    prod(x) - p1(x) * p2(x) + alpha * [frac(x) * g1(x) * ... * gk(x) - f1(x)
-///
+/// * ... * fk(x)] where p1(x) = (1-x1) * frac(x2, ..., xn, 0) + x1 * prod(x2,
-/// Returns proof and Q(x) for testing purpose.
+///   ..., xn, 0), p2(x) = (1-x1) * frac(x2, ..., xn, 1) + x1 * prod(x2, ...,
+///   xn, 1)
+/// Returns proof.
 ///
 /// Cost: O(N)
 pub(super) fn prove_zero_check<F: PrimeField>(
-    fx: &Rc<DenseMultilinearExtension<F>>,
+    fxs: &[Rc<DenseMultilinearExtension<F>>],
-    gx: &Rc<DenseMultilinearExtension<F>>,
+    gxs: &[Rc<DenseMultilinearExtension<F>>],
+    frac_poly: &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 num_vars = frac_poly.num_vars;
 
-    let prod_partial_evals = build_prod_partial_eval(prod_x)?;
+    // compute p1(x) = (1-x1) * frac(x2, ..., xn, 0) + x1 * prod(x2, ..., xn, 0)
-    let prod_0x = prod_partial_evals[0].clone();
+    // compute p2(x) = (1-x1) * frac(x2, ..., xn, 1) + x1 * prod(x2, ..., xn, 1)
-    let prod_1x = prod_partial_evals[1].clone();
+    let mut p1_evals = vec![F::zero(); 1 << num_vars];
-    let prod_x0 = prod_partial_evals[2].clone();
+    let mut p2_evals = vec![F::zero(); 1 << num_vars];
-    let prod_x1 = prod_partial_evals[3].clone();
+    for x in 0..1 << num_vars {
+        let (x0, x1, sign) = get_index(x, num_vars);
+        if !sign {
+            p1_evals[x] = frac_poly.evaluations[x0];
+            p2_evals[x] = frac_poly.evaluations[x1];
+        } else {
+            p1_evals[x] = prod_x.evaluations[x0];
+            p2_evals[x] = prod_x.evaluations[x1];
+        }
+    }
+    let p1 = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars, p1_evals,
+    ));
+    let p2 = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
+        num_vars, p2_evals,
+    ));
 
-    // compute g(x) * prod(0, x) * alpha
+    // compute Q(x)
-    let mut q_x = VirtualPolynomial::new_from_mle(gx, F::one());
+    // prod(x)
-    q_x.mul_by_mle(prod_0x, *alpha)?;
+    let mut q_x = VirtualPolynomial::new_from_mle(prod_x, F::one());
 
-    //   g(x) * prod(0, x) * alpha
+    //   prod(x)
-    // - f(x) * alpha
+    // - p1(x) * p2(x)
-    q_x.add_mle_list([fx.clone()], -*alpha)?;
+    q_x.add_mle_list([p1, p2], -F::one())?;
 
-    // Q(x) := prod(1,x) - prod(x, 0) * prod(x, 1)
+    //   prod(x)
-    //       + alpha * (
+    // - p1(x) * p2(x)
-    //             g(x) * prod(0, x)
+    // + alpha * frac(x) * g1(x) * ... * gk(x)
-    //           - f(x))
+    let mut mle_list = gxs.to_vec();
-    q_x.add_mle_list([prod_x0, prod_x1], -F::one())?;
+    mle_list.push(frac_poly.clone());
-    q_x.add_mle_list([prod_1x], F::one())?;
+    q_x.add_mle_list(mle_list, *alpha)?;
+
+    //   prod(x)
+    // - p1(x) * p2(x)
+    // + alpha * frac(x) * g1(x) * ... * gk(x)
+    // - alpha * f1(x) * ... * fk(x)]
+    q_x.add_mle_list(fxs.to_vec(), -*alpha)?;
 
     let iop_proof = <PolyIOP<F> as ZeroCheck<F>>::prove(&q_x, transcript)?;
 
     end_timer!(start);
     Ok((iop_proof, q_x))
 }
-
-/// Helper function of the IOP.
-///
-/// Input:
-/// - prod(x)
-///
-/// Output: the following 4 polynomials
-/// - prod(0, x)
-/// - prod(1, x)
-/// - prod(x, 0)
-/// - prod(x, 1)
-fn build_prod_partial_eval<F: PrimeField>(
-    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 = Rc::new(DenseMultilinearExtension::from_evaluations_slice(
-        num_vars,
-        &prod_x_eval[0..1 << num_vars],
-    ));
-    // prod(1, x)
-    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)
-    // ===================================
-    //
-    // now we compute eval_x0 and eval_x1
-    // eval_0x will be the odd coefficients of eval
-    // and eval_1x will be the even coefficients of eval
-    let mut eval_x0 = vec![];
-    let mut eval_x1 = vec![];
-    for (x, &prod_x) in prod_x_eval.iter().enumerate() {
-        if x & 1 == 0 {
-            eval_x0.push(prod_x);
-        } else {
-            eval_x1.push(prod_x);
-        }
-    }
-    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);
-
-    Ok([prod_0_x, prod_1_x, prod_x_0, prod_x_1])
-}
-
-/// Returns the evaluations of
-/// - `prod(0,x) := prod(0, x1, …, xn)` which is the MLE over the
-/// evaluations of f(x)/g(x) on the boolean hypercube {0,1}^n:
-///
-/// The caller needs to check num_vars matches in f/g
-/// Cost: linear in N.
-fn compute_prod_0<F: PrimeField>(
-    fx: &DenseMultilinearExtension<F>,
-    gx: &DenseMultilinearExtension<F>,
-) -> Result<Vec<F>, PolyIOPErrors> {
-    let start = start_timer!(|| "compute prod(0,x)");
-
-    let input = fx
-        .iter()
-        .zip(gx.iter())
-        .map(|(&fi, &gi)| (fi, gi))
-        .collect::<Vec<_>>();
-    let prod_0x_evals = input.par_iter().map(|(x, y)| *x / *y).collect::<Vec<_>>();
-
-    end_timer!(start);
-    Ok(prod_0x_evals)
-}