1use ark_ec::{
19 models::short_weierstrass::Affine as SWJAffine, short_weierstrass::SWCurveConfig,
20 AdditiveGroup, AffineRepr, CurveGroup,
21};
22use ark_ff::{BitIteratorBE, Field, One, PrimeField, Zero};
23use itertools::Itertools;
24use mina_poseidon::sponge::ScalarChallenge;
25use rayon::prelude::*;
26use std::ops::AddAssign;
27
28fn add_pairs_in_place<P: SWCurveConfig>(pairs: &mut Vec<SWJAffine<P>>) {
29 let len = if pairs.len().is_multiple_of(2) {
30 pairs.len()
31 } else {
32 pairs.len() - 1
33 };
34 let mut denominators = pairs
35 .chunks_exact_mut(2)
36 .map(|p| {
37 if p[0].x == p[1].x {
38 if p[1].y.is_zero() {
39 P::BaseField::one()
40 } else {
41 p[1].y.double()
42 }
43 } else {
44 p[0].x - p[1].x
45 }
46 })
47 .collect::<Vec<_>>();
48
49 ark_ff::batch_inversion::<P::BaseField>(&mut denominators);
50
51 for (i, d) in (0..len).step_by(2).zip(denominators.iter()) {
52 let j = i / 2;
53 if pairs[i + 1].is_zero() {
54 pairs[j] = pairs[i];
55 } else if pairs[i].is_zero() {
56 pairs[j] = pairs[i + 1];
57 } else if pairs[i + 1].x == pairs[i].x
58 && (pairs[i + 1].y != pairs[i].y || pairs[i + 1].y.is_zero())
59 {
60 pairs[j] = SWJAffine::<P>::zero();
61 } else if pairs[i + 1].x == pairs[i].x && pairs[i + 1].y == pairs[i].y {
62 let sq = pairs[i].x.square();
63 let s = (sq.double() + sq + P::COEFF_A) * d;
64 let x = s.square() - pairs[i].x.double();
65 let y = -pairs[i].y - (s * (x - pairs[i].x));
66 pairs[j].x = x;
67 pairs[j].y = y;
68 } else {
69 let s = (pairs[i].y - pairs[i + 1].y) * d;
70 let x = s.square() - pairs[i].x - pairs[i + 1].x;
71 let y = -pairs[i].y - (s * (x - pairs[i].x));
72 pairs[j].x = x;
73 pairs[j].y = y;
74 }
75 }
76
77 let len = pairs.len();
78 if len % 2 == 1 {
79 pairs[len / 2] = pairs[len - 1];
80 pairs.truncate(len / 2 + 1);
81 } else {
82 pairs.truncate(len / 2);
83 }
84}
85
86fn batch_add_assign_no_branch<P: SWCurveConfig>(
91 denominators: &mut [P::BaseField],
92 v0: &mut [SWJAffine<P>],
93 v1: &[SWJAffine<P>],
94) {
95 denominators
96 .par_iter_mut()
97 .enumerate()
98 .for_each(|(i, denom)| {
99 let p0 = v0[i];
100 let p1 = v1[i];
101 let d = p0.x - p1.x;
102 *denom = d;
103 });
104
105 ark_ff::batch_inversion::<P::BaseField>(denominators);
106
107 denominators
108 .par_iter()
109 .zip(v0.par_iter_mut())
110 .zip(v1.par_iter())
111 .for_each(|((d, p0), p1)| {
112 let s = (p0.y - p1.y) * d;
113 let x = s.square() - p0.x - p1.x;
114 let y = -p0.y - (s * (x - p0.x));
115 p0.x = x;
116 p0.y = y;
117 });
118}
119
120pub fn batch_add_assign<P: SWCurveConfig>(
122 denominators: &mut [P::BaseField],
123 v0: &mut [SWJAffine<P>],
124 v1: &[SWJAffine<P>],
125) {
126 denominators
127 .par_iter_mut()
128 .zip(v0.par_iter())
129 .zip(v1.par_iter())
130 .for_each(|((denom, p0), p1)| {
131 let d = if p0.x == p1.x {
132 if p1.y.is_zero() {
133 P::BaseField::one()
134 } else {
135 p1.y.double()
136 }
137 } else {
138 p0.x - p1.x
139 };
140 *denom = d;
141 });
142
143 ark_ff::batch_inversion::<P::BaseField>(denominators);
144
145 denominators
146 .par_iter()
147 .zip(v0.par_iter_mut())
148 .zip(v1.par_iter())
149 .for_each(|((d, p0), p1)| {
150 if p1.is_zero() {
151 } else if p0.is_zero() {
152 *p0 = *p1;
153 } else if p1.x == p0.x && (p1.y != p0.y || p1.y == P::BaseField::zero()) {
154 *p0 = SWJAffine::<P>::zero();
155 } else if p1.x == p0.x && p1.y == p0.y {
156 let sq = p0.x.square();
157 let s = (sq.double() + sq + P::COEFF_A) * d;
158 let x = s.square() - p0.x.double();
159 let y = -p0.y - (s * (x - p0.x));
160 p0.x = x;
161 p0.y = y;
162 } else {
163 let s = (p0.y - p1.y) * d;
164 let x = s.square() - p0.x - p1.x;
165 let y = -p0.y - (s * (x - p0.x));
166 p0.x = x;
167 p0.y = y;
168 }
169 });
170}
171
172fn affine_window_combine_base<P: SWCurveConfig>(
173 g1: &[SWJAffine<P>],
174 g2: &[SWJAffine<P>],
175 x1: P::ScalarField,
176 x2: P::ScalarField,
177) -> Vec<SWJAffine<P>> {
178 let g1g2 = {
179 let mut v: Vec<_> = (0..2 * g1.len())
180 .map(|i| {
181 let j = i / 2;
182 if i % 2 == 0 {
183 g1[j]
184 } else {
185 g2[j]
186 }
187 })
188 .collect();
189 add_pairs_in_place(&mut v);
190 v
191 };
192 assert!(g1g2.len() == g1.len());
193
194 let windows1 = BitIteratorBE::new(x1.into_bigint()).tuples();
195 let windows2 = BitIteratorBE::new(x2.into_bigint()).tuples();
196
197 let mut points = vec![SWJAffine::<P>::zero(); g1.len()];
198
199 let mut denominators = vec![P::BaseField::zero(); g1.len()];
200
201 let [g01_00, g10_00, g11_00, g00_01, g01_01, g10_01, g11_01, g00_10, g01_10, g10_10, g11_10, g00_11, g01_11, g10_11, g11_11] =
202 affine_shamir_window_table(&mut denominators, g1, g2);
203
204 for ((hi_1, lo_1), (hi_2, lo_2)) in windows1.zip(windows2) {
205 for _ in 0..2 {
207 for i in 0..g1.len() {
208 denominators[i] = points[i].y.double();
209 }
210 ark_ff::batch_inversion::<P::BaseField>(&mut denominators);
211
212 for i in 0..g1.len() {
214 let d = denominators[i];
215 let sq = points[i].x.square();
216 let s = (sq.double() + sq + P::COEFF_A) * d;
217 let x = s.square() - points[i].x.double();
218 let y = -points[i].y - (s * (x - points[i].x));
219 points[i].x = x;
220 points[i].y = y;
221 }
222 }
223
224 match ((hi_1, lo_1), (hi_2, lo_2)) {
225 ((false, false), (false, false)) => (),
226 ((false, true), (false, false)) => {
227 batch_add_assign(&mut denominators, &mut points, &g01_00)
228 }
229 ((true, false), (false, false)) => {
230 batch_add_assign(&mut denominators, &mut points, &g10_00)
231 }
232 ((true, true), (false, false)) => {
233 batch_add_assign(&mut denominators, &mut points, &g11_00)
234 }
235
236 ((false, false), (false, true)) => {
237 batch_add_assign(&mut denominators, &mut points, &g00_01)
238 }
239 ((false, true), (false, true)) => {
240 batch_add_assign(&mut denominators, &mut points, &g01_01)
241 }
242 ((true, false), (false, true)) => {
243 batch_add_assign(&mut denominators, &mut points, &g10_01)
244 }
245 ((true, true), (false, true)) => {
246 batch_add_assign(&mut denominators, &mut points, &g11_01)
247 }
248
249 ((false, false), (true, false)) => {
250 batch_add_assign(&mut denominators, &mut points, &g00_10)
251 }
252 ((false, true), (true, false)) => {
253 batch_add_assign(&mut denominators, &mut points, &g01_10)
254 }
255 ((true, false), (true, false)) => {
256 batch_add_assign(&mut denominators, &mut points, &g10_10)
257 }
258 ((true, true), (true, false)) => {
259 batch_add_assign(&mut denominators, &mut points, &g11_10)
260 }
261
262 ((false, false), (true, true)) => {
263 batch_add_assign(&mut denominators, &mut points, &g00_11)
264 }
265 ((false, true), (true, true)) => {
266 batch_add_assign(&mut denominators, &mut points, &g01_11)
267 }
268 ((true, false), (true, true)) => {
269 batch_add_assign(&mut denominators, &mut points, &g10_11)
270 }
271 ((true, true), (true, true)) => {
272 batch_add_assign(&mut denominators, &mut points, &g11_11)
273 }
274 }
275 }
276 points
277}
278
279fn batch_endo_in_place<P: SWCurveConfig>(endo_coeff: P::BaseField, ps: &mut [SWJAffine<P>]) {
280 ps.par_iter_mut().for_each(|p| p.x *= endo_coeff);
281}
282
283fn batch_negate_in_place<P: SWCurveConfig>(ps: &mut [SWJAffine<P>]) {
284 ps.par_iter_mut().for_each(|p| {
285 p.y = -p.y;
286 });
287}
288
289fn affine_window_combine_one_endo_base<P: SWCurveConfig>(
293 endo_coeff: P::BaseField,
294 g1: &[SWJAffine<P>],
295 g2: &[SWJAffine<P>],
296 chal: ScalarChallenge<P::ScalarField>,
297) -> Vec<SWJAffine<P>> {
298 fn assign<A: Copy>(dst: &mut [A], src: &[A]) {
299 let n = dst.len();
300 dst[..n].clone_from_slice(&src[..n]);
301 }
302
303 fn get_bit(limbs_lsb: &[u64], i: u64) -> u64 {
304 let limb = i / 64;
305 let j = i % 64;
306 (limbs_lsb[limb as usize] >> j) & 1
307 }
308
309 let rep = chal.inner().into_bigint();
310 let r = rep.as_ref();
311
312 let mut denominators = vec![P::BaseField::zero(); g1.len()];
313 let mut points = g2.to_vec();
315 batch_endo_in_place(endo_coeff, &mut points);
316 batch_add_assign_no_branch(&mut denominators, &mut points, g2);
317 batch_double_in_place(&mut denominators, &mut points);
318
319 let mut tmp_s = g2.to_vec();
320 let mut tmp_acc = g2.to_vec();
321 for i in (0..(128 / 2)).rev() {
322 assign(&mut tmp_s, g2);
324 assign(&mut tmp_acc, &points);
326
327 let r_2i = get_bit(r, 2 * i);
328 if r_2i == 0 {
329 batch_negate_in_place(&mut tmp_s);
330 }
331 if get_bit(r, 2 * i + 1) == 1 {
332 batch_endo_in_place(endo_coeff, &mut tmp_s);
333 }
334
335 batch_add_assign_no_branch(&mut denominators, &mut points, &tmp_s);
337 batch_add_assign_no_branch(&mut denominators, &mut points, &tmp_acc);
338 }
339 batch_add_assign(&mut denominators, &mut points, g1);
341 points
342}
343
344fn batch_double_in_place<P: SWCurveConfig>(
346 denominators: &mut Vec<P::BaseField>,
347 points: &mut [SWJAffine<P>],
348) {
349 denominators
350 .par_iter_mut()
351 .zip(points.par_iter())
352 .for_each(|(d, p)| {
353 *d = p.y.double();
354 });
355 ark_ff::batch_inversion::<P::BaseField>(denominators);
356
357 denominators
359 .par_iter()
360 .zip(points.par_iter_mut())
361 .for_each(|(d, p)| {
362 let sq = p.x.square();
363 let s = (sq.double() + sq + P::COEFF_A) * d;
364 let x = s.square() - p.x.double();
365 let y = -p.y - (s * (x - p.x));
366 p.x = x;
367 p.y = y;
368 });
369}
370
371fn affine_window_combine_one_base<P: SWCurveConfig>(
372 g1: &[SWJAffine<P>],
373 g2: &[SWJAffine<P>],
374 x2: P::ScalarField,
375) -> Vec<SWJAffine<P>> {
376 let windows2 = BitIteratorBE::new(x2.into_bigint()).tuples();
377
378 let mut points = vec![SWJAffine::<P>::zero(); g1.len()];
379
380 let mut denominators = vec![P::BaseField::zero(); g1.len()];
381
382 let [g01, g10, g11] = affine_shamir_window_table_one(&mut denominators, g2);
383
384 for (hi_2, lo_2) in windows2 {
385 for _ in 0..2 {
387 for i in 0..g1.len() {
388 denominators[i] = points[i].y.double();
389 }
390 ark_ff::batch_inversion::<P::BaseField>(&mut denominators);
391
392 for i in 0..g1.len() {
394 let d = denominators[i];
395 let sq = points[i].x.square();
396 let s = (sq.double() + sq + P::COEFF_A) * d;
397 let x = s.square() - points[i].x.double();
398 let y = -points[i].y - (s * (x - points[i].x));
399 points[i].x = x;
400 points[i].y = y;
401 }
402 }
403
404 match (hi_2, lo_2) {
405 (false, false) => (),
406 (false, true) => batch_add_assign(&mut denominators, &mut points, &g01),
407 (true, false) => batch_add_assign(&mut denominators, &mut points, &g10),
408 (true, true) => batch_add_assign(&mut denominators, &mut points, &g11),
409 }
410 }
411
412 batch_add_assign(&mut denominators, &mut points, g1);
413
414 points
415}
416
417pub fn affine_window_combine<P: SWCurveConfig>(
418 g1: &[SWJAffine<P>],
419 g2: &[SWJAffine<P>],
420 x1: P::ScalarField,
421 x2: P::ScalarField,
422) -> Vec<SWJAffine<P>> {
423 const CHUNK_SIZE: usize = 10_000;
424 let b: Vec<_> = g1.chunks(CHUNK_SIZE).zip(g2.chunks(CHUNK_SIZE)).collect();
425 let v: Vec<_> = b
426 .into_par_iter()
427 .map(|(v1, v2)| affine_window_combine_base(v1, v2, x1, x2))
428 .collect();
429 v.concat()
430}
431
432pub fn affine_window_combine_one_endo<P: SWCurveConfig>(
437 endo_coeff: P::BaseField,
438 g1: &[SWJAffine<P>],
439 g2: &[SWJAffine<P>],
440 chal: ScalarChallenge<P::ScalarField>,
441) -> Vec<SWJAffine<P>> {
442 const CHUNK_SIZE: usize = 4096;
443 let b: Vec<_> = g1.chunks(CHUNK_SIZE).zip(g2.chunks(CHUNK_SIZE)).collect();
444 let v: Vec<_> = b
445 .into_par_iter()
446 .map(|(v1, v2)| affine_window_combine_one_endo_base(endo_coeff, v1, v2, chal.clone()))
447 .collect();
448 v.concat()
449}
450pub fn affine_window_combine_one<P: SWCurveConfig>(
451 g1: &[SWJAffine<P>],
452 g2: &[SWJAffine<P>],
453 x2: P::ScalarField,
454) -> Vec<SWJAffine<P>> {
455 const CHUNK_SIZE: usize = 10_000;
456 let b: Vec<_> = g1.chunks(CHUNK_SIZE).zip(g2.chunks(CHUNK_SIZE)).collect();
457 let v: Vec<_> = b
458 .into_par_iter()
459 .map(|(v1, v2)| affine_window_combine_one_base(v1, v2, x2))
460 .collect();
461 v.concat()
462}
463
464pub fn window_combine<G: AffineRepr>(
465 g_lo: &[G],
466 g_hi: &[G],
467 x_lo: G::ScalarField,
468 x_hi: G::ScalarField,
469) -> Vec<G> {
470 let mut g_proj: Vec<G::Group> = {
471 let pairs: Vec<_> = g_lo.iter().zip(g_hi).collect();
472 pairs
473 .into_par_iter()
474 .map(|(lo, hi)| window_shamir::<G>(x_lo, *lo, x_hi, *hi))
475 .collect()
476 };
477 G::Group::normalize_batch(g_proj.as_mut_slice())
478}
479
480pub fn affine_shamir_window_table<P: SWCurveConfig>(
481 denominators: &mut [P::BaseField],
482 g1: &[SWJAffine<P>],
483 g2: &[SWJAffine<P>],
484) -> [Vec<SWJAffine<P>>; 15] {
485 fn assign<A: Copy>(dst: &mut [A], src: &[A]) {
486 let n = dst.len();
487 dst[..n].clone_from_slice(&src[..n]);
488 }
489
490 let n = g1.len();
491
492 let mut res: [Vec<_>; 15] = [
493 vec![SWJAffine::<P>::zero(); n],
494 vec![SWJAffine::<P>::zero(); n],
495 vec![SWJAffine::<P>::zero(); n],
496 vec![SWJAffine::<P>::zero(); n],
497 vec![SWJAffine::<P>::zero(); n],
498 vec![SWJAffine::<P>::zero(); n],
499 vec![SWJAffine::<P>::zero(); n],
500 vec![SWJAffine::<P>::zero(); n],
501 vec![SWJAffine::<P>::zero(); n],
502 vec![SWJAffine::<P>::zero(); n],
503 vec![SWJAffine::<P>::zero(); n],
504 vec![SWJAffine::<P>::zero(); n],
505 vec![SWJAffine::<P>::zero(); n],
506 vec![SWJAffine::<P>::zero(); n],
507 vec![SWJAffine::<P>::zero(); n],
508 ];
509
510 let [g01_00, g10_00, g11_00, g00_01, g01_01, g10_01, g11_01, g00_10, g01_10, g10_10, g11_10, g00_11, g01_11, g10_11, g11_11] =
511 &mut res;
512
513 assign(g01_00, g1);
514
515 assign(g10_00, g1);
516 batch_add_assign(denominators, g10_00, g1);
517
518 assign(g11_00, g10_00);
519 batch_add_assign(denominators, g11_00, g1);
520
521 assign(g00_01, g2);
522
523 assign(g01_01, g00_01);
524 batch_add_assign(denominators, g01_01, g1);
525
526 assign(g10_01, g01_01);
527 batch_add_assign(denominators, g10_01, g1);
528
529 assign(g11_01, g10_01);
530 batch_add_assign(denominators, g11_01, g1);
531
532 assign(g00_10, g00_01);
533 batch_add_assign(denominators, g00_10, g2);
534
535 assign(g01_10, g00_10);
536 batch_add_assign(denominators, g01_10, g1);
537
538 assign(g10_10, g01_10);
539 batch_add_assign(denominators, g10_10, g1);
540
541 assign(g11_10, g10_10);
542 batch_add_assign(denominators, g11_10, g1);
543
544 assign(g00_11, g00_10);
545 batch_add_assign(denominators, g00_11, g2);
546
547 assign(g01_11, g00_11);
548 batch_add_assign(denominators, g01_11, g1);
549
550 assign(g10_11, g01_11);
551 batch_add_assign(denominators, g10_11, g1);
552
553 assign(g11_11, g10_11);
554 batch_add_assign(denominators, g11_11, g1);
555
556 res
557}
558
559pub fn affine_shamir_window_table_one<P: SWCurveConfig>(
560 denominators: &mut [P::BaseField],
561 g1: &[SWJAffine<P>],
562) -> [Vec<SWJAffine<P>>; 3] {
563 fn assign<A: Copy>(dst: &mut [A], src: &[A]) {
564 let n = dst.len();
565 dst[..n].clone_from_slice(&src[..n]);
566 }
567
568 let n = g1.len();
569
570 let mut res: [Vec<_>; 3] = [
571 vec![SWJAffine::<P>::zero(); n],
572 vec![SWJAffine::<P>::zero(); n],
573 vec![SWJAffine::<P>::zero(); n],
574 ];
575
576 let [g01, g10, g11] = &mut res;
577
578 assign(g01, g1);
579
580 assign(g10, g1);
581 batch_add_assign(denominators, g10, g1);
582
583 assign(g11, g10);
584 batch_add_assign(denominators, g11, g1);
585
586 res
587}
588
589fn window_shamir<G: AffineRepr>(x1: G::ScalarField, g1: G, x2: G::ScalarField, g2: G) -> G::Group {
590 let [_g00_00, g01_00, g10_00, g11_00, g00_01, g01_01, g10_01, g11_01, g00_10, g01_10, g10_10, g11_10, g00_11, g01_11, g10_11, g11_11] =
591 shamir_window_table(g1, g2);
592
593 let windows1 = BitIteratorBE::new(x1.into_bigint()).tuples();
594 let windows2 = BitIteratorBE::new(x2.into_bigint()).tuples();
595
596 let mut res = G::Group::zero();
597
598 for ((hi_1, lo_1), (hi_2, lo_2)) in windows1.zip(windows2) {
599 res.double_in_place();
600 res.double_in_place();
601 match ((hi_1, lo_1), (hi_2, lo_2)) {
602 ((false, false), (false, false)) => (),
603 ((false, true), (false, false)) => res.add_assign(&g01_00),
604 ((true, false), (false, false)) => res.add_assign(&g10_00),
605 ((true, true), (false, false)) => res.add_assign(&g11_00),
606
607 ((false, false), (false, true)) => res.add_assign(&g00_01),
608 ((false, true), (false, true)) => res.add_assign(&g01_01),
609 ((true, false), (false, true)) => res.add_assign(&g10_01),
610 ((true, true), (false, true)) => res.add_assign(&g11_01),
611
612 ((false, false), (true, false)) => res.add_assign(&g00_10),
613 ((false, true), (true, false)) => res.add_assign(&g01_10),
614 ((true, false), (true, false)) => res.add_assign(&g10_10),
615 ((true, true), (true, false)) => res.add_assign(&g11_10),
616
617 ((false, false), (true, true)) => res.add_assign(&g00_11),
618 ((false, true), (true, true)) => res.add_assign(&g01_11),
619 ((true, false), (true, true)) => res.add_assign(&g10_11),
620 ((true, true), (true, true)) => res.add_assign(&g11_11),
621 }
622 }
623
624 res
625}
626
627pub fn shamir_window_table<G: AffineRepr>(g1: G, g2: G) -> [G; 16] {
628 let g00_00 = G::generator().into_group();
629 let g01_00 = g1.into_group();
630 let g10_00 = {
631 let mut g = g01_00;
632 g.add_assign(&g1);
633 g
634 };
635 let g11_00 = {
636 let mut g = g10_00;
637 g.add_assign(&g1);
638 g
639 };
640
641 let g00_01 = g2.into_group();
642 let g01_01 = {
643 let mut g = g00_01;
644 g.add_assign(&g1);
645 g
646 };
647 let g10_01 = {
648 let mut g = g01_01;
649 g.add_assign(&g1);
650 g
651 };
652 let g11_01 = {
653 let mut g = g10_01;
654 g.add_assign(&g1);
655 g
656 };
657
658 let g00_10 = {
659 let mut g = g00_01;
660 g.add_assign(&g2);
661 g
662 };
663 let g01_10 = {
664 let mut g = g00_10;
665 g.add_assign(&g1);
666 g
667 };
668 let g10_10 = {
669 let mut g = g01_10;
670 g.add_assign(&g1);
671 g
672 };
673 let g11_10 = {
674 let mut g = g10_10;
675 g.add_assign(&g1);
676 g
677 };
678 let g00_11 = {
679 let mut g = g00_10;
680 g.add_assign(&g2);
681 g
682 };
683 let g01_11 = {
684 let mut g = g00_11;
685 g.add_assign(&g1);
686 g
687 };
688 let g10_11 = {
689 let mut g = g01_11;
690 g.add_assign(&g1);
691 g
692 };
693 let g11_11 = {
694 let mut g = g10_11;
695 g.add_assign(&g1);
696 g
697 };
698
699 let mut v = vec![
700 g00_00, g01_00, g10_00, g11_00, g00_01, g01_01, g10_01, g11_01, g00_10, g01_10, g10_10,
701 g11_10, g00_11, g01_11, g10_11, g11_11,
702 ];
703 let v: Vec<_> = G::Group::normalize_batch(v.as_mut_slice());
704 [
705 v[0], v[1], v[2], v[3], v[4], v[5], v[6], v[7], v[8], v[9], v[10], v[11], v[12], v[13],
706 v[14], v[15],
707 ]
708}