impl zero check (#15) and vp operations (#17)

poly-iop/Cargo.toml

@@ -20,6 +20,7 @@ displaydoc = { version = "0.2.3", default-features = false }
 rayon = { version = "1.5.2", default-features = false, optional = true }
 
 [features]
+# default = [ "parallel", "print-trace" ]
 default = [ "parallel" ]
 parallel = [ 
     "rayon",

poly-iop/src/lib.rs

@@ -1,8 +1,5 @@
-#![allow(dead_code)]
-
-use std::marker::PhantomData;
-
 use ark_ff::PrimeField;
+use std::marker::PhantomData;
 
 mod errors;
 mod structs;
@@ -10,14 +7,14 @@ mod sum_check;
 mod transcript;
 mod utils;
 mod virtual_poly;
-// mod zero_check;
+mod zero_check;
 
 pub use virtual_poly::VirtualPolynomial;
 
 /// Struct for PolyIOP protocol.
 /// It is instantiated with
 /// - SumCheck protocol.
-/// - ZeroCheck protocol. (WIP)
+/// - ZeroCheck protocol.
 pub struct PolyIOP<F: PrimeField> {
     phantom: PhantomData<F>,
 }

poly-iop/src/structs.rs

@@ -1,5 +1,6 @@
 //! Structs for polynomials and extensions.
 
+use crate::VirtualPolynomial;
 use ark_ff::PrimeField;
 use std::marker::PhantomData;
 
@@ -38,3 +39,26 @@ pub struct IOPProof {
 pub struct IOPProverMessage<F: PrimeField> {
     pub(crate) evaluations: Vec<F>,
 }
+
+/// Prover State of a PolyIOP
+pub struct IOPProverState<F: PrimeField> {
+    /// sampled randomness given by the verifier
+    pub challenges: Vec<F>,
+    /// the current round number
+    pub(crate) round: usize,
+    /// pointer to the virtual polynomial
+    pub(crate) poly: VirtualPolynomial<F>,
+}
+
+/// Prover State of a PolyIOP
+pub struct IOPVerifierState<F: PrimeField> {
+    pub(crate) round: usize,
+    pub(crate) num_vars: usize,
+    pub(crate) max_degree: usize,
+    pub(crate) finished: bool,
+    /// a list storing the univariate polynomial in evaluation form sent by the
+    /// prover at each round
+    pub(crate) polynomials_received: Vec<Vec<F>>,
+    /// a list storing the randomness sampled by the verifier at each round
+    pub(crate) challenges: Vec<F>,
+}

poly-iop/src/sum_check/mod.rs

@@ -4,7 +4,7 @@
 
 use crate::{
     errors::PolyIOPErrors,
-    structs::{DomainInfo, IOPProof, SubClaim},
+    structs::{DomainInfo, IOPProof, IOPProverState, IOPVerifierState, SubClaim},
     transcript::IOPTranscript,
     virtual_poly::VirtualPolynomial,
     PolyIOP,
@@ -15,9 +15,7 @@ use ark_std::{end_timer, start_timer};
 mod prover;
 mod verifier;
 
-pub use prover::ProverState;
-pub use verifier::VerifierState;
-
+/// Trait for doing sum check protocols.
 pub trait SumCheck<F: PrimeField> {
     type Proof;
     type PolyList;
@@ -36,27 +34,15 @@ pub trait SumCheck {
     /// SumCheck prover/verifier.
     fn init_transcript() -> Self::Transcript;
 
-    /// generate proof of the sum of polynomial over {0,1}^`num_vars`
-    ///
-    /// The polynomial is represented by a list of products of polynomials along
-    /// with its coefficient that is meant to be added together.
-    ///
-    /// This data structure of the polynomial is a list of list of
-    /// `(coefficient, DenseMultilinearExtension)`.
-    /// * Number of products n = `polynomial.products.len()`,
-    /// * Number of multiplicands of ith product m_i =
-    ///   `polynomial.products[i].1.len()`,
-    /// * Coefficient of ith product c_i = `polynomial.products[i].0`
+    /// Generate proof of the sum of polynomial over {0,1}^`num_vars`
     ///
-    /// The resulting polynomial is
-    ///
-    /// $$\sum_{i=0}^{n}C_i\cdot\prod_{j=0}^{m_i}P_{ij}$$
+    /// The polynomial is represented in the form of a VirtualPolynomial.
     fn prove(
         poly: &Self::PolyList,
         transcript: &mut Self::Transcript,
     ) -> Result<Self::Proof, PolyIOPErrors>;
 
-    /// verify the claimed sum using the proof
+    /// Verify the claimed sum using the proof
     fn verify(
         sum: F,
         proof: &Self::Proof,
@@ -65,6 +51,7 @@ pub trait SumCheck {
     ) -> Result<Self::SubClaim, PolyIOPErrors>;
 }
 
+/// Trait for sum check protocol prover side APIs.
 pub trait SumCheckProver<F: PrimeField>
 where
     Self: Sized,
@@ -72,22 +59,10 @@ where
     type PolyList;
     type ProverMessage;
 
-    /// initialize the prover to argue for the sum of polynomial over
+    /// Initialize the prover to argue for the sum of polynomial over
     /// {0,1}^`num_vars`
     ///
-    /// The polynomial is represented by a list of products of polynomials along
-    /// with its coefficient that is meant to be added together.
-    ///
-    /// This data structure of the polynomial is a list of list of
-    /// `(coefficient, DenseMultilinearExtension)`.
-    /// * Number of products n = `polynomial.products.len()`,
-    /// * Number of multiplicands of ith product m_i =
-    ///   `polynomial.products[i].1.len()`,
-    /// * Coefficient of ith product c_i = `polynomial.products[i].0`
-    ///
-    /// The resulting polynomial is
-    ///
-    /// $$\sum_{i=0}^{n}C_i\cdot\prod_{j=0}^{m_i}P_{ij}$$
+    /// The polynomial is represented in the form of a VirtualPolynomial.
     fn prover_init(polynomial: &Self::PolyList) -> Result<Self, PolyIOPErrors>;
 
     /// receive message from verifier, generate prover message, and proceed to
@@ -100,6 +75,7 @@ where
     ) -> Result<Self::ProverMessage, PolyIOPErrors>;
 }
 
+/// Trait for sum check protocol verifier side APIs.
 pub trait SumCheckVerifier<F: PrimeField> {
     type DomainInfo;
     type ProverMessage;
@@ -165,21 +141,9 @@ impl SumCheck for PolyIOP {
         res
     }
 
-    /// generate proof of the sum of polynomial over {0,1}^`num_vars`
-    ///
-    /// The polynomial is represented by a list of products of polynomials along
-    /// with its coefficient that is meant to be added together.
+    /// Generate proof of the sum of polynomial over {0,1}^`num_vars`
     ///
-    /// This data structure of the polynomial is a list of list of
-    /// `(coefficient, DenseMultilinearExtension)`.
-    /// * Number of products n = `polynomial.products.len()`,
-    /// * Number of multiplicands of ith product m_i =
-    ///   `polynomial.products[i].1.len()`,
-    /// * Coefficient of ith product c_i = `polynomial.products[i].0`
-    ///
-    /// The resulting polynomial is
-    ///
-    /// $$\sum_{i=0}^{n}C_i\cdot\prod_{j=0}^{m_i}P_{ij}$$
+    /// The polynomial is represented in the form of a VirtualPolynomial.
     fn prove(
         poly: &Self::PolyList,
         transcript: &mut Self::Transcript,
@@ -188,12 +152,12 @@ impl SumCheck for PolyIOP {
 
         transcript.append_domain_info(&poly.domain_info)?;
 
-        let mut prover_state = ProverState::prover_init(poly)?;
+        let mut prover_state = IOPProverState::prover_init(poly)?;
         let mut challenge = None;
         let mut prover_msgs = Vec::with_capacity(poly.domain_info.num_variables);
         for _ in 0..poly.domain_info.num_variables {
             let prover_msg =
-                ProverState::prove_round_and_update_state(&mut prover_state, &challenge)?;
+                IOPProverState::prove_round_and_update_state(&mut prover_state, &challenge)?;
             transcript.append_prover_message(&prover_msg)?;
             prover_msgs.push(prover_msg);
             challenge = Some(transcript.get_and_append_challenge(b"Internal round")?);
@@ -215,18 +179,18 @@ impl SumCheck for PolyIOP {
         let start = start_timer!(|| "sum check verify");
 
         transcript.append_domain_info(domain_info)?;
-        let mut verifier_state = VerifierState::verifier_init(domain_info);
+        let mut verifier_state = IOPVerifierState::verifier_init(domain_info);
         for i in 0..domain_info.num_variables {
             let prover_msg = proof.proofs.get(i).expect("proof is incomplete");
             transcript.append_prover_message(prover_msg)?;
-            VerifierState::verify_round_and_update_state(
+            IOPVerifierState::verify_round_and_update_state(
                 &mut verifier_state,
                 prover_msg,
                 transcript,
             )?;
         }
 
-        let res = VerifierState::check_and_generate_subclaim(&verifier_state, &claimed_sum);
+        let res = IOPVerifierState::check_and_generate_subclaim(&verifier_state, &claimed_sum);
 
         end_timer!(start);
         res
@@ -237,41 +201,49 @@ impl SumCheck for PolyIOP {
 mod test {
 
     use super::*;
-    use crate::virtual_poly::test::random_list_of_products;
     use ark_bls12_381::Fr;
     use ark_ff::UniformRand;
     use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
     use ark_std::test_rng;
     use std::rc::Rc;
 
-    fn test_sumcheck(nv: usize, num_multiplicands_range: (usize, usize), num_products: usize) {
+    fn test_sumcheck(
+        nv: usize,
+        num_multiplicands_range: (usize, usize),
+        num_products: usize,
+    ) -> Result<(), PolyIOPErrors> {
         let mut rng = test_rng();
-        let mut transcript = PolyIOP::init_transcript();
+        let mut transcript = <PolyIOP<Fr> as SumCheck<Fr>>::init_transcript();
 
         let (poly, asserted_sum) =
-            random_list_of_products::<Fr, _>(nv, num_multiplicands_range, num_products, &mut rng);
-        let proof = PolyIOP::prove(&poly, &mut transcript).expect("fail to prove");
+            VirtualPolynomial::rand(nv, num_multiplicands_range, num_products, &mut rng)?;
+        let proof = <PolyIOP<Fr> as SumCheck<Fr>>::prove(&poly, &mut transcript)?;
         let poly_info = poly.domain_info.clone();
-        let mut transcript = PolyIOP::init_transcript();
-        let subclaim = PolyIOP::verify(asserted_sum, &proof, &poly_info, &mut transcript)
-            .expect("fail to verify");
+        let mut transcript = <PolyIOP<Fr> as SumCheck<Fr>>::init_transcript();
+        let subclaim = <PolyIOP<Fr> as SumCheck<Fr>>::verify(
+            asserted_sum,
+            &proof,
+            &poly_info,
+            &mut transcript,
+        )?;
         assert!(
             poly.evaluate(&subclaim.point).unwrap() == subclaim.expected_evaluation,
             "wrong subclaim"
         );
+        Ok(())
     }
 
     fn test_sumcheck_internal(
         nv: usize,
         num_multiplicands_range: (usize, usize),
         num_products: usize,
-    ) {
+    ) -> Result<(), PolyIOPErrors> {
         let mut rng = test_rng();
         let (poly, asserted_sum) =
-            random_list_of_products::<Fr, _>(nv, num_multiplicands_range, num_products, &mut rng);
+            VirtualPolynomial::<Fr>::rand(nv, num_multiplicands_range, num_products, &mut rng)?;
         let poly_info = poly.domain_info.clone();
-        let mut prover_state = ProverState::prover_init(&poly).unwrap();
-        let mut verifier_state = VerifierState::verifier_init(&poly_info);
+        let mut prover_state = IOPProverState::prover_init(&poly)?;
+        let mut verifier_state = IOPVerifierState::verifier_init(&poly_info);
         let mut challenge = None;
         let mut transcript = IOPTranscript::new(b"a test transcript");
         transcript
@@ -279,10 +251,11 @@ mod test {
             .unwrap();
         for _ in 0..poly.domain_info.num_variables {
             let prover_message =
-                ProverState::prove_round_and_update_state(&mut prover_state, &challenge).unwrap();
+                IOPProverState::prove_round_and_update_state(&mut prover_state, &challenge)
+                    .unwrap();
 
             challenge = Some(
-                VerifierState::verify_round_and_update_state(
+                IOPVerifierState::verify_round_and_update_state(
                     &mut verifier_state,
                     &prover_message,
                     &mut transcript,
@@ -290,115 +263,120 @@ mod test {
                 .unwrap(),
             );
         }
-        let subclaim = VerifierState::check_and_generate_subclaim(&verifier_state, &asserted_sum)
-            .expect("fail to generate subclaim");
+        let subclaim =
+            IOPVerifierState::check_and_generate_subclaim(&verifier_state, &asserted_sum)
+                .expect("fail to generate subclaim");
         assert!(
             poly.evaluate(&subclaim.point).unwrap() == subclaim.expected_evaluation,
             "wrong subclaim"
         );
+        Ok(())
     }
 
     #[test]
-    fn test_trivial_polynomial() {
+    fn test_trivial_polynomial() -> Result<(), PolyIOPErrors> {
         let nv = 1;
         let num_multiplicands_range = (4, 13);
         let num_products = 5;
 
-        test_sumcheck(nv, num_multiplicands_range, num_products);
-        test_sumcheck_internal(nv, num_multiplicands_range, num_products);
+        test_sumcheck(nv, num_multiplicands_range, num_products)?;
+        test_sumcheck_internal(nv, num_multiplicands_range, num_products)
     }
     #[test]
-    fn test_normal_polynomial() {
+    fn test_normal_polynomial() -> Result<(), PolyIOPErrors> {
         let nv = 12;
         let num_multiplicands_range = (4, 9);
         let num_products = 5;
 
-        test_sumcheck(nv, num_multiplicands_range, num_products);
-        test_sumcheck_internal(nv, num_multiplicands_range, num_products);
+        test_sumcheck(nv, num_multiplicands_range, num_products)?;
+        test_sumcheck_internal(nv, num_multiplicands_range, num_products)
     }
     #[test]
-    #[should_panic]
     fn zero_polynomial_should_error() {
         let nv = 0;
         let num_multiplicands_range = (4, 13);
         let num_products = 5;
 
-        test_sumcheck(nv, num_multiplicands_range, num_products);
-        test_sumcheck_internal(nv, num_multiplicands_range, num_products);
+        assert!(test_sumcheck(nv, num_multiplicands_range, num_products).is_err());
+        assert!(test_sumcheck_internal(nv, num_multiplicands_range, num_products).is_err());
     }
 
     #[test]
-    fn test_extract_sum() {
+    fn test_extract_sum() -> Result<(), PolyIOPErrors> {
         let mut rng = test_rng();
-        let mut transcript = PolyIOP::init_transcript();
-        let (poly, asserted_sum) = random_list_of_products::<Fr, _>(8, (3, 4), 3, &mut rng);
+        let mut transcript = <PolyIOP<Fr> as SumCheck<Fr>>::init_transcript();
+        let (poly, asserted_sum) = VirtualPolynomial::<Fr>::rand(8, (3, 4), 3, &mut rng)?;
 
-        let proof = PolyIOP::prove(&poly, &mut transcript).expect("fail to prove");
-        assert_eq!(PolyIOP::extract_sum(&proof), asserted_sum);
+        let proof = <PolyIOP<Fr> as SumCheck<Fr>>::prove(&poly, &mut transcript)?;
+        assert_eq!(
+            <PolyIOP<Fr> as SumCheck<Fr>>::extract_sum(&proof),
+            asserted_sum
+        );
+        Ok(())
     }
 
     #[test]
     /// Test that the memory usage of shared-reference is linear to number of
     /// unique MLExtensions instead of total number of multiplicands.
-    fn test_shared_reference() {
+    fn test_shared_reference() -> Result<(), PolyIOPErrors> {
         let mut rng = test_rng();
         let ml_extensions: Vec<_> = (0..5)
             .map(|_| Rc::new(DenseMultilinearExtension::<Fr>::rand(8, &mut rng)))
             .collect();
         let mut poly = VirtualPolynomial::new(8);
-        poly.add_product(
+        poly.add_mle_list(
             vec![
                 ml_extensions[2].clone(),
                 ml_extensions[3].clone(),
                 ml_extensions[0].clone(),
             ],
             Fr::rand(&mut rng),
-        )
-        .unwrap();
-        poly.add_product(
+        )?;
+        poly.add_mle_list(
             vec![
                 ml_extensions[1].clone(),
                 ml_extensions[4].clone(),
                 ml_extensions[4].clone(),
             ],
             Fr::rand(&mut rng),
-        )
-        .unwrap();
-        poly.add_product(
+        )?;
+        poly.add_mle_list(
             vec![
                 ml_extensions[3].clone(),
                 ml_extensions[2].clone(),
                 ml_extensions[1].clone(),
             ],
             Fr::rand(&mut rng),
-        )
-        .unwrap();
-        poly.add_product(
+        )?;
+        poly.add_mle_list(
             vec![ml_extensions[0].clone(), ml_extensions[0].clone()],
             Fr::rand(&mut rng),
-        )
-        .unwrap();
-        poly.add_product(vec![ml_extensions[4].clone()], Fr::rand(&mut rng))
-            .unwrap();
+        )?;
+        poly.add_mle_list(vec![ml_extensions[4].clone()], Fr::rand(&mut rng))?;
 
         assert_eq!(poly.flattened_ml_extensions.len(), 5);
 
         // test memory usage for prover
-        let prover = ProverState::prover_init(&poly).unwrap();
+        let prover = IOPProverState::prover_init(&poly).unwrap();
         assert_eq!(prover.poly.flattened_ml_extensions.len(), 5);
         drop(prover);
 
-        let mut transcript = PolyIOP::init_transcript();
+        let mut transcript = <PolyIOP<Fr> as SumCheck<Fr>>::init_transcript();
         let poly_info = poly.domain_info.clone();
-        let proof = PolyIOP::prove(&poly, &mut transcript).expect("fail to prove");
-        let asserted_sum = PolyIOP::extract_sum(&proof);
-
-        let mut transcript = PolyIOP::init_transcript();
-        let subclaim = PolyIOP::verify(asserted_sum, &proof, &poly_info, &mut transcript)
-            .expect("fail to verify");
+        let proof = <PolyIOP<Fr> as SumCheck<Fr>>::prove(&poly, &mut transcript)?;
+        let asserted_sum = <PolyIOP<Fr> as SumCheck<Fr>>::extract_sum(&proof);
+
+        let mut transcript = <PolyIOP<Fr> as SumCheck<Fr>>::init_transcript();
+        let subclaim = <PolyIOP<Fr> as SumCheck<Fr>>::verify(
+            asserted_sum,
+            &proof,
+            &poly_info,
+            &mut transcript,
+        )?;
         assert!(
-            poly.evaluate(&subclaim.point).unwrap() == subclaim.expected_evaluation,
+            poly.evaluate(&subclaim.point)? == subclaim.expected_evaluation,
             "wrong subclaim"
         );
+        Ok(())
     }
 }

poly-iop/src/sum_check/prover.rs

@@ -1,49 +1,27 @@
 //! Prover
-use std::rc::Rc;
 
-// TODO: some of the struct is generic for Sum Checks and Zero Checks.
-// If so move them to src/structs.rs
 use super::SumCheckProver;
-use crate::{errors::PolyIOPErrors, structs::IOPProverMessage, virtual_poly::VirtualPolynomial};
+use crate::{
+    errors::PolyIOPErrors,
+    structs::{IOPProverMessage, IOPProverState},
+    virtual_poly::VirtualPolynomial,
+};
 use ark_ff::PrimeField;
 use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
 use ark_std::{end_timer, start_timer, vec::Vec};
+use std::rc::Rc;
 
 #[cfg(feature = "parallel")]
 use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
 
-/// Prover State
-pub struct ProverState<F: PrimeField> {
-    /// sampled randomness given by the verifier
-    pub challenges: Vec<F>,
-    /// the current round number
-    pub(crate) round: usize,
-    /// pointer to the virtual polynomial
-    pub(crate) poly: VirtualPolynomial<F>,
-}
-
-impl<F: PrimeField> SumCheckProver<F> for ProverState<F> {
+impl<F: PrimeField> SumCheckProver<F> for IOPProverState<F> {
     type PolyList = VirtualPolynomial<F>;
     type ProverMessage = IOPProverMessage<F>;
 
-    /// initialize the prover to argue for the sum of polynomial over
+    /// Initialize the prover to argue for the sum of polynomial over
     /// {0,1}^`num_vars`
-    ///
-    /// The polynomial is represented by a list of products of polynomials along
-    /// with its coefficient that is meant to be added together.
-    ///
-    /// This data structure of the polynomial is a list of list of
-    /// `(coefficient, DenseMultilinearExtension)`.
-    /// * Number of products n = `polynomial.products.len()`,
-    /// * Number of multiplicands of ith product m_i =
-    ///   `polynomial.products[i].1.len()`,
-    /// * Coefficient of ith product c_i = `polynomial.products[i].0`
-    ///
-    /// The resulting polynomial is
-    ///
-    /// $$\sum_{i=0}^{n}C_i\cdot\prod_{j=0}^{m_i}P_{ij}$$
     fn prover_init(polynomial: &Self::PolyList) -> Result<Self, PolyIOPErrors> {
-        let start = start_timer!(|| "prover init");
+        let start = start_timer!(|| "sum check prover init");
         if polynomial.domain_info.num_variables == 0 {
             return Err(PolyIOPErrors::InvalidParameters(
                 "Attempt to prove a constant.".to_string(),
@@ -51,14 +29,14 @@ impl SumCheckProver for ProverState {
         }
         end_timer!(start);
 
-        Ok(ProverState {
+        Ok(Self {
             challenges: Vec::with_capacity(polynomial.domain_info.num_variables),
             round: 0,
             poly: polynomial.clone(),
         })
     }
 
-    /// receive message from verifier, generate prover message, and proceed to
+    /// Receive message from verifier, generate prover message, and proceed to
     /// next round
     ///
     /// Main algorithm used is from section 3.2 of [XZZPS19](https://eprint.iacr.org/2019/317.pdf#subsection.3.2).
@@ -66,7 +44,8 @@ impl SumCheckProver for ProverState {
         &mut self,
         challenge: &Option<F>,
     ) -> Result<Self::ProverMessage, PolyIOPErrors> {
-        let start = start_timer!(|| format!("prove {}-th round and update state", self.round));
+        let start =
+            start_timer!(|| format!("sum check prove {}-th round and update state", self.round));
 
         let fix_argument = start_timer!(|| "fix argument");
 

poly-iop/src/sum_check/verifier.rs

@@ -4,7 +4,7 @@
 use super::SumCheckVerifier;
 use crate::{
     errors::PolyIOPErrors,
-    structs::{DomainInfo, IOPProverMessage, SubClaim},
+    structs::{DomainInfo, IOPProverMessage, IOPVerifierState, SubClaim},
     transcript::IOPTranscript,
 };
 use ark_ff::PrimeField;
@@ -13,20 +13,7 @@ use ark_std::{end_timer, start_timer};
 #[cfg(feature = "parallel")]
 use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
 
-/// Verifier State
-pub struct VerifierState<F: PrimeField> {
-    round: usize,
-    num_vars: usize,
-    max_degree: usize,
-    finished: bool,
-    /// a list storing the univariate polynomial in evaluation form sent by the
-    /// prover at each round
-    polynomials_received: Vec<Vec<F>>,
-    /// a list storing the randomness sampled by the verifier at each round
-    challenges: Vec<F>,
-}
-
-impl<F: PrimeField> SumCheckVerifier<F> for VerifierState<F> {
+impl<F: PrimeField> SumCheckVerifier<F> for IOPVerifierState<F> {
     type DomainInfo = DomainInfo<F>;
     type ProverMessage = IOPProverMessage<F>;
     type Challenge = F;
@@ -35,8 +22,8 @@ impl SumCheckVerifier for VerifierState {
 
     /// initialize the verifier
     fn verifier_init(index_info: &Self::DomainInfo) -> Self {
-        let start = start_timer!(|| "verifier init");
-        let res = VerifierState {
+        let start = start_timer!(|| "sum check verifier init");
+        let res = Self {
             round: 1,
             num_vars: index_info.num_variables,
             max_degree: index_info.max_degree,
@@ -59,7 +46,8 @@ impl SumCheckVerifier for VerifierState {
         prover_msg: &Self::ProverMessage,
         transcript: &mut Self::Transcript,
     ) -> Result<Self::Challenge, PolyIOPErrors> {
-        let start = start_timer!(|| format!("verify {}-th round and update state", self.round));
+        let start =
+            start_timer!(|| format!("sum check verify {}-th round and update state", self.round));
 
         if self.finished {
             return Err(PolyIOPErrors::InvalidVerifier(
@@ -101,7 +89,7 @@ impl SumCheckVerifier for VerifierState {
         &self,
         asserted_sum: &F,
     ) -> Result<Self::SubClaim, PolyIOPErrors> {
-        let start = start_timer!(|| "check_and_generate_subclaim");
+        let start = start_timer!(|| "sum check check and generate subclaim");
         if !self.finished {
             return Err(PolyIOPErrors::InvalidVerifier(
                 "Incorrect verifier state: Verifier has not finished.".to_string(),
@@ -173,10 +161,10 @@ impl SumCheckVerifier for VerifierState {
     }
 }
 
-/// interpolate a uni-variate degree-`p_i.len()-1` polynomial and evaluate this
+/// Interpolate a uni-variate degree-`p_i.len()-1` polynomial and evaluate this
 /// polynomial at `eval_at`.
 pub(crate) fn interpolate_uni_poly<F: PrimeField>(p_i: &[F], eval_at: F) -> F {
-    let start = start_timer!(|| "interpolate_uni_poly");
+    let start = start_timer!(|| "sum check interpolate uni poly");
     let mut result = F::zero();
     let mut i = F::zero();
     for term in p_i.iter() {

poly-iop/src/virtual_poly.rs

@@ -1,10 +1,40 @@
 use crate::{errors::PolyIOPErrors, structs::DomainInfo};
 use ark_ff::PrimeField;
 use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
-use ark_std::{end_timer, start_timer};
-use std::{cmp::max, collections::HashMap, marker::PhantomData, rc::Rc};
+use ark_std::{
+    end_timer,
+    rand::{Rng, RngCore},
+    start_timer,
+};
+use std::{cmp::max, collections::HashMap, marker::PhantomData, ops::Add, rc::Rc};
 
-/// A virtual polynomial is a list of multilinear polynomials
+#[rustfmt::skip]
+/// A virtual polynomial is a sum of products of multilinear polynomials;
+/// where the multilinear polynomials are stored via their multilinear
+/// extensions:  `(coefficient, DenseMultilinearExtension)`
+///
+/// * Number of products n = `polynomial.products.len()`,
+/// * Number of multiplicands of ith product m_i =
+///   `polynomial.products[i].1.len()`,
+/// * Coefficient of ith product c_i = `polynomial.products[i].0`
+///
+/// The resulting polynomial is
+///
+/// $$\sum_{i=0}^{n} c_i \cdot \prod_{j=0}^{m_i} P_{ij} $$
+///
+/// Example:
+///  f = c0 * f0 * f1 * f2 + c1 * f3 * f4
+/// where f0 ... f4 are multilinear polynomials
+///
+/// - flattened_ml_extensions stores the multilinear extension representation of
+///   f0, f1, f2, f3 and f4
+/// - products is 
+///     \[ 
+///         (c0, \[0, 1, 2\]), 
+///         (c1, \[3, 4\]) 
+///     \]
+/// - raw_pointers_lookup_table maps fi to i
+///
 #[derive(Clone, Debug, Default, PartialEq)]
 pub struct VirtualPolynomial<F: PrimeField> {
     /// Aux information about the multilinear polynomial
@@ -18,6 +48,26 @@ pub struct VirtualPolynomial {
     raw_pointers_lookup_table: HashMap<*const DenseMultilinearExtension<F>, usize>,
 }
 
+impl<F: PrimeField> Add for &VirtualPolynomial<F> {
+    type Output = VirtualPolynomial<F>;
+    fn add(self, other: &VirtualPolynomial<F>) -> Self::Output {
+        let start = start_timer!(|| "virtual poly add");
+        let mut res = self.clone();
+        for products in other.products.iter() {
+            let cur: Vec<Rc<DenseMultilinearExtension<F>>> = products
+                .1
+                .iter()
+                .map(|&x| other.flattened_ml_extensions[x].clone())
+                .collect();
+
+            res.add_mle_list(cur, products.0)
+                .expect("add product failed");
+        }
+        end_timer!(start);
+        res
+    }
+}
+
 impl<F: PrimeField> VirtualPolynomial<F> {
     /// Returns an empty polynomial
     pub fn new(num_variables: usize) -> Self {
@@ -33,33 +83,52 @@ impl VirtualPolynomial {
         }
     }
 
+    /// Returns an new virtual polynomial from a 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);
+
+        VirtualPolynomial {
+            domain_info: DomainInfo {
+                // The max degree is the max degree of any individual variable
+                max_degree: 1,
+                num_variables: mle.num_vars,
+                phantom: PhantomData::default(),
+            },
+            products: vec![(coefficient, vec![0])],
+            flattened_ml_extensions: vec![mle],
+            raw_pointers_lookup_table: hm,
+        }
+    }
+
     /// Add a list of multilinear extensions that is meant to be multiplied
     /// together. The resulting polynomial will be multiplied by the scalar
     /// `coefficient`.
-    pub fn add_product(
+    pub fn add_mle_list(
         &mut self,
-        product: impl IntoIterator<Item = Rc<DenseMultilinearExtension<F>>>,
+        mle_list: impl IntoIterator<Item = Rc<DenseMultilinearExtension<F>>>,
         coefficient: F,
     ) -> Result<(), PolyIOPErrors> {
-        let product: Vec<Rc<DenseMultilinearExtension<F>>> = product.into_iter().collect();
-        let mut indexed_product = Vec::with_capacity(product.len());
-        assert!(!product.is_empty());
-        self.domain_info.max_degree = max(self.domain_info.max_degree, product.len());
-        for m in product {
-            if m.num_vars != self.domain_info.num_variables {
+        let mle_list: Vec<Rc<DenseMultilinearExtension<F>>> = mle_list.into_iter().collect();
+        let mut indexed_product = Vec::with_capacity(mle_list.len());
+        assert!(!mle_list.is_empty());
+        self.domain_info.max_degree = max(self.domain_info.max_degree, mle_list.len());
+        for mle in mle_list {
+            if mle.num_vars != self.domain_info.num_variables {
                 return Err(PolyIOPErrors::InvalidParameters(format!(
                     "product has a multiplicand with wrong number of variables {} vs {}",
-                    m.num_vars, self.domain_info.num_variables
+                    mle.num_vars, self.domain_info.num_variables
                 )));
             }
 
-            let m_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(&m);
-            if let Some(index) = self.raw_pointers_lookup_table.get(&m_ptr) {
+            let mle_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(&mle);
+            if let Some(index) = self.raw_pointers_lookup_table.get(&mle_ptr) {
                 indexed_product.push(*index)
             } else {
                 let curr_index = self.flattened_ml_extensions.len();
-                self.flattened_ml_extensions.push(m.clone());
-                self.raw_pointers_lookup_table.insert(m_ptr, curr_index);
+                self.flattened_ml_extensions.push(mle.clone());
+                self.raw_pointers_lookup_table.insert(mle_ptr, curr_index);
                 indexed_product.push(curr_index);
             }
         }
@@ -67,9 +136,39 @@ impl VirtualPolynomial {
         Ok(())
     }
 
+    /// Multiple the current VirtualPolynomial by an MLE:
+    /// - add the MLE to the MLE list
+    /// - multiple each product by MLE and its coefficient
+    pub fn mul_by_mle(
+        &mut self,
+        mle: Rc<DenseMultilinearExtension<F>>,
+        coefficient: F,
+    ) -> Result<(), PolyIOPErrors> {
+        let start = start_timer!(|| "mul by mle");
+        let mle_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(&mle);
+
+        let mle_index = match self.raw_pointers_lookup_table.get(&mle_ptr) {
+            Some(&p) => p,
+            None => {
+                self.raw_pointers_lookup_table
+                    .insert(mle_ptr, self.flattened_ml_extensions.len());
+                self.flattened_ml_extensions.push(mle);
+                self.flattened_ml_extensions.len() - 1
+            },
+        };
+
+        for (prod_coef, indices) in self.products.iter_mut() {
+            indices.push(mle_index);
+            *prod_coef *= coefficient;
+        }
+        self.domain_info.max_degree += 1;
+        end_timer!(start);
+        Ok(())
+    }
+
     /// Evaluate the polynomial at point `point`
     pub fn evaluate(&self, point: &[F]) -> Result<F, PolyIOPErrors> {
-        let start = start_timer!(|| "begin evaluation");
+        let start = start_timer!(|| "evaluation");
 
         if self.domain_info.num_variables != point.len() {
             return Err(PolyIOPErrors::InvalidParameters(format!(
@@ -98,114 +197,171 @@ impl VirtualPolynomial {
         end_timer!(start);
         Ok(res)
     }
-}
 
-#[cfg(test)]
-pub(crate) mod test {
-    use super::*;
-    use ark_std::rand::{Rng, RngCore};
-
-    pub fn random_product<F: PrimeField, R: RngCore>(
+    /// Sample a random virtual polynomial, return the polynomial and its sum.
+    pub fn rand<R: RngCore>(
         nv: usize,
-        num_multiplicands: usize,
+        num_multiplicands_range: (usize, usize),
+        num_products: usize,
         rng: &mut R,
-    ) -> (Vec<Rc<DenseMultilinearExtension<F>>>, F) {
-        let mut multiplicands = Vec::with_capacity(num_multiplicands);
-        for _ in 0..num_multiplicands {
-            multiplicands.push(Vec::with_capacity(1 << nv))
-        }
-        let mut sum = F::zero();
+    ) -> Result<(Self, F), PolyIOPErrors> {
+        let start = start_timer!("sample random virtual polynomial");
 
-        for _ in 0..(1 << nv) {
-            let mut product = F::one();
-            for i in 0..num_multiplicands {
-                let val = F::rand(rng);
-                multiplicands[i].push(val);
-                product *= val;
-            }
-            sum += product;
+        let mut sum = F::zero();
+        let mut poly = VirtualPolynomial::new(nv);
+        for _ in 0..num_products {
+            let num_multiplicands =
+                rng.gen_range(num_multiplicands_range.0..num_multiplicands_range.1);
+            let (product, product_sum) = random_mle_list(nv, num_multiplicands, rng);
+            let coefficient = F::rand(rng);
+            poly.add_mle_list(product.into_iter(), coefficient)?;
+            sum += product_sum * coefficient;
         }
 
-        (
-            multiplicands
-                .into_iter()
-                .map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
-                .collect(),
-            sum,
-        )
+        end_timer!(start);
+        Ok((poly, sum))
     }
 
-    pub(crate) fn random_list_of_products<F: PrimeField, R: RngCore>(
+    /// Sample a random virtual polynomial that evaluates to zero everywhere on
+    /// the boolean hypercube.
+    pub fn rand_zero<R: RngCore>(
         nv: usize,
         num_multiplicands_range: (usize, usize),
         num_products: usize,
         rng: &mut R,
-    ) -> (VirtualPolynomial<F>, F) {
-        let mut sum = F::zero();
+    ) -> Result<Self, PolyIOPErrors> {
         let mut poly = VirtualPolynomial::new(nv);
         for _ in 0..num_products {
             let num_multiplicands =
                 rng.gen_range(num_multiplicands_range.0..num_multiplicands_range.1);
-            let (product, product_sum) = random_product(nv, num_multiplicands, rng);
+            let product = random_zero_mle_list(nv, num_multiplicands, rng);
             let coefficient = F::rand(rng);
-            poly.add_product(product.into_iter(), coefficient).unwrap();
-            sum += product_sum * coefficient;
+            poly.add_mle_list(product.into_iter(), coefficient).unwrap();
+        }
+
+        Ok(poly)
+    }
+}
+
+/// Sample a random list of multilinear polynomials.
+/// Returns
+/// - the list of polynomials,
+/// - its sum of polynomial evaluations over the boolean hypercube.
+fn random_mle_list<F: PrimeField, R: RngCore>(
+    nv: usize,
+    degree: usize,
+    rng: &mut R,
+) -> (Vec<Rc<DenseMultilinearExtension<F>>>, F) {
+    let start = start_timer!("sample random mle list");
+    let mut multiplicands = Vec::with_capacity(degree);
+    for _ in 0..degree {
+        multiplicands.push(Vec::with_capacity(1 << nv))
+    }
+    let mut sum = F::zero();
+
+    for _ in 0..(1 << nv) {
+        let mut product = F::one();
+
+        for e in multiplicands.iter_mut() {
+            let val = F::rand(rng);
+            e.push(val);
+            product *= val;
+        }
+        sum += product;
+    }
+
+    let list = multiplicands
+        .into_iter()
+        .map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
+        .collect();
+
+    end_timer!(start);
+    (list, sum)
+}
+
+// Build a randomize list of mle-s whose sum is zero.
+pub fn random_zero_mle_list<F: PrimeField, R: RngCore>(
+    nv: usize,
+    degree: usize,
+    rng: &mut R,
+) -> Vec<Rc<DenseMultilinearExtension<F>>> {
+    let start = start_timer!("sample random zero mle list");
+
+    let mut multiplicands = Vec::with_capacity(degree);
+    for _ in 0..degree {
+        multiplicands.push(Vec::with_capacity(1 << nv))
+    }
+    for _ in 0..(1 << nv) {
+        multiplicands[0].push(F::zero());
+        for e in multiplicands.iter_mut().skip(1) {
+            e.push(F::rand(rng));
+        }
+    }
+
+    let list = multiplicands
+        .into_iter()
+        .map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
+        .collect();
+
+    end_timer!(start);
+    list
+}
+
+#[cfg(test)]
+pub(crate) mod test {
+    use super::*;
+    use ark_bls12_381::Fr;
+    use ark_ff::UniformRand;
+    use ark_std::test_rng;
+
+    #[test]
+    fn test_virtual_polynomial_additions() -> Result<(), PolyIOPErrors> {
+        let mut rng = test_rng();
+        for nv in 2..5 {
+            for num_products in 2..5 {
+                let base: Vec<Fr> = (0..nv).map(|_| Fr::rand(&mut rng)).collect();
+
+                let (a, _a_sum) =
+                    VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
+                let (b, _b_sum) =
+                    VirtualPolynomial::<Fr>::rand(nv, (2, 3), num_products, &mut rng)?;
+                let c = &a + &b;
+
+                assert_eq!(
+                    a.evaluate(base.as_ref())? + b.evaluate(base.as_ref())?,
+                    c.evaluate(base.as_ref())?
+                );
+            }
         }
 
-        (poly, sum)
+        Ok(())
     }
 
-    // pub fn random_zero_product<F: PrimeField, R: RngCore>(
-    //     nv: usize,
-    //     num_multiplicands: usize,
-    //     rng: &mut R,
-    // ) -> Vec<Rc<DenseMultilinearExtension<F>>> {
-    //     let degree = 2;
-    //     let mut multiplicands = Vec::with_capacity(degree);
-    //     for _ in 0..degree {
-    //         multiplicands.push(Vec::with_capacity(1 << nv))
-    //     }
-    //     let mut sum = F::zero();
-
-    //     for _ in 0..(1 << nv) {
-    //         let mut product = F::one();
-    //         for i in 0..degree {
-    //             let val = F::zero(); // F::rand(rng);
-    //             multiplicands[i].push(val);
-    //             product *= val;
-    //         }
-    //         sum += product;
-    //     }
-
-    //     // // last nv offsets the poly to 0
-    //     // for i in 0..num_multiplicands - 1 {
-    //     //     multiplicands[i].push(F::one());
-    //     // }
-    //     // multiplicands[num_multiplicands - 1].push(-sum);
-
-    //     multiplicands
-    //         .into_iter()
-    //         .map(|x|
-    // Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
-    //         .collect()
-    // }
-
-    // pub(crate) fn random_zero_list_of_products<F: PrimeField, R: RngCore>(
-    //     nv: usize,
-    //     num_multiplicands_range: (usize, usize),
-    //     num_products: usize,
-    //     rng: &mut R,
-    // ) -> VirtualPolynomial<F> {
-    //     let mut poly = VirtualPolynomial::new(nv);
-    //     for _ in 0..num_products {
-    //         let num_multiplicands =
-    //
-    // rng.gen_range(num_multiplicands_range.0..num_multiplicands_range.1);
-    //         let product = random_zero_product(nv, num_multiplicands, rng);
-    //         let coefficient = F::rand(rng);
-    //         poly.add_product(product.into_iter(), coefficient);
-    //     }
-
-    //     poly
-    // }
+    #[test]
+    fn test_virtual_polynomial_mul_by_mle() -> Result<(), PolyIOPErrors> {
+        let mut rng = test_rng();
+        for nv in 2..5 {
+            for num_products in 2..5 {
+                let base: Vec<Fr> = (0..nv).map(|_| Fr::rand(&mut rng)).collect();
+
+                let (a, _a_sum) =
+                    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 mut c = a.clone();
+
+                c.mul_by_mle(b_mle, coeff)?;
+
+                assert_eq!(
+                    a.evaluate(base.as_ref())? * b_vp.evaluate(base.as_ref())?,
+                    c.evaluate(base.as_ref())?
+                );
+            }
+        }
+
+        Ok(())
+    }
 }

poly-iop/src/zero_check/mod.rs

@@ -1,11 +1,7 @@
-mod prover;
-mod verifier;
-
 use ark_ff::PrimeField;
 use ark_poly::DenseMultilinearExtension;
-pub use prover::ProverState;
+use ark_std::{end_timer, start_timer};
 use std::rc::Rc;
-pub use verifier::VerifierState;
 
 use crate::{
     errors::PolyIOPErrors,
@@ -31,6 +27,8 @@ pub trait ZeroCheck {
     /// ZeroCheck prover/verifier.
     fn init_transcript() -> Self::Transcript;
 
+    /// initialize the prover to argue for the sum of polynomial over
+    /// {0,1}^`num_vars` is zero.
     fn prove(
         poly: &Self::PolyList,
         transcript: &mut Self::Transcript,
@@ -48,7 +46,11 @@ impl ZeroCheck for PolyIOP {
     type Proof = IOPProof<F>;
     type PolyList = VirtualPolynomial<F>;
     type DomainInfo = DomainInfo<F>;
-    type SubClaim = SubClaim<F>;
+
+    /// A ZeroCheck SubClaim consists of
+    /// - the SubClaim from the ZeroCheck
+    /// - the initial challenge r which is used to build eq(x, r)
+    type SubClaim = (SubClaim<F>, Vec<F>);
     type Transcript = IOPTranscript<F>;
 
     /// Initialize the system with a transcript
@@ -61,29 +63,53 @@ impl ZeroCheck for PolyIOP {
         IOPTranscript::<F>::new(b"Initializing ZeroCheck transcript")
     }
 
+    /// initialize the prover to argue for the sum of polynomial over
+    /// {0,1}^`num_vars` is zero.
     fn prove(
         poly: &Self::PolyList,
         transcript: &mut Self::Transcript,
     ) -> Result<Self::Proof, PolyIOPErrors> {
+        let start = start_timer!(|| "zero check prove");
+
         let length = poly.domain_info.num_variables;
         let r = transcript.get_and_append_challenge_vectors(b"vector r", length)?;
-        let f_hat = build_f_hat(poly, r.as_ref());
-        <Self as SumCheck<F>>::prove(&f_hat, transcript)
+        let f_hat = build_f_hat(poly, r.as_ref())?;
+        let res = <Self as SumCheck<F>>::prove(&f_hat, transcript);
+
+        end_timer!(start);
+        res
     }
 
     /// Verify the claimed sum using the proof.
-    /// Caller needs to makes sure that `\hat f = f * eq(x, r)`
+    /// the initial challenge `r` is also returned.
+    /// The caller needs to makes sure that `\hat f = f * eq(x, r)`
     fn verify(
         proof: &Self::Proof,
-        domain_info: &Self::DomainInfo,
+        fx_domain_info: &Self::DomainInfo,
         transcript: &mut Self::Transcript,
     ) -> Result<Self::SubClaim, PolyIOPErrors> {
-        println!(
-            "sum: {}",
-            proof.proofs[0].evaluations[0] + proof.proofs[0].evaluations[1]
-        );
+        let start = start_timer!(|| "zero check verify");
+
+        // check that the sum is zero
+        if proof.proofs[0].evaluations[0] + proof.proofs[0].evaluations[1] != F::zero() {
+            return Err(PolyIOPErrors::InvalidProof(format!(
+                "zero check: sum {} is not zero",
+                proof.proofs[0].evaluations[0] + proof.proofs[0].evaluations[1]
+            )));
+        }
+
+        // generate `r` and pass it to the caller for correctness check
+        let length = fx_domain_info.num_variables;
+        let r = transcript.get_and_append_challenge_vectors(b"vector r", length)?;
 
-        <Self as SumCheck<F>>::verify(F::zero(), proof, domain_info, transcript)
+        // hat_fx's max degree is increased by eq(x, r).degree() which is 1
+        let mut hat_fx_domain_info = fx_domain_info.clone();
+        hat_fx_domain_info.max_degree += 1;
+        let subclaim =
+            <Self as SumCheck<F>>::verify(F::zero(), proof, &hat_fx_domain_info, transcript)?;
+
+        end_timer!(start);
+        Ok((subclaim, r))
     }
 }
 
@@ -91,39 +117,20 @@ impl ZeroCheck for PolyIOP {
 //      \hat f(x) = \sum_{x_i \in eval_x} f(x_i) eq(x, r)
 // where
 //      eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
-fn build_f_hat<F: PrimeField>(poly: &VirtualPolynomial<F>, r: &[F]) -> VirtualPolynomial<F> {
+fn build_f_hat<F: PrimeField>(
+    poly: &VirtualPolynomial<F>,
+    r: &[F],
+) -> Result<VirtualPolynomial<F>, PolyIOPErrors> {
+    let start = start_timer!(|| "zero check build hat f");
+
     assert_eq!(poly.domain_info.num_variables, r.len());
-    let mut res = poly.clone();
+
     let eq_x_r = build_eq_x_r(r);
-    res.add_product([eq_x_r; 1], F::one());
-    // // First, we build array for {1 - r_i}
-    // let one_minus_r: Vec<F> = r.iter().map(|ri| F::one() - ri).collect();
-
-    // let mut eval = vec![];
-    // // let eq_x_r = build_eq_x_r(r);
-    // let num_var = r.len();
-    // let mut res = VirtualPolynomial::new(num_var);
-    // // res.add_product([eq_x_r; 1], F::one());
-
-    // for i in 0..1 << num_var {
-    //     let bit_sequence = bit_decompose(i, num_var);
-    //     let bit_points: Vec<F> = bit_sequence.iter().map(|&x| F::from(x as
-    // u64)).collect();     let mut eq_eval = F::one();
-    //     for (&bit, (ri, one_minus_ri)) in
-    // bit_sequence.iter().zip(r.iter().zip(one_minus_r.iter()))     {
-    //         if bit {
-    //             eq_eval *= ri;
-    //         } else {
-    //             eq_eval *= one_minus_ri;
-    //         }
-    //     }
-    //     eval.push(eq_eval * poly.evaluate(&bit_points))
-    // }
-    // let hat_f = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
-    //     num_var, eval,
-    // ));
-    // res.add_product([hat_f; 1], F::one());
-    res
+    let mut res = poly.clone();
+    res.mul_by_mle(eq_x_r, F::one())?;
+
+    end_timer!(start);
+    Ok(res)
 }
 
 // Evaluate
@@ -131,6 +138,8 @@ fn build_f_hat(poly: &VirtualPolynomial, r: &[F]) -> VirtualPo
 // over r, which is
 //      eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))
 fn build_eq_x_r<F: PrimeField>(r: &[F]) -> Rc<DenseMultilinearExtension<F>> {
+    let start = start_timer!(|| "zero check build build eq_x_r");
+
     // we build eq(x,r) from its evaluations
     // we want to evaluate eq(x,r) over x \in {0, 1}^num_vars
     // for example, with num_vars = 4, x is a binary vector of 4, then
@@ -157,18 +166,14 @@ fn build_eq_x_r(r: &[F]) -> Rc> {
 
         for (&bit, (ri, one_minus_ri)) in bit_sequence.iter().zip(r.iter().zip(one_minus_r.iter()))
         {
-            if bit {
-                current_eval *= *ri;
-            } else {
-                current_eval *= *one_minus_ri;
-            }
+            current_eval *= if bit { *ri } else { *one_minus_ri };
         }
         eval.push(current_eval);
     }
-    let res = Rc::new(DenseMultilinearExtension::from_evaluations_vec(
-        num_var, eval,
-    ));
+    let mle = DenseMultilinearExtension::from_evaluations_vec(num_var, eval);
 
+    let res = Rc::new(mle);
+    end_timer!(start);
     res
 }
 
@@ -186,131 +191,83 @@ fn bit_decompose(input: u64, num_var: usize) -> Vec {
 mod test {
 
     use super::ZeroCheck;
-    use crate::{virtual_poly::test::random_zero_list_of_products, PolyIOP, VirtualPolynomial};
+    use crate::{errors::PolyIOPErrors, PolyIOP, VirtualPolynomial};
     use ark_bls12_381::Fr;
-    use ark_ff::UniformRand;
-    use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
     use ark_std::test_rng;
-    use std::rc::Rc;
 
-    fn test_polynomial(nv: usize, num_multiplicands_range: (usize, usize), num_products: usize) {
+    fn test_zerocheck(
+        nv: usize,
+        num_multiplicands_range: (usize, usize),
+        num_products: usize,
+    ) -> Result<(), PolyIOPErrors> {
         let mut rng = test_rng();
-        let mut transcript = PolyIOP::init_transcript();
-        transcript
-            .append_message(b"testing", b"initializing transcript for testing")
-            .unwrap();
-        let poly = random_zero_list_of_products::<Fr, _>(
-            nv,
-            num_multiplicands_range,
-            num_products,
-            &mut rng,
-        );
-        // println!("{:?}", poly);
-
-        let proof = PolyIOP::prove(&poly, &mut transcript).expect("fail to prove");
-        println!(
-            "{}",
-            proof.proofs[0].evaluations[0] + proof.proofs[0].evaluations[1]
-        );
-
-        let poly_info = poly.domain_info.clone();
-        let mut transcript = PolyIOP::init_transcript();
-        transcript
-            .append_message(b"testing", b"initializing transcript for testing")
-            .unwrap();
-        let subclaim =
-            PolyIOP::verify(&proof, &poly_info, &mut transcript).expect("fail to verify");
-        assert!(
-            poly.evaluate(&subclaim.point) == subclaim.expected_evaluation,
-            "wrong subclaim"
-        );
+
+        {
+            // good path: zero virtual poly
+            let poly =
+                VirtualPolynomial::rand_zero(nv, num_multiplicands_range, num_products, &mut rng)?;
+
+            let mut transcript = <PolyIOP<Fr> as ZeroCheck<Fr>>::init_transcript();
+            transcript.append_message(b"testing", b"initializing transcript for testing")?;
+            let proof = <PolyIOP<Fr> as ZeroCheck<Fr>>::prove(&poly, &mut transcript)?;
+
+            let poly_info = poly.domain_info.clone();
+            let mut transcript = <PolyIOP<Fr> as ZeroCheck<Fr>>::init_transcript();
+            transcript.append_message(b"testing", b"initializing transcript for testing")?;
+            let subclaim =
+                <PolyIOP<Fr> as ZeroCheck<Fr>>::verify(&proof, &poly_info, &mut transcript)?.0;
+            assert!(
+                poly.evaluate(&subclaim.point)? == subclaim.expected_evaluation,
+                "wrong subclaim"
+            );
+        }
+
+        {
+            // bad path: random virtual poly whose sum is not zero
+            let (poly, _sum) =
+                VirtualPolynomial::rand(nv, num_multiplicands_range, num_products, &mut rng)?;
+
+            let mut transcript = <PolyIOP<Fr> as ZeroCheck<Fr>>::init_transcript();
+            transcript.append_message(b"testing", b"initializing transcript for testing")?;
+            let proof = <PolyIOP<Fr> as ZeroCheck<Fr>>::prove(&poly, &mut transcript)?;
+
+            let poly_info = poly.domain_info.clone();
+            let mut transcript = <PolyIOP<Fr> as ZeroCheck<Fr>>::init_transcript();
+            transcript.append_message(b"testing", b"initializing transcript for testing")?;
+
+            assert!(
+                <PolyIOP<Fr> as ZeroCheck<Fr>>::verify(&proof, &poly_info, &mut transcript)
+                    .is_err()
+            );
+        }
+
+        Ok(())
     }
 
     #[test]
-    fn test_trivial_polynomial() {
+    fn test_trivial_polynomial() -> Result<(), PolyIOPErrors> {
         let nv = 1;
         let num_multiplicands_range = (4, 5);
         let num_products = 1;
 
-        test_polynomial(nv, num_multiplicands_range, num_products);
+        test_zerocheck(nv, num_multiplicands_range, num_products)
     }
     #[test]
-    fn test_normal_polynomial() {
-        let nv = 16;
+    fn test_normal_polynomial() -> Result<(), PolyIOPErrors> {
+        let nv = 5;
         let num_multiplicands_range = (4, 9);
         let num_products = 5;
 
-        test_polynomial(nv, num_multiplicands_range, num_products);
+        test_zerocheck(nv, num_multiplicands_range, num_products)
     }
     #[test]
-    #[should_panic]
-    fn zero_polynomial_should_error() {
+
+    fn zero_polynomial_should_error() -> Result<(), PolyIOPErrors> {
         let nv = 0;
         let num_multiplicands_range = (4, 13);
         let num_products = 5;
 
-        test_polynomial(nv, num_multiplicands_range, num_products);
-    }
-
-    #[test]
-    /// Test that the memory usage of shared-reference is linear to number of
-    /// unique MLExtensions instead of total number of multiplicands.
-    fn test_shared_reference() {
-        let mut rng = test_rng();
-        let ml_extensions: Vec<_> = (0..5)
-            .map(|_| Rc::new(DenseMultilinearExtension::<Fr>::rand(8, &mut rng)))
-            .collect();
-        let mut poly = VirtualPolynomial::new(8);
-        poly.add_product(
-            vec![
-                ml_extensions[2].clone(),
-                ml_extensions[3].clone(),
-                ml_extensions[0].clone(),
-            ],
-            Fr::rand(&mut rng),
-        );
-        poly.add_product(
-            vec![
-                ml_extensions[1].clone(),
-                ml_extensions[4].clone(),
-                ml_extensions[4].clone(),
-            ],
-            Fr::rand(&mut rng),
-        );
-        poly.add_product(
-            vec![
-                ml_extensions[3].clone(),
-                ml_extensions[2].clone(),
-                ml_extensions[1].clone(),
-            ],
-            Fr::rand(&mut rng),
-        );
-        poly.add_product(
-            vec![ml_extensions[0].clone(), ml_extensions[0].clone()],
-            Fr::rand(&mut rng),
-        );
-        poly.add_product(vec![ml_extensions[4].clone()], Fr::rand(&mut rng));
-
-        assert_eq!(poly.flattened_ml_extensions.len(), 5);
-
-        let mut transcript = PolyIOP::init_transcript();
-
-        transcript
-            .append_message(b"testing", b"initializing transcript for testing")
-            .unwrap();
-        let poly_info = poly.domain_info.clone();
-        let proof = PolyIOP::prove(&poly, &mut transcript).expect("fail to prove");
-
-        let mut transcript = PolyIOP::init_transcript();
-
-        transcript
-            .append_message(b"testing", b"initializing transcript for testing")
-            .unwrap();
-        let subclaim =
-            PolyIOP::verify(&proof, &poly_info, &mut transcript).expect("fail to verify");
-        assert!(
-            poly.evaluate(&subclaim.point) == subclaim.expected_evaluation,
-            "wrong subclaim"
-        );
+        assert!(test_zerocheck(nv, num_multiplicands_range, num_products).is_err());
+        Ok(())
     }
 }

poly-iop/src/zero_check/prover.rs

@@ -1,17 +0,0 @@
-use ark_ff::PrimeField;
-use ark_poly::DenseMultilinearExtension;
-
-/// Prover State
-pub struct ProverState<F: PrimeField> {
-    /// sampled randomness given by the verifier
-    pub challenges: Vec<F>,
-    /// Stores the list of products that is meant to be added together. Each
-    /// multiplicand is represented by the index in flattened_ml_extensions
-    pub list_of_products: Vec<(F, Vec<usize>)>,
-    /// Stores a list of multilinear extensions in which `self.list_of_products`
-    /// points to
-    pub flattened_ml_extensions: Vec<DenseMultilinearExtension<F>>,
-    pub(crate) num_vars: usize,
-    pub(crate) max_degree: usize,
-    pub(crate) round: usize,
-}

poly-iop/src/zero_check/verifier.rs

@@ -1,14 +0,0 @@
-use ark_ff::PrimeField;
-
-/// Verifier State
-pub struct VerifierState<F: PrimeField> {
-    round: usize,
-    nv: usize,
-    max_degree: usize,
-    finished: bool,
-    /// a list storing the univariate polynomial in evaluation form sent by the
-    /// prover at each round
-    polynomials_received: Vec<Vec<F>>,
-    /// a list storing the randomness sampled by the verifier at each round
-    challenges: Vec<F>,
-}