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        #[cfg(feature = "parallel")]
126        products_sum.par_iter_mut().enumerate().for_each(|(t, e)| {
127            for b in 0..1 << (nv - i) {
128                // evaluate P_round(t)
129                for (coefficient, products) in products.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        #[cfg(not(feature = "parallel"))]
143        for b in 0..1 << (nv - i) {
144            products_sum
145                .iter_mut()
146                .take(degree + 1)
147                .enumerate()
148                .for_each(|(t, e)| {
149                    // evaluate P_round(t)
150                    for (coefficient, products) in products.iter() {
151                        let num_multiplicands = products.len();
152                        let mut product = *coefficient;
153                        for &f in products.iter().take(num_multiplicands) {
154                            let table = &flattened_ml_extensions[f]; // f's range is checked in init
155                            product *= table[b << 1] * (F::one() - F::from(t as u64))
156                                + table[(b << 1) + 1] * F::from(t as u64);
157                        }
158                        *e += product;
159                    }
160                });
161        }
162
163        self.poly.flattened_ml_extensions = flattened_ml_extensions
164            .iter()
165            .map(|x| Rc::new(x.clone()))
166            .collect();
167
168        end_timer!(compute_sum);
169        end_timer!(start);
170        Ok(IOPProverMessage {
171            evaluations: products_sum,
172        })
173    }
174}