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