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