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