testudo/src/product_tree.rs
src/product_tree.rsRaw
1#![allow(dead_code)]
2use super::dense_mlpoly::DensePolynomial;
3use super::dense_mlpoly::EqPolynomial;
4use super::scalar::Scalar;
5use super::sumcheck::SumcheckInstanceProof;
6use super::transcript::ProofTranscript;
7use merlin::Transcript;
8use serde::{Deserialize, Serialize};
9
10#[derive(Debug)]
11pub struct ProductCircuit {
12 left_vec: Vec<DensePolynomial>,
13 right_vec: Vec<DensePolynomial>,
14}
15
16impl ProductCircuit {
17 fn compute_layer(
18 inp_left: &DensePolynomial,
19 inp_right: &DensePolynomial,
20 ) -> (DensePolynomial, DensePolynomial) {
21 let len = inp_left.len() + inp_right.len();
22 let outp_left = (0..len / 4)
23 .map(|i| inp_left[i] * inp_right[i])
24 .collect::<Vec<Scalar>>();
25 let outp_right = (len / 4..len / 2)
26 .map(|i| inp_left[i] * inp_right[i])
27 .collect::<Vec<Scalar>>();
28
29 (
30 DensePolynomial::new(outp_left),
31 DensePolynomial::new(outp_right),
32 )
33 }
34
35 pub fn new(poly: &DensePolynomial) -> Self {
36 let mut left_vec: Vec<DensePolynomial> = Vec::new();
37 let mut right_vec: Vec<DensePolynomial> = Vec::new();
38
39 let num_layers = poly.len().log2() as usize;
40 let (outp_left, outp_right) = poly.split(poly.len() / 2);
41
42 left_vec.push(outp_left);
43 right_vec.push(outp_right);
44
45 for i in 0..num_layers - 1 {
46 let (outp_left, outp_right) = ProductCircuit::compute_layer(&left_vec[i], &right_vec[i]);
47 left_vec.push(outp_left);
48 right_vec.push(outp_right);
49 }
50
51 ProductCircuit {
52 left_vec,
53 right_vec,
54 }
55 }
56
57 pub fn evaluate(&self) -> Scalar {
58 let len = self.left_vec.len();
59 assert_eq!(self.left_vec[len - 1].get_num_vars(), 0);
60 assert_eq!(self.right_vec[len - 1].get_num_vars(), 0);
61 self.left_vec[len - 1][0] * self.right_vec[len - 1][0]
62 }
63}
64
65pub struct DotProductCircuit {
66 left: DensePolynomial,
67 right: DensePolynomial,
68 weight: DensePolynomial,
69}
70
71impl DotProductCircuit {
72 pub fn new(left: DensePolynomial, right: DensePolynomial, weight: DensePolynomial) -> Self {
73 assert_eq!(left.len(), right.len());
74 assert_eq!(left.len(), weight.len());
75 DotProductCircuit {
76 left,
77 right,
78 weight,
79 }
80 }
81
82 pub fn evaluate(&self) -> Scalar {
83 (0..self.left.len())
84 .map(|i| self.left[i] * self.right[i] * self.weight[i])
85 .sum()
86 }
87
88 pub fn split(&mut self) -> (DotProductCircuit, DotProductCircuit) {
89 let idx = self.left.len() / 2;
90 assert_eq!(idx * 2, self.left.len());
91 let (l1, l2) = self.left.split(idx);
92 let (r1, r2) = self.right.split(idx);
93 let (w1, w2) = self.weight.split(idx);
94 (
95 DotProductCircuit {
96 left: l1,
97 right: r1,
98 weight: w1,
99 },
100 DotProductCircuit {
101 left: l2,
102 right: r2,
103 weight: w2,
104 },
105 )
106 }
107}
108
109#[allow(dead_code)]
110#[derive(Debug, Serialize, Deserialize)]
111pub struct LayerProof {
112 pub proof: SumcheckInstanceProof,
113 pub claims: Vec<Scalar>,
114}
115
116#[allow(dead_code)]
117impl LayerProof {
118 pub fn verify(
119 &self,
120 claim: Scalar,
121 num_rounds: usize,
122 degree_bound: usize,
123 transcript: &mut Transcript,
124 ) -> (Scalar, Vec<Scalar>) {
125 self
126 .proof
127 .verify(claim, num_rounds, degree_bound, transcript)
128 .unwrap()
129 }
130}
131
132#[allow(dead_code)]
133#[derive(Debug, Serialize, Deserialize)]
134pub struct LayerProofBatched {
135 pub proof: SumcheckInstanceProof,
136 pub claims_prod_left: Vec<Scalar>,
137 pub claims_prod_right: Vec<Scalar>,
138}
139
140#[allow(dead_code)]
141impl LayerProofBatched {
142 pub fn verify(
143 &self,
144 claim: Scalar,
145 num_rounds: usize,
146 degree_bound: usize,
147 transcript: &mut Transcript,
148 ) -> (Scalar, Vec<Scalar>) {
149 self
150 .proof
151 .verify(claim, num_rounds, degree_bound, transcript)
152 .unwrap()
153 }
154}
155
156#[derive(Debug, Serialize, Deserialize)]
157pub struct ProductCircuitEvalProof {
158 proof: Vec<LayerProof>,
159}
160
161#[derive(Debug, Serialize, Deserialize)]
162pub struct ProductCircuitEvalProofBatched {
163 proof: Vec<LayerProofBatched>,
164 claims_dotp: (Vec<Scalar>, Vec<Scalar>, Vec<Scalar>),
165}
166
167impl ProductCircuitEvalProof {
168 #![allow(dead_code)]
169 pub fn prove(
170 circuit: &mut ProductCircuit,
171 transcript: &mut Transcript,
172 ) -> (Self, Scalar, Vec<Scalar>) {
173 let mut proof: Vec<LayerProof> = Vec::new();
174 let num_layers = circuit.left_vec.len();
175
176 let mut claim = circuit.evaluate();
177 let mut rand = Vec::new();
178 for layer_id in (0..num_layers).rev() {
179 let len = circuit.left_vec[layer_id].len() + circuit.right_vec[layer_id].len();
180
181 let mut poly_C = DensePolynomial::new(EqPolynomial::new(rand.clone()).evals());
182 assert_eq!(poly_C.len(), len / 2);
183
184 let num_rounds_prod = poly_C.len().log2() as usize;
185 let comb_func_prod = |poly_A_comp: &Scalar,
186 poly_B_comp: &Scalar,
187 poly_C_comp: &Scalar|
188 -> Scalar { poly_A_comp * poly_B_comp * poly_C_comp };
189 let (proof_prod, rand_prod, claims_prod) = SumcheckInstanceProof::prove_cubic(
190 &claim,
191 num_rounds_prod,
192 &mut circuit.left_vec[layer_id],
193 &mut circuit.right_vec[layer_id],
194 &mut poly_C,
195 comb_func_prod,
196 transcript,
197 );
198
199 transcript.append_scalar(b"claim_prod_left", &claims_prod[0]);
200 transcript.append_scalar(b"claim_prod_right", &claims_prod[1]);
201
202 // produce a random challenge
203 let r_layer = transcript.challenge_scalar(b"challenge_r_layer");
204 claim = claims_prod[0] + r_layer * (claims_prod[1] - claims_prod[0]);
205
206 let mut ext = vec![r_layer];
207 ext.extend(rand_prod);
208 rand = ext;
209
210 proof.push(LayerProof {
211 proof: proof_prod,
212 claims: claims_prod[0..claims_prod.len() - 1].to_vec(),
213 });
214 }
215
216 (ProductCircuitEvalProof { proof }, claim, rand)
217 }
218
219 pub fn verify(
220 &self,
221 eval: Scalar,
222 len: usize,
223 transcript: &mut Transcript,
224 ) -> (Scalar, Vec<Scalar>) {
225 let num_layers = len.log2() as usize;
226 let mut claim = eval;
227 let mut rand: Vec<Scalar> = Vec::new();
228 //let mut num_rounds = 0;
229 assert_eq!(self.proof.len(), num_layers);
230 for (num_rounds, i) in (0..num_layers).enumerate() {
231 let (claim_last, rand_prod) = self.proof[i].verify(claim, num_rounds, 3, transcript);
232
233 let claims_prod = &self.proof[i].claims;
234 transcript.append_scalar(b"claim_prod_left", &claims_prod[0]);
235 transcript.append_scalar(b"claim_prod_right", &claims_prod[1]);
236
237 assert_eq!(rand.len(), rand_prod.len());
238 let eq: Scalar = (0..rand.len())
239 .map(|i| {
240 rand[i] * rand_prod[i] + (Scalar::one() - rand[i]) * (Scalar::one() - rand_prod[i])
241 })
242 .product();
243 assert_eq!(claims_prod[0] * claims_prod[1] * eq, claim_last);
244
245 // produce a random challenge
246 let r_layer = transcript.challenge_scalar(b"challenge_r_layer");
247 claim = (Scalar::one() - r_layer) * claims_prod[0] + r_layer * claims_prod[1];
248 let mut ext = vec![r_layer];
249 ext.extend(rand_prod);
250 rand = ext;
251 }
252
253 (claim, rand)
254 }
255}
256
257impl ProductCircuitEvalProofBatched {
258 pub fn prove(
259 prod_circuit_vec: &mut Vec<&mut ProductCircuit>,
260 dotp_circuit_vec: &mut Vec<&mut DotProductCircuit>,
261 transcript: &mut Transcript,
262 ) -> (Self, Vec<Scalar>) {
263 assert!(!prod_circuit_vec.is_empty());
264
265 let mut claims_dotp_final = (Vec::new(), Vec::new(), Vec::new());
266
267 let mut proof_layers: Vec<LayerProofBatched> = Vec::new();
268 let num_layers = prod_circuit_vec[0].left_vec.len();
269 let mut claims_to_verify = (0..prod_circuit_vec.len())
270 .map(|i| prod_circuit_vec[i].evaluate())
271 .collect::<Vec<Scalar>>();
272 let mut rand = Vec::new();
273 for layer_id in (0..num_layers).rev() {
274 // prepare paralell instance that share poly_C first
275 let len = prod_circuit_vec[0].left_vec[layer_id].len()
276 + prod_circuit_vec[0].right_vec[layer_id].len();
277
278 let mut poly_C_par = DensePolynomial::new(EqPolynomial::new(rand.clone()).evals());
279 assert_eq!(poly_C_par.len(), len / 2);
280
281 let num_rounds_prod = poly_C_par.len().log2() as usize;
282 let comb_func_prod = |poly_A_comp: &Scalar,
283 poly_B_comp: &Scalar,
284 poly_C_comp: &Scalar|
285 -> Scalar { poly_A_comp * poly_B_comp * poly_C_comp };
286
287 let mut poly_A_batched_par: Vec<&mut DensePolynomial> = Vec::new();
288 let mut poly_B_batched_par: Vec<&mut DensePolynomial> = Vec::new();
289 for prod_circuit in prod_circuit_vec.iter_mut() {
290 poly_A_batched_par.push(&mut prod_circuit.left_vec[layer_id]);
291 poly_B_batched_par.push(&mut prod_circuit.right_vec[layer_id])
292 }
293 let poly_vec_par = (
294 &mut poly_A_batched_par,
295 &mut poly_B_batched_par,
296 &mut poly_C_par,
297 );
298
299 // prepare sequential instances that don't share poly_C
300 let mut poly_A_batched_seq: Vec<&mut DensePolynomial> = Vec::new();
301 let mut poly_B_batched_seq: Vec<&mut DensePolynomial> = Vec::new();
302 let mut poly_C_batched_seq: Vec<&mut DensePolynomial> = Vec::new();
303 if layer_id == 0 && !dotp_circuit_vec.is_empty() {
304 // add additional claims
305 for item in dotp_circuit_vec.iter() {
306 claims_to_verify.push(item.evaluate());
307 assert_eq!(len / 2, item.left.len());
308 assert_eq!(len / 2, item.right.len());
309 assert_eq!(len / 2, item.weight.len());
310 }
311
312 for dotp_circuit in dotp_circuit_vec.iter_mut() {
313 poly_A_batched_seq.push(&mut dotp_circuit.left);
314 poly_B_batched_seq.push(&mut dotp_circuit.right);
315 poly_C_batched_seq.push(&mut dotp_circuit.weight);
316 }
317 }
318 let poly_vec_seq = (
319 &mut poly_A_batched_seq,
320 &mut poly_B_batched_seq,
321 &mut poly_C_batched_seq,
322 );
323
324 // produce a fresh set of coeffs and a joint claim
325 let coeff_vec =
326 transcript.challenge_vector(b"rand_coeffs_next_layer", claims_to_verify.len());
327 let claim = (0..claims_to_verify.len())
328 .map(|i| claims_to_verify[i] * coeff_vec[i])
329 .sum();
330
331 let (proof, rand_prod, claims_prod, claims_dotp) = SumcheckInstanceProof::prove_cubic_batched(
332 &claim,
333 num_rounds_prod,
334 poly_vec_par,
335 poly_vec_seq,
336 &coeff_vec,
337 comb_func_prod,
338 transcript,
339 );
340
341 let (claims_prod_left, claims_prod_right, _claims_eq) = claims_prod;
342 for i in 0..prod_circuit_vec.len() {
343 transcript.append_scalar(b"claim_prod_left", &claims_prod_left[i]);
344 transcript.append_scalar(b"claim_prod_right", &claims_prod_right[i]);
345 }
346
347 if layer_id == 0 && !dotp_circuit_vec.is_empty() {
348 let (claims_dotp_left, claims_dotp_right, claims_dotp_weight) = claims_dotp;
349 for i in 0..dotp_circuit_vec.len() {
350 transcript.append_scalar(b"claim_dotp_left", &claims_dotp_left[i]);
351 transcript.append_scalar(b"claim_dotp_right", &claims_dotp_right[i]);
352 transcript.append_scalar(b"claim_dotp_weight", &claims_dotp_weight[i]);
353 }
354 claims_dotp_final = (claims_dotp_left, claims_dotp_right, claims_dotp_weight);
355 }
356
357 // produce a random challenge to condense two claims into a single claim
358 let r_layer = transcript.challenge_scalar(b"challenge_r_layer");
359
360 claims_to_verify = (0..prod_circuit_vec.len())
361 .map(|i| claims_prod_left[i] + r_layer * (claims_prod_right[i] - claims_prod_left[i]))
362 .collect::<Vec<Scalar>>();
363
364 let mut ext = vec![r_layer];
365 ext.extend(rand_prod);
366 rand = ext;
367
368 proof_layers.push(LayerProofBatched {
369 proof,
370 claims_prod_left,
371 claims_prod_right,
372 });
373 }
374
375 (
376 ProductCircuitEvalProofBatched {
377 proof: proof_layers,
378 claims_dotp: claims_dotp_final,
379 },
380 rand,
381 )
382 }
383
384 pub fn verify(
385 &self,
386 claims_prod_vec: &[Scalar],
387 claims_dotp_vec: &[Scalar],
388 len: usize,
389 transcript: &mut Transcript,
390 ) -> (Vec<Scalar>, Vec<Scalar>, Vec<Scalar>) {
391 let num_layers = len.log2() as usize;
392 let mut rand: Vec<Scalar> = Vec::new();
393 //let mut num_rounds = 0;
394 assert_eq!(self.proof.len(), num_layers);
395
396 let mut claims_to_verify = claims_prod_vec.to_owned();
397 let mut claims_to_verify_dotp: Vec<Scalar> = Vec::new();
398 for (num_rounds, i) in (0..num_layers).enumerate() {
399 if i == num_layers - 1 {
400 claims_to_verify.extend(claims_dotp_vec);
401 }
402
403 // produce random coefficients, one for each instance
404 let coeff_vec =
405 transcript.challenge_vector(b"rand_coeffs_next_layer", claims_to_verify.len());
406
407 // produce a joint claim
408 let claim = (0..claims_to_verify.len())
409 .map(|i| claims_to_verify[i] * coeff_vec[i])
410 .sum();
411
412 let (claim_last, rand_prod) = self.proof[i].verify(claim, num_rounds, 3, transcript);
413
414 let claims_prod_left = &self.proof[i].claims_prod_left;
415 let claims_prod_right = &self.proof[i].claims_prod_right;
416 assert_eq!(claims_prod_left.len(), claims_prod_vec.len());
417 assert_eq!(claims_prod_right.len(), claims_prod_vec.len());
418
419 for i in 0..claims_prod_vec.len() {
420 transcript.append_scalar(b"claim_prod_left", &claims_prod_left[i]);
421 transcript.append_scalar(b"claim_prod_right", &claims_prod_right[i]);
422 }
423
424 assert_eq!(rand.len(), rand_prod.len());
425 let eq: Scalar = (0..rand.len())
426 .map(|i| {
427 rand[i] * rand_prod[i] + (Scalar::one() - rand[i]) * (Scalar::one() - rand_prod[i])
428 })
429 .product();
430 let mut claim_expected: Scalar = (0..claims_prod_vec.len())
431 .map(|i| coeff_vec[i] * (claims_prod_left[i] * claims_prod_right[i] * eq))
432 .sum();
433
434 // add claims from the dotp instances
435 if i == num_layers - 1 {
436 let num_prod_instances = claims_prod_vec.len();
437 let (claims_dotp_left, claims_dotp_right, claims_dotp_weight) = &self.claims_dotp;
438 for i in 0..claims_dotp_left.len() {
439 transcript.append_scalar(b"claim_dotp_left", &claims_dotp_left[i]);
440 transcript.append_scalar(b"claim_dotp_right", &claims_dotp_right[i]);
441 transcript.append_scalar(b"claim_dotp_weight", &claims_dotp_weight[i]);
442
443 claim_expected += coeff_vec[i + num_prod_instances]
444 * claims_dotp_left[i]
445 * claims_dotp_right[i]
446 * claims_dotp_weight[i];
447 }
448 }
449
450 assert_eq!(claim_expected, claim_last);
451
452 // produce a random challenge
453 let r_layer = transcript.challenge_scalar(b"challenge_r_layer");
454
455 claims_to_verify = (0..claims_prod_left.len())
456 .map(|i| claims_prod_left[i] + r_layer * (claims_prod_right[i] - claims_prod_left[i]))
457 .collect::<Vec<Scalar>>();
458
459 // add claims to verify for dotp circuit
460 if i == num_layers - 1 {
461 let (claims_dotp_left, claims_dotp_right, claims_dotp_weight) = &self.claims_dotp;
462
463 for i in 0..claims_dotp_vec.len() / 2 {
464 // combine left claims
465 let claim_left = claims_dotp_left[2 * i]
466 + r_layer * (claims_dotp_left[2 * i + 1] - claims_dotp_left[2 * i]);
467
468 let claim_right = claims_dotp_right[2 * i]
469 + r_layer * (claims_dotp_right[2 * i + 1] - claims_dotp_right[2 * i]);
470
471 let claim_weight = claims_dotp_weight[2 * i]
472 + r_layer * (claims_dotp_weight[2 * i + 1] - claims_dotp_weight[2 * i]);
473 claims_to_verify_dotp.push(claim_left);
474 claims_to_verify_dotp.push(claim_right);
475 claims_to_verify_dotp.push(claim_weight);
476 }
477 }
478
479 let mut ext = vec![r_layer];
480 ext.extend(rand_prod);
481 rand = ext;
482 }
483 (claims_to_verify, claims_to_verify_dotp, rand)
484 }
485}