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