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 might 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
17pub struct SumCheck<
18 F: PrimeField,
19 UV: Polynomial<F> + DenseUVPolynomial<F>,
20 MV: Polynomial<F> + DenseMVPolynomial<F>,
21> {
22 _f: PhantomData<F>,
23 _uv: PhantomData<UV>,
24 _mv: PhantomData<MV>,
25}
26
27impl<
28 F: PrimeField,
29 UV: Polynomial<F> + DenseUVPolynomial<F>,
30 MV: Polynomial<F> + DenseMVPolynomial<F>,
31 > SumCheck<F, UV, MV>
32{
33 fn partial_evaluate(g: &MV, point: &[Option<F>]) -> UV {
34 assert!(point.len() >= g.num_vars(), "Invalid evaluation domain");
35 // TODO: add check: there can only be 1 'None' value in point
36 if g.is_zero() {
37 return UV::from_coefficients_vec(vec![F::zero()]);
38 }
39
40 // note: this can be parallelized with cfg_into_iter
41 let mut univ_terms: Vec<(F, SparseTerm)> = vec![];
42 for (coef, term) in g.terms().iter() {
43 // partial_evaluate each term
44 let mut new_coef = F::one();
45 let mut new_term = Vec::new();
46 for (var, power) in term.iter() {
47 match point[*var] {
48 Some(v) => {
49 if v.is_zero() {
50 new_coef = F::zero();
51 new_term = vec![];
52 break;
53 } else {
54 new_coef = new_coef * v.pow([(*power) as u64]);
55 }
56 }
57 _ => {
58 new_term.push((*var, *power));
59 }
60 };
61 }
62 let new_term = SparseTerm::new(new_term);
63 let new_coef = new_coef * coef;
64 univ_terms.push((new_coef, new_term));
65 }
66 let mv_poly: SparsePolynomial<F, SparseTerm> =
67 DenseMVPolynomial::<F>::from_coefficients_vec(g.num_vars(), univ_terms.clone());
68
69 let mut univ_coeffs: Vec<F> = vec![F::zero(); mv_poly.degree() + 1];
70 for (coef, term) in univ_terms {
71 if term.is_empty() {
72 univ_coeffs[0] += coef;
73 continue;
74 }
75 for (_, power) in term.iter() {
76 univ_coeffs[*power] += coef;
77 }
78 }
79 UV::from_coefficients_vec(univ_coeffs)
80 }
81
82 fn point_complete(challenges: Vec<F>, n_elems: usize, iter_num: usize) -> Vec<F> {
83 let p = Self::point(challenges, false, n_elems, iter_num);
84 // let mut r = Vec::new();
85 let mut r = vec![F::zero(); n_elems];
86 for i in 0..n_elems {
87 // r.push(p[i].unwrap());
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");
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 let v = g.num_vars();
129
130 let r = vec![F::from(2_u32), F::from(3_u32), F::from(6_u32)]; // TMP will come from transcript
131
132 // compute H
133 let mut H = F::zero();
134 for i in 0..(2_u64.pow(v as u32) as usize) {
135 let p = Self::point_complete(vec![], v, i);
136
137 H = H + g.evaluate(&p.into());
138 }
139
140 let mut ss: Vec<UV> = Vec::new();
141 for i in 0..v {
142 let var_slots = v - 1 - i;
143 let n_points = 2_u64.pow(var_slots as u32) as usize;
144 let mut s_i = UV::zero();
145 let r_round = r.as_slice()[..i].to_vec();
146 for j in 0..n_points {
147 let point = Self::point(r_round.clone(), true, v, j);
148 s_i = s_i + Self::partial_evaluate(&g, &point);
149 }
150 ss.push(s_i);
151 }
152
153 let point_last = r;
154 let last_g_eval = g.evaluate(&point_last.into());
155 (H, ss, last_g_eval)
156 }
157
158 pub fn verify(proof: (F, Vec<UV>, F)) -> bool {
159 // let c: F, ss: Vec<UV>;
160 let (c, ss, last_g_eval) = proof;
161
162 let r = vec![F::from(2_u32), F::from(3_u32), F::from(6_u32)]; // TMP will come from transcript
163
164 for (i, s) in ss.iter().enumerate() {
165 // TODO check degree
166 if i == 0 {
167 if c != s.evaluate(&F::zero()) + s.evaluate(&F::one()) {
168 return false;
169 }
170 continue;
171 }
172
173 if ss[i - 1].evaluate(&r[i - 1]) != s.evaluate(&F::zero()) + s.evaluate(&F::one()) {
174 return false;
175 }
176 }
177 // last round
178 if ss[ss.len() - 1].evaluate(&r[r.len() - 1]) != last_g_eval {
179 return false;
180 }
181
182 true
183 }
184}
185
186#[cfg(test)]
187mod tests {
188 use super::*;
189 use ark_bn254::Fr; // scalar field
190
191 #[test]
192 fn test_new_point() {
193 let f4 = Fr::from(4_u32);
194 let f1 = Fr::from(1);
195 let f0 = Fr::from(0);
196 type SC = SumCheck<Fr, DensePolynomial<Fr>, SparsePolynomial<Fr, SparseTerm>>;
197
198 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 0);
199 assert_eq!(vec![Some(f4), None, Some(f0), Some(f0), Some(f0),], p);
200
201 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 1);
202 assert_eq!(vec![Some(f4), None, Some(f1), Some(f0), Some(f0),], p);
203
204 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 2);
205 assert_eq!(vec![Some(f4), None, Some(f0), Some(f1), Some(f0),], p);
206
207 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 3);
208 assert_eq!(vec![Some(f4), None, Some(f1), Some(f1), Some(f0),], p);
209
210 let p = SC::point(vec![Fr::from(4_u32)], true, 5, 4);
211 assert_eq!(vec![Some(f4), None, Some(f0), Some(f0), Some(f1),], p);
212
213 // without None
214 let p = SC::point(vec![], false, 4, 0);
215 assert_eq!(vec![Some(f0), Some(f0), Some(f0), Some(f0),], p);
216
217 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 0);
218 assert_eq!(vec![Some(f4), Some(f0), Some(f0), Some(f0), Some(f0),], p);
219
220 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 1);
221 assert_eq!(vec![Some(f4), Some(f1), Some(f0), Some(f0), Some(f0),], p);
222
223 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 3);
224 assert_eq!(vec![Some(f4), Some(f1), Some(f1), Some(f0), Some(f0),], p);
225
226 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 4);
227 assert_eq!(vec![Some(f4), Some(f0), Some(f0), Some(f1), Some(f0),], p);
228
229 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 10);
230 assert_eq!(vec![Some(f4), Some(f0), Some(f1), Some(f0), Some(f1),], p);
231
232 let p = SC::point(vec![Fr::from(4_u32)], false, 5, 15);
233 assert_eq!(vec![Some(f4), Some(f1), Some(f1), Some(f1), Some(f1),], p);
234
235 // let p = SC::point(vec![Fr::from(4_u32)], false, 4, 16); // TODO expect error
236 }
237
238 #[test]
239 fn test_partial_evaluate() {
240 // g(X_0, X_1, X_2) = 2 X_0^3 + X_0 X_2 + X_1 X_2
241 let terms = vec![
242 (Fr::from(2u32), SparseTerm::new(vec![(0_usize, 3)])),
243 (
244 Fr::from(1u32),
245 SparseTerm::new(vec![(0_usize, 1), (2_usize, 1)]),
246 ),
247 (
248 Fr::from(1u32),
249 SparseTerm::new(vec![(1_usize, 1), (2_usize, 1)]),
250 ),
251 ];
252 let p = SparsePolynomial::from_coefficients_slice(3, &terms);
253
254 type SC = SumCheck<Fr, DensePolynomial<Fr>, SparsePolynomial<Fr, SparseTerm>>;
255
256 let e0 = SC::partial_evaluate(&p, &[Some(Fr::from(2_u32)), None, Some(Fr::from(0_u32))]);
257 assert_eq!(e0.coeffs(), vec![Fr::from(16_u32)]);
258 let e1 = SC::partial_evaluate(&p, &[Some(Fr::from(2_u32)), None, Some(Fr::from(1_u32))]);
259 assert_eq!(e1.coeffs(), vec![Fr::from(18_u32), Fr::from(1)]);
260
261 assert_eq!((e0 + e1).coeffs(), vec![Fr::from(34_u32), Fr::from(1)]);
262 }
263
264 #[test]
265 fn test_flow_hardcoded_values() {
266 let mut rng = ark_std::test_rng();
267 // let p = SparsePolynomial::<Fr, SparseTerm>::rand(deg, 3, &mut rng);
268 // let p = rand_poly(3, 3, &mut rng);
269
270 // g(X_0, X_1, X_2) = 2 X_0^3 + X_0 X_2 + X_1 X_2
271 let terms = vec![
272 (Fr::from(2u32), SparseTerm::new(vec![(0_usize, 3)])),
273 (
274 Fr::from(1u32),
275 SparseTerm::new(vec![(0_usize, 1), (2_usize, 1)]),
276 ),
277 (
278 Fr::from(1u32),
279 SparseTerm::new(vec![(1_usize, 1), (2_usize, 1)]),
280 ),
281 ];
282 let p = SparsePolynomial::from_coefficients_slice(3, &terms);
283 // println!("p {:?}", p);
284
285 type SC = SumCheck<Fr, DensePolynomial<Fr>, SparsePolynomial<Fr, SparseTerm>>;
286
287 let proof = SC::prove(p);
288 assert_eq!(proof.0, Fr::from(12_u32));
289 println!("proof {:?}", proof);
290
291 let v = SC::verify(proof);
292 assert!(v);
293 }
294
295 #[test]
296 fn test_flow_rng() {
297 let mut rng = ark_std::test_rng();
298 let p = SparsePolynomial::<Fr, SparseTerm>::rand(3, 3, &mut rng);
299 println!("p {:?}", p);
300
301 type SC = SumCheck<Fr, DensePolynomial<Fr>, SparsePolynomial<Fr, SparseTerm>>;
302
303 let proof = SC::prove(p);
304 println!("proof.s len {:?}", proof.1.len());
305
306 let v = SC::verify(proof);
307 assert!(v);
308 }
309}