nova-study/src/sumcheck.rs
src/sumcheck.rsRaw
1// Sum-check protocol initial implementation, not used by the rest of the repo but implemented as
2// an exercise and it will probably be used in the future.
3
4use ark_ff::{BigInteger, PrimeField};
5use ark_poly::{
6 multivariate::{SparsePolynomial, SparseTerm, Term},
7 univariate::DensePolynomial,
8 DenseMVPolynomial, DenseUVPolynomial, EvaluationDomain, GeneralEvaluationDomain, Polynomial,
9 SparseMultilinearExtension,
10};
11use ark_std::cfg_into_iter;
12use ark_std::log2;
13use ark_std::marker::PhantomData;
14use ark_std::ops::Mul;
15use ark_std::{rand::Rng, UniformRand};
16
17use crate::transcript::Transcript;
18
19pub struct SumCheck<
20 F: PrimeField,
21 UV: Polynomial<F> + DenseUVPolynomial<F>,
22 MV: Polynomial<F> + DenseMVPolynomial<F>,
23> {
24 _f: PhantomData<F>,
25 _uv: PhantomData<UV>,
26 _mv: PhantomData<MV>,
27}
28
29impl<
30 F: PrimeField,
31 UV: Polynomial<F> + DenseUVPolynomial<F>,
32 MV: Polynomial<F> + DenseMVPolynomial<F>,
33 > SumCheck<F, UV, MV>
34{
35 fn partial_evaluate(g: &MV, point: &[Option<F>]) -> UV {
36 assert!(point.len() >= g.num_vars(), "Invalid evaluation domain");
37 // TODO: add check: there can only be 1 'None' value in point
38 if g.is_zero() {
39 return UV::from_coefficients_vec(vec![F::zero()]);
40 }
41
42 // note: this can be parallelized with cfg_into_iter
43 let mut univ_terms: Vec<(F, SparseTerm)> = vec![];
44 for (coef, term) in g.terms().iter() {
45 // partial_evaluate each term
46 let mut new_coef = F::one();
47 let mut new_term = Vec::new();
48 for (var, power) in term.iter() {
49 match point[*var] {
50 Some(v) => {
51 if v.is_zero() {
52 new_coef = F::zero();
53 new_term = vec![];
54 break;
55 } else {
56 new_coef = new_coef * v.pow([(*power) as u64]);
57 }
58 }
59 _ => {
60 new_term.push((*var, *power));
61 }
62 };
63 }
64 let new_term = SparseTerm::new(new_term);
65 let new_coef = new_coef * coef;
66 univ_terms.push((new_coef, new_term));
67 }
68 let mv_poly: SparsePolynomial<F, SparseTerm> =
69 DenseMVPolynomial::<F>::from_coefficients_vec(g.num_vars(), univ_terms.clone());
70
71 let mut univ_coeffs: Vec<F> = vec![F::zero(); mv_poly.degree() + 1];
72 for (coef, term) in univ_terms {
73 if term.is_empty() {
74 univ_coeffs[0] += coef;
75 continue;
76 }
77 for (_, power) in term.iter() {
78 univ_coeffs[*power] += coef;
79 }
80 }
81 UV::from_coefficients_vec(univ_coeffs)
82 }
83
84 fn point_complete(challenges: Vec<F>, n_elems: usize, iter_num: usize) -> Vec<F> {
85 let p = Self::point(challenges, false, n_elems, iter_num);
86 let mut r = vec![F::zero(); n_elems];
87 for i in 0..n_elems {
88 r[i] = p[i].unwrap();
89 }
90 r
91 }
92 fn point(challenges: Vec<F>, none: bool, n_elems: usize, iter_num: usize) -> Vec<Option<F>> {
93 let mut n_vars = n_elems - challenges.len();
94 assert!(n_vars >= log2(iter_num + 1) as usize);
95
96 if none {
97 // WIP
98 if n_vars == 0 {
99 panic!("err"); // or return directly challenges vector
100 }
101 n_vars -= 1;
102 }
103 let none_pos = if none {
104 challenges.len() + 1
105 } else {
106 challenges.len()
107 };
108 let mut p: Vec<Option<F>> = vec![None; n_elems];
109 for i in 0..challenges.len() {
110 p[i] = Some(challenges[i]);
111 }
112 for i in 0..n_vars {
113 let k = F::from(iter_num as u64).into_bigint().to_bytes_le();
114 let bit = k[(i / 8) as usize] & (1 << (i % 8));
115 if bit == 0 {
116 p[none_pos + i] = Some(F::zero());
117 } else {
118 p[none_pos + i] = Some(F::one());
119 }
120 }
121 p
122 }
123
124 pub fn prove(g: MV) -> (F, Vec<UV>, F)
125 where
126 <MV as Polynomial<F>>::Point: From<Vec<F>>,
127 {
128 // init transcript
129 let mut transcript: Transcript<F> = Transcript::<F>::new();
130
131 let v = g.num_vars();
132
133 // compute H
134 let mut H = F::zero();
135 for i in 0..(2_u64.pow(v as u32) as usize) {
136 let p = Self::point_complete(vec![], v, i);
137
138 H = H + g.evaluate(&p.into());
139 }
140 transcript.add(b"H", &H);
141
142 let mut ss: Vec<UV> = Vec::new();
143 let mut r: Vec<F> = vec![];
144 for i in 0..v {
145 let r_i = transcript.get_challenge(b"r_i");
146 r.push(r_i);
147
148 let var_slots = v - 1 - i;
149 let n_points = 2_u64.pow(var_slots as u32) as usize;
150
151 let mut s_i = UV::zero();
152 for j in 0..n_points {
153 let point = Self::point(r[..i].to_vec(), true, v, j);
154 s_i = s_i + Self::partial_evaluate(&g, &point);
155 }
156 transcript.add(b"s_i", &s_i);
157 ss.push(s_i);
158 }
159
160 let last_g_eval = g.evaluate(&r.into());
161 (H, ss, last_g_eval)
162 }
163
164 pub fn verify(proof: (F, Vec<UV>, F)) -> bool {
165 // init transcript
166 let mut transcript: Transcript<F> = Transcript::<F>::new();
167 transcript.add(b"H", &proof.0);
168
169 let (c, ss, last_g_eval) = proof;
170
171 let mut r: Vec<F> = vec![];
172 for (i, s) in ss.iter().enumerate() {
173 // TODO check degree
174 if i == 0 {
175 if c != s.evaluate(&F::zero()) + s.evaluate(&F::one()) {
176 return false;
177 }
178 let r_i = transcript.get_challenge(b"r_i");
179 r.push(r_i);
180 transcript.add(b"s_i", s);
181 continue;
182 }
183
184 let r_i = transcript.get_challenge(b"r_i");
185 r.push(r_i);
186 if ss[i - 1].evaluate(&r[i - 1]) != s.evaluate(&F::zero()) + s.evaluate(&F::one()) {
187 return false;
188 }
189 transcript.add(b"s_i", s);
190 }
191 // last round
192 if ss[ss.len() - 1].evaluate(&r[r.len() - 1]) != last_g_eval {
193 return false;
194 }
195
196 true
197 }
198}
199
200#[cfg(test)]
201mod tests {
202 use super::*;
203 use ark_bn254::Fr; // scalar field
204
205 #[test]
206 fn test_new_point() {
207 let f4 = Fr::from(4_u32);
208 let f1 = Fr::from(1);
209 let f0 = Fr::from(0);
210 type SC = SumCheck<Fr, DensePolynomial<Fr>, SparsePolynomial<Fr, SparseTerm>>;
211
212 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 0);
213 assert_eq!(vec![Some(f4), None, Some(f0), Some(f0), Some(f0),], p);
214
215 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 1);
216 assert_eq!(vec![Some(f4), None, Some(f1), Some(f0), Some(f0),], p);
217
218 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 2);
219 assert_eq!(vec![Some(f4), None, Some(f0), Some(f1), Some(f0),], p);
220
221 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 3);
222 assert_eq!(vec![Some(f4), None, Some(f1), Some(f1), Some(f0),], p);
223
224 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 4);
225 assert_eq!(vec![Some(f4), None, Some(f0), Some(f0), Some(f1),], p);
226
227 // without None
228 let p = SC::point(vec![], false, 4, 0);
229 assert_eq!(vec![Some(f0), Some(f0), Some(f0), Some(f0),], p);
230
231 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 0);
232 assert_eq!(vec![Some(f4), Some(f0), Some(f0), Some(f0), Some(f0),], p);
233
234 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 1);
235 assert_eq!(vec![Some(f4), Some(f1), Some(f0), Some(f0), Some(f0),], p);
236
237 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 3);
238 assert_eq!(vec![Some(f4), Some(f1), Some(f1), Some(f0), Some(f0),], p);
239
240 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 4);
241 assert_eq!(vec![Some(f4), Some(f0), Some(f0), Some(f1), Some(f0),], p);
242
243 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 10);
244 assert_eq!(vec![Some(f4), Some(f0), Some(f1), Some(f0), Some(f1),], p);
245
246 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 15);
247 assert_eq!(vec![Some(f4), Some(f1), Some(f1), Some(f1), Some(f1),], p);
248
249 // let p = SC::point(vec![Fr::from(4_u32)], false, 4, 16); // TODO expect error
250 }
251
252 #[test]
253 fn test_partial_evaluate() {
254 // g(X_0, X_1, X_2) = 2 X_0^3 + X_0 X_2 + X_1 X_2
255 let terms = vec![
256 (Fr::from(2u32), SparseTerm::new(vec![(0_usize, 3)])),
257 (
258 Fr::from(1u32),
259 SparseTerm::new(vec![(0_usize, 1), (2_usize, 1)]),
260 ),
261 (
262 Fr::from(1u32),
263 SparseTerm::new(vec![(1_usize, 1), (2_usize, 1)]),
264 ),
265 ];
266 let p = SparsePolynomial::from_coefficients_slice(3, &terms);
267
268 type SC = SumCheck<Fr, DensePolynomial<Fr>, SparsePolynomial<Fr, SparseTerm>>;
269
270 let e0 = SC::partial_evaluate(&p, &[Some(Fr::from(2_u32)), None, Some(Fr::from(0_u32))]);
271 assert_eq!(e0.coeffs(), vec![Fr::from(16_u32)]);
272 let e1 = SC::partial_evaluate(&p, &[Some(Fr::from(2_u32)), None, Some(Fr::from(1_u32))]);
273 assert_eq!(e1.coeffs(), vec![Fr::from(18_u32), Fr::from(1)]);
274
275 assert_eq!((e0 + e1).coeffs(), vec![Fr::from(34_u32), Fr::from(1)]);
276 }
277
278 #[test]
279 fn test_flow_hardcoded_values() {
280 let mut rng = ark_std::test_rng();
281 // g(X_0, X_1, X_2) = 2 X_0^3 + X_0 X_2 + X_1 X_2
282 let terms = vec![
283 (Fr::from(2u32), SparseTerm::new(vec![(0_usize, 3)])),
284 (
285 Fr::from(1u32),
286 SparseTerm::new(vec![(0_usize, 1), (2_usize, 1)]),
287 ),
288 (
289 Fr::from(1u32),
290 SparseTerm::new(vec![(1_usize, 1), (2_usize, 1)]),
291 ),
292 ];
293 let p = SparsePolynomial::from_coefficients_slice(3, &terms);
294 // println!("p {:?}", p);
295
296 type SC = SumCheck<Fr, DensePolynomial<Fr>, SparsePolynomial<Fr, SparseTerm>>;
297
298 let proof = SC::prove(p);
299 assert_eq!(proof.0, Fr::from(12_u32));
300 // println!("proof {:?}", proof);
301
302 let v = SC::verify(proof);
303 assert!(v);
304 }
305
306 fn rand_poly<R: Rng>(l: usize, d: usize, rng: &mut R) -> SparsePolynomial<Fr, SparseTerm> {
307 // This method is from the arkworks/algebra/poly/multivariate test:
308 // https://github.com/arkworks-rs/algebra/blob/bc991d44c5e579025b7ed56df3d30267a7b9acac/poly/src/polynomial/multivariate/sparse.rs#L303
309 let mut random_terms = Vec::new();
310 let num_terms = rng.gen_range(1..1000);
311 // For each term, randomly select up to `l` variables with degree
312 // in [1,d] and random coefficient
313 random_terms.push((Fr::rand(rng), SparseTerm::new(vec![])));
314 for _ in 1..num_terms {
315 let term = (0..l)
316 .map(|i| {
317 if rng.gen_bool(0.5) {
318 Some((i, rng.gen_range(1..(d + 1))))
319 } else {
320 None
321 }
322 })
323 .flatten()
324 .collect();
325 let coeff = Fr::rand(rng);
326 random_terms.push((coeff, SparseTerm::new(term)));
327 }
328 SparsePolynomial::from_coefficients_slice(l, &random_terms)
329 }
330
331 #[test]
332 fn test_flow_rng() {
333 let mut rng = ark_std::test_rng();
334 // let p = SparsePolynomial::<Fr, SparseTerm>::rand(3, 3, &mut rng);
335 let p = rand_poly(3, 3, &mut rng);
336 // println!("p {:?}", p);
337
338 type SC = SumCheck<Fr, DensePolynomial<Fr>, SparsePolynomial<Fr, SparseTerm>>;
339
340 let proof = SC::prove(p);
341 println!("proof.s len {:?}", proof.1.len());
342
343 let v = SC::verify(proof);
344 assert!(v);
345 }
346}