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