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;
5use ark_ff::{BigInteger, PrimeField};
6use ark_poly::{
7 multivariate::{SparsePolynomial, SparseTerm, Term},
8 DenseMVPolynomial, DenseUVPolynomial, Polynomial,
9};
10use ark_std::log2;
11use ark_std::marker::PhantomData;
12
13use ark_crypto_primitives::sponge::{poseidon::PoseidonConfig, Absorb};
14
15use crate::transcript::Transcript;
16
17pub struct SumCheck<
18 F: PrimeField + Absorb,
19 C: CurveGroup,
20 UV: DenseUVPolynomial<F>,
21 MV: DenseMVPolynomial<F>,
22> {
23 _f: PhantomData<F>,
24 _c: PhantomData<C>,
25 _uv: PhantomData<UV>,
26 _mv: PhantomData<MV>,
27}
28
29impl<
30 F: PrimeField + Absorb,
31 C: CurveGroup,
32 UV: Polynomial<F> + DenseUVPolynomial<F>,
33 MV: Polynomial<F> + DenseMVPolynomial<F>,
34 > SumCheck<F, C, UV, MV>
35where
36 <C as CurveGroup>::BaseField: Absorb,
37{
38 fn partial_evaluate(g: &MV, point: &[Option<F>]) -> UV {
39 assert!(point.len() >= g.num_vars(), "Invalid evaluation domain");
40 // TODO: add check: there can only be 1 'None' value in point
41 if g.is_zero() {
42 return UV::from_coefficients_vec(vec![F::zero()]);
43 }
44
45 // note: this can be parallelized with cfg_into_iter
46 let mut univ_terms: Vec<(F, SparseTerm)> = vec![];
47 for (coef, term) in g.terms().iter() {
48 // partial_evaluate each term
49 let mut new_coef = F::one();
50 let mut new_term = Vec::new();
51 for (var, power) in term.iter() {
52 match point[*var] {
53 Some(v) => {
54 if v.is_zero() {
55 new_coef = F::zero();
56 new_term = vec![];
57 break;
58 } else {
59 new_coef *= v.pow([(*power) as u64]);
60 }
61 }
62 _ => {
63 new_term.push((*var, *power));
64 }
65 };
66 }
67 let new_term = SparseTerm::new(new_term);
68 let new_coef = new_coef * coef;
69 univ_terms.push((new_coef, new_term));
70 }
71 let mv_poly: SparsePolynomial<F, SparseTerm> =
72 DenseMVPolynomial::<F>::from_coefficients_vec(g.num_vars(), univ_terms.clone());
73
74 let mut univ_coeffs: Vec<F> = vec![F::zero(); mv_poly.degree() + 1];
75 for (coef, term) in univ_terms {
76 if term.is_empty() {
77 univ_coeffs[0] += coef;
78 continue;
79 }
80 for (_, power) in term.iter() {
81 univ_coeffs[*power] += coef;
82 }
83 }
84 UV::from_coefficients_vec(univ_coeffs)
85 }
86
87 fn point_complete(challenges: Vec<F>, n_elems: usize, iter_num: usize) -> Vec<F> {
88 let p = Self::point(challenges, false, n_elems, iter_num);
89 let mut r = vec![F::zero(); n_elems];
90 for i in 0..n_elems {
91 r[i] = p[i].unwrap();
92 }
93 r
94 }
95 fn point(challenges: Vec<F>, none: bool, n_elems: usize, iter_num: usize) -> Vec<Option<F>> {
96 let mut n_vars = n_elems - challenges.len();
97 assert!(n_vars >= log2(iter_num + 1) as usize);
98
99 if none {
100 // WIP
101 if n_vars == 0 {
102 panic!("err"); // or return directly challenges vector
103 }
104 n_vars -= 1;
105 }
106 let none_pos = if none {
107 challenges.len() + 1
108 } else {
109 challenges.len()
110 };
111 let mut p: Vec<Option<F>> = vec![None; n_elems];
112 for i in 0..challenges.len() {
113 p[i] = Some(challenges[i]);
114 }
115 for i in 0..n_vars {
116 let k = F::from(iter_num as u64).into_bigint().to_bytes_le();
117 let bit = k[i / 8] & (1 << (i % 8));
118 if bit == 0 {
119 p[none_pos + i] = Some(F::zero());
120 } else {
121 p[none_pos + i] = Some(F::one());
122 }
123 }
124 p
125 }
126
127 pub fn prove(poseidon_config: &PoseidonConfig<F>, g: MV) -> (F, Vec<UV>, F)
128 where
129 <MV as Polynomial<F>>::Point: From<Vec<F>>,
130 {
131 // init transcript
132 let mut transcript = Transcript::<F, C>::new(poseidon_config);
133
134 let v = g.num_vars();
135
136 // compute H
137 let mut H = F::zero();
138 for i in 0..(2_u64.pow(v as u32) as usize) {
139 let p = Self::point_complete(vec![], v, i);
140
141 H += g.evaluate(&p.into());
142 }
143 transcript.add(&H);
144
145 let mut ss: Vec<UV> = Vec::new();
146 let mut r: Vec<F> = vec![];
147 for i in 0..v {
148 let r_i = transcript.get_challenge();
149 r.push(r_i);
150
151 let var_slots = v - 1 - i;
152 let n_points = 2_u64.pow(var_slots as u32) as usize;
153
154 let mut s_i = UV::zero();
155 for j in 0..n_points {
156 let point = Self::point(r[..i].to_vec(), true, v, j);
157 s_i = s_i + Self::partial_evaluate(&g, &point);
158 }
159 transcript.add_vec(s_i.coeffs());
160 ss.push(s_i);
161 }
162
163 let last_g_eval = g.evaluate(&r.into());
164 (H, ss, last_g_eval)
165 }
166
167 pub fn verify(poseidon_config: &PoseidonConfig<F>, proof: (F, Vec<UV>, F)) -> bool {
168 // init transcript
169 let mut transcript = Transcript::<F, C>::new(poseidon_config);
170 transcript.add(&proof.0);
171
172 let (c, ss, last_g_eval) = proof;
173
174 let mut r: Vec<F> = vec![];
175 for (i, s) in ss.iter().enumerate() {
176 // TODO check degree
177 if i == 0 {
178 if c != s.evaluate(&F::zero()) + s.evaluate(&F::one()) {
179 return false;
180 }
181 let r_i = transcript.get_challenge();
182 r.push(r_i);
183 transcript.add_vec(s.coeffs());
184 continue;
185 }
186
187 let r_i = transcript.get_challenge();
188 r.push(r_i);
189 if ss[i - 1].evaluate(&r[i - 1]) != s.evaluate(&F::zero()) + s.evaluate(&F::one()) {
190 return false;
191 }
192 transcript.add_vec(s.coeffs());
193 }
194 // last round
195 if ss[ss.len() - 1].evaluate(&r[r.len() - 1]) != last_g_eval {
196 return false;
197 }
198
199 true
200 }
201}
202
203#[cfg(test)]
204mod tests {
205 use super::*;
206 use crate::transcript::poseidon_test_config;
207 use ark_mnt4_298::{Fr, G1Projective}; // scalar field
208 use ark_poly::{
209 multivariate::{SparsePolynomial, SparseTerm, Term},
210 univariate::DensePolynomial,
211 DenseMVPolynomial, DenseUVPolynomial,
212 };
213 use ark_std::{rand::Rng, UniformRand};
214
215 #[test]
216 fn test_new_point() {
217 let f4 = Fr::from(4_u32);
218 let f1 = Fr::from(1);
219 let f0 = Fr::from(0);
220 type SC = SumCheck<Fr, G1Projective, DensePolynomial<Fr>, SparsePolynomial<Fr, SparseTerm>>;
221
222 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 0);
223 assert_eq!(vec![Some(f4), None, Some(f0), Some(f0), Some(f0),], p);
224
225 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 1);
226 assert_eq!(vec![Some(f4), None, Some(f1), Some(f0), Some(f0),], p);
227
228 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 2);
229 assert_eq!(vec![Some(f4), None, Some(f0), Some(f1), Some(f0),], p);
230
231 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 3);
232 assert_eq!(vec![Some(f4), None, Some(f1), Some(f1), Some(f0),], p);
233
234 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 4);
235 assert_eq!(vec![Some(f4), None, Some(f0), Some(f0), Some(f1),], p);
236
237 // without None
238 let p = SC::point(vec![], false, 4, 0);
239 assert_eq!(vec![Some(f0), Some(f0), Some(f0), Some(f0),], p);
240
241 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 0);
242 assert_eq!(vec![Some(f4), Some(f0), Some(f0), Some(f0), Some(f0),], p);
243
244 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 1);
245 assert_eq!(vec![Some(f4), Some(f1), Some(f0), Some(f0), Some(f0),], p);
246
247 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 3);
248 assert_eq!(vec![Some(f4), Some(f1), Some(f1), Some(f0), Some(f0),], p);
249
250 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 4);
251 assert_eq!(vec![Some(f4), Some(f0), Some(f0), Some(f1), Some(f0),], p);
252
253 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 10);
254 assert_eq!(vec![Some(f4), Some(f0), Some(f1), Some(f0), Some(f1),], p);
255
256 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 15);
257 assert_eq!(vec![Some(f4), Some(f1), Some(f1), Some(f1), Some(f1),], p);
258
259 // let p = SC::point(vec![Fr::from(4_u32)], false, 4, 16); // TODO expect error
260 }
261
262 #[test]
263 fn test_partial_evaluate() {
264 // g(X_0, X_1, X_2) = 2 X_0^3 + X_0 X_2 + X_1 X_2
265 let terms = vec![
266 (Fr::from(2u32), SparseTerm::new(vec![(0_usize, 3)])),
267 (
268 Fr::from(1u32),
269 SparseTerm::new(vec![(0_usize, 1), (2_usize, 1)]),
270 ),
271 (
272 Fr::from(1u32),
273 SparseTerm::new(vec![(1_usize, 1), (2_usize, 1)]),
274 ),
275 ];
276 let p = SparsePolynomial::from_coefficients_slice(3, &terms);
277
278 type SC = SumCheck<Fr, G1Projective, DensePolynomial<Fr>, SparsePolynomial<Fr, SparseTerm>>;
279
280 let e0 = SC::partial_evaluate(&p, &[Some(Fr::from(2_u32)), None, Some(Fr::from(0_u32))]);
281 assert_eq!(e0.coeffs(), vec![Fr::from(16_u32)]);
282 let e1 = SC::partial_evaluate(&p, &[Some(Fr::from(2_u32)), None, Some(Fr::from(1_u32))]);
283 assert_eq!(e1.coeffs(), vec![Fr::from(18_u32), Fr::from(1)]);
284
285 assert_eq!((e0 + e1).coeffs(), vec![Fr::from(34_u32), Fr::from(1)]);
286 }
287
288 #[test]
289 fn test_flow_hardcoded_values() {
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}