1//! Prover
  2
  3use super::SumCheckProver;
  4use crate::{
  5    errors::PolyIOPErrors,
  6    structs::{IOPProverMessage, IOPProverState},
  7    virtual_poly::VirtualPolynomial,
  8};
  9use ark_ff::PrimeField;
 10use ark_poly::{DenseMultilinearExtension, MultilinearExtension};
 11use ark_std::{end_timer, start_timer, vec::Vec};
 12use std::rc::Rc;
 13
 14#[cfg(feature = "parallel")]
 15use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
 16
 17impl<F: PrimeField> SumCheckProver<F> for IOPProverState<F> {
 18    type PolyList = VirtualPolynomial<F>;
 19    type ProverMessage = IOPProverMessage<F>;
 20
 21    /// Initialize the prover to argue for the sum of polynomial over
 22    /// {0,1}^`num_vars`
 23    fn prover_init(polynomial: &Self::PolyList) -> Result<Self, PolyIOPErrors> {
 24        let start = start_timer!(|| "sum check prover init");
 25        if polynomial.domain_info.num_variables == 0 {
 26            return Err(PolyIOPErrors::InvalidParameters(
 27                "Attempt to prove a constant.".to_string(),
 28            ));
 29        }
 30        end_timer!(start);
 31
 32        Ok(Self {
 33            challenges: Vec::with_capacity(polynomial.domain_info.num_variables),
 34            round: 0,
 35            poly: polynomial.clone(),
 36        })
 37    }
 38
 39    /// Receive message from verifier, generate prover message, and proceed to
 40    /// next round
 41    ///
 42    /// Main algorithm used is from section 3.2 of [XZZPS19](https://eprint.iacr.org/2019/317.pdf#subsection.3.2).
 43    fn prove_round_and_update_state(
 44        &mut self,
 45        challenge: &Option<F>,
 46    ) -> Result<Self::ProverMessage, PolyIOPErrors> {
 47        let start =
 48            start_timer!(|| format!("sum check prove {}-th round and update state", self.round));
 49
 50        let fix_argument = start_timer!(|| "fix argument");
 51
 52        let mut flattened_ml_extensions: Vec<DenseMultilinearExtension<F>> = self
 53            .poly
 54            .flattened_ml_extensions
 55            .iter()
 56            .map(|x| x.as_ref().clone())
 57            .collect();
 58
 59        if let Some(chal) = challenge {
 60            if self.round == 0 {
 61                return Err(PolyIOPErrors::InvalidProver(
 62                    "first round should be prover first.".to_string(),
 63                ));
 64            }
 65            self.challenges.push(*chal);
 66
 67            // fix argument
 68            let i = self.round;
 69            let r = self.challenges[i - 1];
 70            #[cfg(feature = "parallel")]
 71            flattened_ml_extensions
 72                .par_iter_mut()
 73                .for_each(|multiplicand| *multiplicand = multiplicand.fix_variables(&[r]));
 74
 75            #[cfg(not(feature = "parallel"))]
 76            flattened_ml_extensions
 77                .iter_mut()
 78                .for_each(|multiplicand| *multiplicand = multiplicand.fix_variables(&[r]));
 79        } else if self.round > 0 {
 80            return Err(PolyIOPErrors::InvalidProver(
 81                "verifier message is empty".to_string(),
 82            ));
 83        }
 84        end_timer!(fix_argument);
 85
 86        self.round += 1;
 87
 88        if self.round > self.poly.domain_info.num_variables {
 89            return Err(PolyIOPErrors::InvalidProver(
 90                "Prover is not active".to_string(),
 91            ));
 92        }
 93
 94        let products_list = self.poly.products.clone();
 95        let i = self.round;
 96        let nv = self.poly.domain_info.num_variables;
 97        let degree = self.poly.domain_info.max_degree; // the degree of univariate polynomial sent by prover at this round
 98
 99        let mut products_sum = Vec::with_capacity(degree + 1);
100        products_sum.resize(degree + 1, F::zero());
101
102        let compute_sum = start_timer!(|| "compute sum");
103        // generate sum
104        #[cfg(feature = "parallel")]
105        products_sum.par_iter_mut().enumerate().for_each(|(t, e)| {
106            for b in 0..1 << (nv - i) {
107                // evaluate P_round(t)
108                for (coefficient, products) in products_list.iter() {
109                    let num_multiplicands = products.len();
110                    let mut product = *coefficient;
111                    for &f in products.iter().take(num_multiplicands) {
112                        let table = &flattened_ml_extensions[f]; // f's range is checked in init
113                        product *= table[b << 1] * (F::one() - F::from(t as u64))
114                            + table[(b << 1) + 1] * F::from(t as u64);
115                    }
116                    *e += product;
117                }
118            }
119        });
120
121        #[cfg(not(feature = "parallel"))]
122        for b in 0..1 << (nv - i) {
123            products_sum
124                .iter_mut()
125                .take(degree + 1)
126                .enumerate()
127                .for_each(|(t, e)| {
128                    // evaluate P_round(t)
129                    for (coefficient, products) in products_list.iter() {
130                        let num_multiplicands = products.len();
131                        let mut product = *coefficient;
132                        for &f in products.iter().take(num_multiplicands) {
133                            let table = &flattened_ml_extensions[f]; // f's range is checked in init
134                            product *= table[b << 1] * (F::one() - F::from(t as u64))
135                                + table[(b << 1) + 1] * F::from(t as u64);
136                        }
137                        *e += product;
138                    }
139                });
140        }
141
142        self.poly.flattened_ml_extensions = flattened_ml_extensions
143            .iter()
144            .map(|x| Rc::new(x.clone()))
145            .collect();
146
147        end_timer!(compute_sum);
148        end_timer!(start);
149        Ok(IOPProverMessage {
150            evaluations: products_sum,
151        })
152    }
153}