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