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