hyperplonk/subroutines/src/poly_iop/sum_check/prover.rs
subroutines/src/poly_iop/sum_check/prover.rsRaw
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::PrimeField;
16use ark_poly::DenseMultilinearExtension;
17use ark_std::{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 })
44 }
45
46 /// Receive message from verifier, generate prover message, and proceed to
47 /// next round.
48 ///
49 /// Main algorithm used is from section 3.2 of [XZZPS19](https://eprint.iacr.org/2019/317.pdf#subsection.3.2).
50 fn prove_round_and_update_state(
51 &mut self,
52 challenge: &Option<F>,
53 ) -> Result<Self::ProverMessage, PolyIOPErrors> {
54 // let start =
55 // start_timer!(|| format!("sum check prove {}-th round and update state",
56 // self.round));
57
58 if self.round >= self.poly.aux_info.num_variables {
59 return Err(PolyIOPErrors::InvalidProver(
60 "Prover is not active".to_string(),
61 ));
62 }
63
64 // let fix_argument = start_timer!(|| "fix argument");
65
66 // Step 1:
67 // fix argument and evaluate f(x) over x_m = r; where r is the challenge
68 // for the current round, and m is the round number, indexed from 1
69 //
70 // i.e.:
71 // at round m <= n, for each mle g(x_1, ... x_n) within the flattened_mle
72 // which has already been evaluated to
73 //
74 // g(r_1, ..., r_{m-1}, x_m ... x_n)
75 //
76 // eval g over r_m, and mutate g to g(r_1, ... r_m,, x_{m+1}... x_n)
77 let mut flattened_ml_extensions: Vec<DenseMultilinearExtension<F>> = self
78 .poly
79 .flattened_ml_extensions
80 .par_iter()
81 .map(|x| x.as_ref().clone())
82 .collect();
83
84 if let Some(chal) = challenge {
85 if self.round == 0 {
86 return Err(PolyIOPErrors::InvalidProver(
87 "first round should be prover first.".to_string(),
88 ));
89 }
90 self.challenges.push(*chal);
91
92 let r = self.challenges[self.round - 1];
93 #[cfg(feature = "parallel")]
94 flattened_ml_extensions
95 .par_iter_mut()
96 .for_each(|mle| *mle = fix_variables(mle, &[r]));
97 #[cfg(not(feature = "parallel"))]
98 flattened_ml_extensions
99 .iter_mut()
100 .for_each(|mle| *mle = fix_variables(mle, &[r]));
101 } else if self.round > 0 {
102 return Err(PolyIOPErrors::InvalidProver(
103 "verifier message is empty".to_string(),
104 ));
105 }
106 // end_timer!(fix_argument);
107
108 self.round += 1;
109
110 let products_list = self.poly.products.clone();
111 let mut products_sum = vec![F::zero(); self.poly.aux_info.max_degree + 1];
112
113 // let compute_sum = start_timer!(|| "compute sum");
114
115 // Step 2: generate sum for the partial evaluated polynomial:
116 // f(r_1, ... r_m,, x_{m+1}... x_n)
117
118 #[cfg(feature = "parallel")]
119 {
120 let flag = (self.poly.aux_info.max_degree == 2)
121 && (products_list.len() == 1)
122 && (products_list[0].0 == F::one());
123 if flag {
124 for (t, e) in products_sum.iter_mut().enumerate() {
125 let evals = (0..1 << (self.poly.aux_info.num_variables - self.round))
126 .into_par_iter()
127 .map(|b| {
128 // evaluate P_round(t)
129 let table0 = &flattened_ml_extensions[products_list[0].1[0]];
130 let table1 = &flattened_ml_extensions[products_list[0].1[1]];
131 if t == 0 {
132 table0[b << 1] * table1[b << 1]
133 } else if t == 1 {
134 table0[(b << 1) + 1] * table1[(b << 1) + 1]
135 } else {
136 (table0[(b << 1) + 1] + table0[(b << 1) + 1] - table0[b << 1])
137 * (table1[(b << 1) + 1] + table1[(b << 1) + 1] - table1[b << 1])
138 }
139 })
140 .collect::<Vec<F>>();
141 *e += evals.par_iter().sum::<F>();
142 }
143 } else {
144 for (t, e) in products_sum.iter_mut().enumerate() {
145 let t = F::from(t as u128);
146 let products = (0..1 << (self.poly.aux_info.num_variables - self.round))
147 .into_par_iter()
148 .map(|b| {
149 // evaluate P_round(t)
150 let mut tmp = F::zero();
151 products_list.iter().for_each(|(coefficient, products)| {
152 let num_mles = products.len();
153 let mut product = *coefficient;
154 for &f in products.iter().take(num_mles) {
155 let table = &flattened_ml_extensions[f]; // f's range is checked in init
156 // TODO: Could be done faster by cashing the results from the
157 // previous t and adding the diff
158 // Also possible to use Karatsuba multiplication
159 product *=
160 table[b << 1] + (table[(b << 1) + 1] - table[b << 1]) * t;
161 }
162 tmp += product;
163 });
164
165 tmp
166 })
167 .collect::<Vec<F>>();
168 *e += products.par_iter().sum::<F>();
169 }
170 }
171 }
172
173 #[cfg(not(feature = "parallel"))]
174 products_sum.iter_mut().enumerate().for_each(|(t, e)| {
175 let t = F::from(t as u64);
176 let one_minus_t = F::one() - t;
177
178 for b in 0..1 << (self.poly.aux_info.num_variables - self.round) {
179 // evaluate P_round(t)
180 for (coefficient, products) in products_list.iter() {
181 let num_mles = products.len();
182 let mut product = *coefficient;
183 for &f in products.iter().take(num_mles) {
184 let table = &flattened_ml_extensions[f]; // f's range is checked in init
185 product *= table[b << 1] + (table[(b << 1) + 1] - table[b << 1]) * t;
186 }
187 *e += product;
188 }
189 }
190 });
191
192 // update prover's state to the partial evaluated polynomial
193 self.poly.flattened_ml_extensions = flattened_ml_extensions
194 .par_iter()
195 .map(|x| Arc::new(x.clone()))
196 .collect();
197
198 // end_timer!(compute_sum);
199 // end_timer!(start);
200 Ok(IOPProverMessage {
201 evaluations: products_sum,
202 })
203 }
204}