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