1//! Prover subroutines for a SumCheck protocol.
  2
  3use super::SumCheckProver;
  4use crate::{
  5    errors::PolyIOPErrors,
  6    structs::{IOPProverMessage, IOPProverState},
  7};
  8use arithmetic::VirtualPolynomial;
  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 VirtualPolynomial = VirtualPolynomial<F>;
 19    type ProverMessage = IOPProverMessage<F>;
 20
 21    /// Initialize the prover state to argue for the sum of the input polynomial
 22    /// over {0,1}^`num_vars`.
 23    fn prover_init(polynomial: &Self::VirtualPolynomial) -> Result<Self, PolyIOPErrors> {
 24        let start = start_timer!(|| "sum check prover init");
 25        if polynomial.aux_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.aux_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        if self.round >= self.poly.aux_info.num_variables {
 51            return Err(PolyIOPErrors::InvalidProver(
 52                "Prover is not active".to_string(),
 53            ));
 54        }
 55
 56        let fix_argument = start_timer!(|| "fix argument");
 57
 58        // Step 1:
 59        // fix argument and evaluate f(x) over x_m = r; where r is the challenge
 60        // for the current round, and m is the round number, indexed from 1
 61        //
 62        // i.e.:
 63        // at round m <= n, for each mle g(x_1, ... x_n) within the flattened_mle
 64        // which has already been evaluated to
 65        //
 66        //    g(r_1, ..., r_{m-1}, x_m ... x_n)
 67        //
 68        // eval g over r_m, and mutate g to g(r_1, ... r_m,, x_{m+1}... x_n)
 69        let mut flattened_ml_extensions: Vec<DenseMultilinearExtension<F>> = self
 70            .poly
 71            .flattened_ml_extensions
 72            .iter()
 73            .map(|x| x.as_ref().clone())
 74            .collect();
 75
 76        if let Some(chal) = challenge {
 77            if self.round == 0 {
 78                return Err(PolyIOPErrors::InvalidProver(
 79                    "first round should be prover first.".to_string(),
 80                ));
 81            }
 82            self.challenges.push(*chal);
 83
 84            let r = self.challenges[self.round - 1];
 85            #[cfg(feature = "parallel")]
 86            flattened_ml_extensions
 87                .par_iter_mut()
 88                .for_each(|mle| *mle = mle.fix_variables(&[r]));
 89
 90            #[cfg(not(feature = "parallel"))]
 91            flattened_ml_extensions
 92                .iter_mut()
 93                .for_each(|mle| *mle = mle.fix_variables(&[r]));
 94        } else if self.round > 0 {
 95            return Err(PolyIOPErrors::InvalidProver(
 96                "verifier message is empty".to_string(),
 97            ));
 98        }
 99        end_timer!(fix_argument);
100
101        self.round += 1;
102
103        let products_list = self.poly.products.clone();
104        let mut products_sum = Vec::with_capacity(self.poly.aux_info.max_degree + 1);
105        products_sum.resize(self.poly.aux_info.max_degree + 1, F::zero());
106
107        let compute_sum = start_timer!(|| "compute sum");
108        // Step 2: generate sum for the partial evaluated polynomial:
109        // f(r_1, ... r_m,, x_{m+1}... x_n)
110
111        #[cfg(feature = "parallel")]
112        products_sum.par_iter_mut().enumerate().for_each(|(t, e)| {
113            for b in 0..1 << (self.poly.aux_info.num_variables - self.round) {
114                // evaluate P_round(t)
115                for (coefficient, products) in products_list.iter() {
116                    let num_mles = products.len();
117                    let mut product = *coefficient;
118                    for &f in products.iter().take(num_mles) {
119                        let table = &flattened_ml_extensions[f]; // f's range is checked in init
120                        product *= table[b << 1] * (F::one() - F::from(t as u64))
121                            + table[(b << 1) + 1] * F::from(t as u64);
122                    }
123                    *e += product;
124                }
125            }
126        });
127
128        #[cfg(not(feature = "parallel"))]
129        products_sum.iter_mut().enumerate().for_each(|(t, e)| {
130            for b in 0..1 << (self.poly.aux_info.num_variables - self.round) {
131                // evaluate P_round(t)
132                for (coefficient, products) in products_list.iter() {
133                    let num_mles = products.len();
134                    let mut product = *coefficient;
135                    for &f in products.iter().take(num_mles) {
136                        let table = &flattened_ml_extensions[f]; // f's range is checked in init
137                        product *= table[b << 1] * (F::one() - F::from(t as u64))
138                            + table[(b << 1) + 1] * F::from(t as u64);
139                    }
140                    *e += product;
141                }
142            }
143        });
144
145        // update prover's state to the partial evaluated polynomial
146        self.poly.flattened_ml_extensions = flattened_ml_extensions
147            .iter()
148            .map(|x| Rc::new(x.clone()))
149            .collect();
150
151        end_timer!(compute_sum);
152        end_timer!(start);
153        Ok(IOPProverMessage {
154            evaluations: products_sum,
155        })
156    }
157}