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