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;
13use std::rc::Rc;
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 .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 for val in evals.iter() {
136 *e += val
137 }
138 }
139 } else {
140 for (t, e) in products_sum.iter_mut().enumerate() {
141 let t = F::from(t as u64);
142 let products = (0..1 << (self.poly.aux_info.num_variables - self.round))
143 .into_par_iter()
144 .map(|b| {
145 // evaluate P_round(t)
146 let mut tmp = F::zero();
147 products_list.iter().for_each(|(coefficient, products)| {
148 let num_mles = products.len();
149 let mut product = *coefficient;
150 for &f in products.iter().take(num_mles) {
151 let table = &flattened_ml_extensions[f]; // f's range is checked in init
152 product *=
153 table[b << 1] + (table[(b << 1) + 1] - table[b << 1]) * t;
154 }
155 tmp += product;
156 });
157
158 tmp
159 })
160 .collect::<Vec<F>>();
161
162 for i in products.iter() {
163 *e += i
164 }
165 }
166 }
167 }
168
169 #[cfg(not(feature = "parallel"))]
170 products_sum.iter_mut().enumerate().for_each(|(t, e)| {
171 let t = F::from(t as u64);
172 let one_minus_t = F::one() - t;
173
174 for b in 0..1 << (self.poly.aux_info.num_variables - self.round) {
175 // evaluate P_round(t)
176 for (coefficient, products) in products_list.iter() {
177 let num_mles = products.len();
178 let mut product = *coefficient;
179 for &f in products.iter().take(num_mles) {
180 let table = &flattened_ml_extensions[f]; // f's range is checked in init
181 product *= table[b << 1] + (table[(b << 1) + 1] - table[b << 1]) * t;
182 }
183 *e += product;
184 }
185 }
186 });
187
188 // update prover's state to the partial evaluated polynomial
189 self.poly.flattened_ml_extensions = flattened_ml_extensions
190 .iter()
191 .map(|x| Rc::new(x.clone()))
192 .collect();
193
194 // end_timer!(compute_sum);
195 // end_timer!(start);
196 Ok(IOPProverMessage {
197 evaluations: products_sum,
198 })
199 }
200}