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::{
 14    lagrange_poly::compute_lagrange_interpolated_poly, multilinear_polynomial::fix_variables,
 15    virtual_polynomial::VirtualPolynomial,
 16};
 17use ark_ff::{batch_inversion, PrimeField};
 18use ark_poly::DenseMultilinearExtension;
 19use ark_std::{cfg_into_iter, end_timer, start_timer};
 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<F: PrimeField> SumCheckProver<F> for IOPProverState<F> {
 30    type VirtualPolynomial = VirtualPolynomial<F>;
 31    type ProverMessage = IOPProverMessage<F>;
 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).map(F::from).collect::<Vec<_>>();
 51                    let weights = barycentric_weights(&points);
 52                    (points, weights)
 53                })
 54                .collect(),
 55        })
 56    }
 57
 58    /// Receive message from verifier, generate prover message, and proceed to
 59    /// next round.
 60    ///
 61    /// Main algorithm used is from section 3.2 of [XZZPS19](https://eprint.iacr.org/2019/317.pdf#subsection.3.2).
 62    fn prove_round_and_update_state(
 63        &mut self,
 64        challenge: &Option<F>,
 65    ) -> Result<Self::ProverMessage, PolyIOPErrors> {
 66        // let start =
 67        //     start_timer!(|| format!("sum check prove {}-th round and update state",
 68        // self.round));
 69
 70        if self.round >= self.poly.aux_info.num_variables {
 71            return Err(PolyIOPErrors::InvalidProver(
 72                "Prover is not active".to_string(),
 73            ));
 74        }
 75
 76        // let fix_argument = start_timer!(|| "fix argument");
 77
 78        // Step 1:
 79        // fix argument and evaluate f(x) over x_m = r; where r is the challenge
 80        // for the current round, and m is the round number, indexed from 1
 81        //
 82        // i.e.:
 83        // at round m <= n, for each mle g(x_1, ... x_n) within the flattened_mle
 84        // which has already been evaluated to
 85        //
 86        //    g(r_1, ..., r_{m-1}, x_m ... x_n)
 87        //
 88        // eval g over r_m, and mutate g to g(r_1, ... r_m,, x_{m+1}... x_n)
 89        let mut flattened_ml_extensions: Vec<DenseMultilinearExtension<F>> = self
 90            .poly
 91            .flattened_ml_extensions
 92            .par_iter()
 93            .map(|x| x.as_ref().clone())
 94            .collect();
 95
 96        if let Some(chal) = challenge {
 97            if self.round == 0 {
 98                return Err(PolyIOPErrors::InvalidProver(
 99                    "first round should be prover first.".to_string(),
100                ));
101            }
102            self.challenges.push(*chal);
103
104            let r = self.challenges[self.round - 1];
105            // #[cfg(feature = "parallel")]
106            flattened_ml_extensions
107                .par_iter_mut()
108                .for_each(|mle| *mle = fix_variables(mle, &[r]));
109            // #[cfg(not(feature = "parallel"))]
110            // flattened_ml_extensions
111            //     .iter_mut()
112            //     .for_each(|mle| *mle = fix_variables(mle, &[r]));
113        } else if self.round > 0 {
114            return Err(PolyIOPErrors::InvalidProver(
115                "verifier message is empty".to_string(),
116            ));
117        }
118        // end_timer!(fix_argument);
119
120        self.round += 1;
121
122        let products_list = self.poly.products.clone();
123        let mut products_sum = vec![F::ZERO; self.poly.aux_info.max_degree + 1];
124
125        // Step 2: generate sum for the partial evaluated polynomial:
126        // f(r_1, ... r_m,, x_{m+1}... x_n)
127
128        products_list.iter().for_each(|(coefficient, products)| {
129            let mut sum = cfg_into_iter!(0..1 << (self.poly.aux_info.num_variables - self.round))
130                .fold(
131                    || {
132                        (
133                            vec![(F::ZERO, F::ZERO); products.len()],
134                            vec![F::ZERO; products.len() + 1],
135                        )
136                    },
137                    |(mut buf, mut acc), b| {
138                        buf.iter_mut()
139                            .zip(products.iter())
140                            .for_each(|((eval, step), f)| {
141                                let table = &flattened_ml_extensions[*f];
142                                *eval = table[b << 1];
143                                *step = table[(b << 1) + 1] - table[b << 1];
144                            });
145                        acc[0] += buf.iter().map(|(eval, _)| eval).product::<F>();
146                        acc[1..].iter_mut().for_each(|acc| {
147                            buf.iter_mut().for_each(|(eval, step)| *eval += step as &_);
148                            *acc += buf.iter().map(|(eval, _)| eval).product::<F>();
149                        });
150                        (buf, acc)
151                    },
152                )
153                .map(|(_, partial)| partial)
154                .reduce(
155                    || vec![F::ZERO; products.len() + 1],
156                    |mut sum, partial| {
157                        sum.iter_mut()
158                            .zip(partial.iter())
159                            .for_each(|(sum, partial)| *sum += partial);
160                        sum
161                    },
162                );
163            sum.iter_mut().for_each(|sum| *sum *= coefficient);
164            let extraploation = cfg_into_iter!(0..self.poly.aux_info.max_degree - products.len())
165                .map(|i| {
166                    let (points, weights) = &self.extrapolation_aux[products.len() - 1];
167                    let at = F::from((products.len() + 1 + i) as u64);
168                    extrapolate(points, weights, &sum, &at)
169                })
170                .collect::<Vec<_>>();
171            products_sum
172                .iter_mut()
173                .zip(sum.iter().chain(extraploation.iter()))
174                .for_each(|(products_sum, sum)| *products_sum += sum);
175        });
176
177        // update prover's state to the partial evaluated polynomial
178        self.poly.flattened_ml_extensions = flattened_ml_extensions
179            .par_iter()
180            .map(|x| Arc::new(x.clone()))
181            .collect();
182
183        let prover_poly = compute_lagrange_interpolated_poly::<F>(&products_sum);
184        Ok(IOPProverMessage {
185            coeffs: prover_poly.coeffs,
186        })
187    }
188}
189
190#[allow(clippy::filter_map_bool_then)]
191fn barycentric_weights<F: PrimeField>(points: &[F]) -> Vec<F> {
192    let mut weights = points
193        .iter()
194        .enumerate()
195        .map(|(j, point_j)| {
196            points
197                .iter()
198                .enumerate()
199                .filter_map(|(i, point_i)| (i != j).then(|| *point_j - point_i))
200                .reduce(|acc, value| acc * value)
201                .unwrap_or_else(F::one)
202        })
203        .collect::<Vec<_>>();
204    batch_inversion(&mut weights);
205    weights
206}
207
208fn extrapolate<F: PrimeField>(points: &[F], weights: &[F], evals: &[F], at: &F) -> F {
209    let (coeffs, sum_inv) = {
210        let mut coeffs = points.iter().map(|point| *at - point).collect::<Vec<_>>();
211        batch_inversion(&mut coeffs);
212        coeffs.iter_mut().zip(weights).for_each(|(coeff, weight)| {
213            *coeff *= weight;
214        });
215        let sum_inv = coeffs.iter().sum::<F>().inverse().unwrap_or_default();
216        (coeffs, sum_inv)
217    };
218    coeffs
219        .iter()
220        .zip(evals)
221        .map(|(coeff, eval)| *coeff * eval)
222        .sum::<F>()
223        * sum_inv
224}