1// Copyright (c) 2023 Espresso Systems (espressosys.com)
  2// This file is part of the HyperPlonk library.
  3
  4// You should have received a copy of the MIT License
  5// along with the HyperPlonk library. If not, see <https://mit-license.org/>.
  6
  7//! Prover subroutines for a SumCheck protocol.
  8
  9use super::SumCheckProver;
 10use crate::poly_iop::{
 11    errors::PolyIOPErrors,
 12    structs::{IOPProverMessage, IOPProverState},
 13};
 14use arithmetic::{fix_variables, VirtualPolynomial};
 15use ark_ff::{batch_inversion, PrimeField};
 16use ark_poly::DenseMultilinearExtension;
 17use ark_std::{cfg_into_iter, end_timer, start_timer, vec::Vec};
 18use rayon::prelude::{IntoParallelIterator, IntoParallelRefIterator};
 19use std::sync::Arc;
 20
 21#[cfg(feature = "parallel")]
 22use rayon::iter::{IntoParallelRefMutIterator, ParallelIterator};
 23
 24impl<F: PrimeField> SumCheckProver<F> for IOPProverState<F> {
 25    type VirtualPolynomial = VirtualPolynomial<F>;
 26    type ProverMessage = IOPProverMessage<F>;
 27
 28    /// Initialize the prover state to argue for the sum of the input polynomial
 29    /// over {0,1}^`num_vars`.
 30    fn prover_init(polynomial: &Self::VirtualPolynomial) -> Result<Self, PolyIOPErrors> {
 31        let start = start_timer!(|| "sum check prover init");
 32        if polynomial.aux_info.num_variables == 0 {
 33            return Err(PolyIOPErrors::InvalidParameters(
 34                "Attempt to prove a constant.".to_string(),
 35            ));
 36        }
 37        end_timer!(start);
 38
 39        Ok(Self {
 40            challenges: Vec::with_capacity(polynomial.aux_info.num_variables),
 41            round: 0,
 42            poly: polynomial.clone(),
 43            extrapolation_aux: (1..polynomial.aux_info.max_degree)
 44                .map(|degree| {
 45                    let points = (0..1 + degree as u64).map(F::from).collect::<Vec<_>>();
 46                    let weights = barycentric_weights(&points);
 47                    (points, weights)
 48                })
 49                .collect(),
 50        })
 51    }
 52
 53    /// Receive message from verifier, generate prover message, and proceed to
 54    /// next round.
 55    ///
 56    /// Main algorithm used is from section 3.2 of [XZZPS19](https://eprint.iacr.org/2019/317.pdf#subsection.3.2).
 57    fn prove_round_and_update_state(
 58        &mut self,
 59        challenge: &Option<F>,
 60    ) -> Result<Self::ProverMessage, PolyIOPErrors> {
 61        // let start =
 62        //     start_timer!(|| format!("sum check prove {}-th round and update state",
 63        // self.round));
 64
 65        if self.round >= self.poly.aux_info.num_variables {
 66            return Err(PolyIOPErrors::InvalidProver(
 67                "Prover is not active".to_string(),
 68            ));
 69        }
 70
 71        // let fix_argument = start_timer!(|| "fix argument");
 72
 73        // Step 1:
 74        // fix argument and evaluate f(x) over x_m = r; where r is the challenge
 75        // for the current round, and m is the round number, indexed from 1
 76        //
 77        // i.e.:
 78        // at round m <= n, for each mle g(x_1, ... x_n) within the flattened_mle
 79        // which has already been evaluated to
 80        //
 81        //    g(r_1, ..., r_{m-1}, x_m ... x_n)
 82        //
 83        // eval g over r_m, and mutate g to g(r_1, ... r_m,, x_{m+1}... x_n)
 84        let mut flattened_ml_extensions: Vec<DenseMultilinearExtension<F>> = self
 85            .poly
 86            .flattened_ml_extensions
 87            .par_iter()
 88            .map(|x| x.as_ref().clone())
 89            .collect();
 90
 91        if let Some(chal) = challenge {
 92            if self.round == 0 {
 93                return Err(PolyIOPErrors::InvalidProver(
 94                    "first round should be prover first.".to_string(),
 95                ));
 96            }
 97            self.challenges.push(*chal);
 98
 99            let r = self.challenges[self.round - 1];
100            #[cfg(feature = "parallel")]
101            flattened_ml_extensions
102                .par_iter_mut()
103                .for_each(|mle| *mle = fix_variables(mle, &[r]));
104            #[cfg(not(feature = "parallel"))]
105            flattened_ml_extensions
106                .iter_mut()
107                .for_each(|mle| *mle = fix_variables(mle, &[r]));
108        } else if self.round > 0 {
109            return Err(PolyIOPErrors::InvalidProver(
110                "verifier message is empty".to_string(),
111            ));
112        }
113        // end_timer!(fix_argument);
114
115        self.round += 1;
116
117        let products_list = self.poly.products.clone();
118        let mut products_sum = vec![F::zero(); self.poly.aux_info.max_degree + 1];
119
120        // Step 2: generate sum for the partial evaluated polynomial:
121        // f(r_1, ... r_m,, x_{m+1}... x_n)
122
123        products_list.iter().for_each(|(coefficient, products)| {
124            let mut sum = cfg_into_iter!(0..1 << (self.poly.aux_info.num_variables - self.round))
125                .fold(
126                    || {
127                        (
128                            vec![(F::zero(), F::zero()); products.len()],
129                            vec![F::zero(); products.len() + 1],
130                        )
131                    },
132                    |(mut buf, mut acc), b| {
133                        buf.iter_mut()
134                            .zip(products.iter())
135                            .for_each(|((eval, step), f)| {
136                                let table = &flattened_ml_extensions[*f];
137                                *eval = table[b << 1];
138                                *step = table[(b << 1) + 1] - table[b << 1];
139                            });
140                        acc[0] += buf.iter().map(|(eval, _)| eval).product::<F>();
141                        acc[1..].iter_mut().for_each(|acc| {
142                            buf.iter_mut().for_each(|(eval, step)| *eval += step as &_);
143                            *acc += buf.iter().map(|(eval, _)| eval).product::<F>();
144                        });
145                        (buf, acc)
146                    },
147                )
148                .map(|(_, partial)| partial)
149                .reduce(
150                    || vec![F::zero(); products.len() + 1],
151                    |mut sum, partial| {
152                        sum.iter_mut()
153                            .zip(partial.iter())
154                            .for_each(|(sum, partial)| *sum += partial);
155                        sum
156                    },
157                );
158            sum.iter_mut().for_each(|sum| *sum *= coefficient);
159            let extraploation = cfg_into_iter!(0..self.poly.aux_info.max_degree - products.len())
160                .map(|i| {
161                    let (points, weights) = &self.extrapolation_aux[products.len() - 1];
162                    let at = F::from((products.len() + 1 + i) as u64);
163                    extrapolate(points, weights, &sum, &at)
164                })
165                .collect::<Vec<_>>();
166            products_sum
167                .iter_mut()
168                .zip(sum.iter().chain(extraploation.iter()))
169                .for_each(|(products_sum, sum)| *products_sum += sum);
170        });
171
172        // update prover's state to the partial evaluated polynomial
173        self.poly.flattened_ml_extensions = flattened_ml_extensions
174            .par_iter()
175            .map(|x| Arc::new(x.clone()))
176            .collect();
177
178        Ok(IOPProverMessage {
179            evaluations: products_sum,
180        })
181    }
182}
183
184fn barycentric_weights<F: PrimeField>(points: &[F]) -> Vec<F> {
185    let mut weights = points
186        .iter()
187        .enumerate()
188        .map(|(j, point_j)| {
189            points
190                .iter()
191                .enumerate()
192                .filter_map(|(i, point_i)| (i != j).then(|| *point_j - point_i))
193                .reduce(|acc, value| acc * value)
194                .unwrap_or_else(F::one)
195        })
196        .collect::<Vec<_>>();
197    batch_inversion(&mut weights);
198    weights
199}
200
201fn extrapolate<F: PrimeField>(points: &[F], weights: &[F], evals: &[F], at: &F) -> F {
202    let (coeffs, sum_inv) = {
203        let mut coeffs = points.iter().map(|point| *at - point).collect::<Vec<_>>();
204        batch_inversion(&mut coeffs);
205        coeffs.iter_mut().zip(weights).for_each(|(coeff, weight)| {
206            *coeff *= weight;
207        });
208        let sum_inv = coeffs.iter().sum::<F>().inverse().unwrap_or_default();
209        (coeffs, sum_inv)
210    };
211    coeffs
212        .iter()
213        .zip(evals)
214        .map(|(coeff, eval)| *coeff * eval)
215        .sum::<F>()
216        * sum_inv
217}