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) => {
409 batch_add_assign(&mut denominators, &mut points, &g11);
410 }
411 }
412 }
413
414 batch_add_assign(&mut denominators, &mut points, g1);
415
416 points
417}
418
419pub fn affine_window_combine<P: SWCurveConfig>(
420 g1: &[SWJAffine<P>],
421 g2: &[SWJAffine<P>],
422 x1: P::ScalarField,
423 x2: P::ScalarField,
424) -> Vec<SWJAffine<P>> {
425 const CHUNK_SIZE: usize = 10_000;
426 let b: Vec<_> = g1.chunks(CHUNK_SIZE).zip(g2.chunks(CHUNK_SIZE)).collect();
427 let v: Vec<_> = b
428 .into_par_iter()
429 .map(|(v1, v2)| affine_window_combine_base(v1, v2, x1, x2))
430 .collect();
431 v.concat()
432}
433
434pub fn affine_window_combine_one_endo<P: SWCurveConfig>(
439 endo_coeff: P::BaseField,
440 g1: &[SWJAffine<P>],
441 g2: &[SWJAffine<P>],
442 chal: &ScalarChallenge<P::ScalarField>,
443) -> Vec<SWJAffine<P>> {
444 const CHUNK_SIZE: usize = 4096;
445 let b: Vec<_> = g1.chunks(CHUNK_SIZE).zip(g2.chunks(CHUNK_SIZE)).collect();
446 let v: Vec<_> = b
447 .into_par_iter()
448 .map(|(v1, v2)| affine_window_combine_one_endo_base(endo_coeff, v1, v2, chal))
449 .collect();
450 v.concat()
451}
452pub fn affine_window_combine_one<P: SWCurveConfig>(
453 g1: &[SWJAffine<P>],
454 g2: &[SWJAffine<P>],
455 x2: P::ScalarField,
456) -> Vec<SWJAffine<P>> {
457 const CHUNK_SIZE: usize = 10_000;
458 let b: Vec<_> = g1.chunks(CHUNK_SIZE).zip(g2.chunks(CHUNK_SIZE)).collect();
459 let v: Vec<_> = b
460 .into_par_iter()
461 .map(|(v1, v2)| affine_window_combine_one_base(v1, v2, x2))
462 .collect();
463 v.concat()
464}
465
466pub fn window_combine<G: AffineRepr>(
467 g_lo: &[G],
468 g_hi: &[G],
469 x_lo: G::ScalarField,
470 x_hi: G::ScalarField,
471) -> Vec<G> {
472 let mut g_proj: Vec<G::Group> = {
473 let pairs: Vec<_> = g_lo.iter().zip(g_hi).collect();
474 pairs
475 .into_par_iter()
476 .map(|(lo, hi)| window_shamir::<G>(x_lo, *lo, x_hi, *hi))
477 .collect()
478 };
479 G::Group::normalize_batch(g_proj.as_mut_slice())
480}
481
482pub fn affine_shamir_window_table<P: SWCurveConfig>(
483 denominators: &mut [P::BaseField],
484 g1: &[SWJAffine<P>],
485 g2: &[SWJAffine<P>],
486) -> [Vec<SWJAffine<P>>; 15] {
487 fn assign<A: Copy>(dst: &mut [A], src: &[A]) {
488 let n = dst.len();
489 dst[..n].clone_from_slice(&src[..n]);
490 }
491
492 let n = g1.len();
493
494 let mut res: [Vec<_>; 15] = [
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 vec![SWJAffine::<P>::zero(); n],
509 vec![SWJAffine::<P>::zero(); n],
510 ];
511
512 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] =
513 &mut res;
514
515 assign(g01_00, g1);
516
517 assign(g10_00, g1);
518 batch_add_assign(denominators, g10_00, g1);
519
520 assign(g11_00, g10_00);
521 batch_add_assign(denominators, g11_00, g1);
522
523 assign(g00_01, g2);
524
525 assign(g01_01, g00_01);
526 batch_add_assign(denominators, g01_01, g1);
527
528 assign(g10_01, g01_01);
529 batch_add_assign(denominators, g10_01, g1);
530
531 assign(g11_01, g10_01);
532 batch_add_assign(denominators, g11_01, g1);
533
534 assign(g00_10, g00_01);
535 batch_add_assign(denominators, g00_10, g2);
536
537 assign(g01_10, g00_10);
538 batch_add_assign(denominators, g01_10, g1);
539
540 assign(g10_10, g01_10);
541 batch_add_assign(denominators, g10_10, g1);
542
543 assign(g11_10, g10_10);
544 batch_add_assign(denominators, g11_10, g1);
545
546 assign(g00_11, g00_10);
547 batch_add_assign(denominators, g00_11, g2);
548
549 assign(g01_11, g00_11);
550 batch_add_assign(denominators, g01_11, g1);
551
552 assign(g10_11, g01_11);
553 batch_add_assign(denominators, g10_11, g1);
554
555 assign(g11_11, g10_11);
556 batch_add_assign(denominators, g11_11, g1);
557
558 res
559}
560
561pub fn affine_shamir_window_table_one<P: SWCurveConfig>(
562 denominators: &mut [P::BaseField],
563 g1: &[SWJAffine<P>],
564) -> [Vec<SWJAffine<P>>; 3] {
565 fn assign<A: Copy>(dst: &mut [A], src: &[A]) {
566 let n = dst.len();
567 dst[..n].clone_from_slice(&src[..n]);
568 }
569
570 let n = g1.len();
571
572 let mut res: [Vec<_>; 3] = [
573 vec![SWJAffine::<P>::zero(); n],
574 vec![SWJAffine::<P>::zero(); n],
575 vec![SWJAffine::<P>::zero(); n],
576 ];
577
578 let [g01, g10, g11] = &mut res;
579
580 assign(g01, g1);
581
582 assign(g10, g1);
583 batch_add_assign(denominators, g10, g1);
584
585 assign(g11, g10);
586 batch_add_assign(denominators, g11, g1);
587
588 res
589}
590
591fn window_shamir<G: AffineRepr>(x1: G::ScalarField, g1: G, x2: G::ScalarField, g2: G) -> G::Group {
592 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] =
593 shamir_window_table(g1, g2);
594
595 let windows1 = BitIteratorBE::new(x1.into_bigint()).tuples();
596 let windows2 = BitIteratorBE::new(x2.into_bigint()).tuples();
597
598 let mut res = G::Group::zero();
599
600 for ((hi_1, lo_1), (hi_2, lo_2)) in windows1.zip(windows2) {
601 res.double_in_place();
602 res.double_in_place();
603 match ((hi_1, lo_1), (hi_2, lo_2)) {
604 ((false, false), (false, false)) => (),
605 ((false, true), (false, false)) => res.add_assign(&g01_00),
606 ((true, false), (false, false)) => res.add_assign(&g10_00),
607 ((true, true), (false, false)) => res.add_assign(&g11_00),
608
609 ((false, false), (false, true)) => res.add_assign(&g00_01),
610 ((false, true), (false, true)) => res.add_assign(&g01_01),
611 ((true, false), (false, true)) => res.add_assign(&g10_01),
612 ((true, true), (false, true)) => res.add_assign(&g11_01),
613
614 ((false, false), (true, false)) => res.add_assign(&g00_10),
615 ((false, true), (true, false)) => res.add_assign(&g01_10),
616 ((true, false), (true, false)) => res.add_assign(&g10_10),
617 ((true, true), (true, false)) => res.add_assign(&g11_10),
618
619 ((false, false), (true, true)) => res.add_assign(&g00_11),
620 ((false, true), (true, true)) => res.add_assign(&g01_11),
621 ((true, false), (true, true)) => res.add_assign(&g10_11),
622 ((true, true), (true, true)) => res.add_assign(&g11_11),
623 }
624 }
625
626 res
627}
628
629pub fn shamir_window_table<G: AffineRepr>(g1: G, g2: G) -> [G; 16] {
630 let g00_00 = G::generator().into_group();
631 let g01_00 = g1.into_group();
632 let g10_00 = {
633 let mut g = g01_00;
634 g.add_assign(&g1);
635 g
636 };
637 let g11_00 = {
638 let mut g = g10_00;
639 g.add_assign(&g1);
640 g
641 };
642
643 let g00_01 = g2.into_group();
644 let g01_01 = {
645 let mut g = g00_01;
646 g.add_assign(&g1);
647 g
648 };
649 let g10_01 = {
650 let mut g = g01_01;
651 g.add_assign(&g1);
652 g
653 };
654 let g11_01 = {
655 let mut g = g10_01;
656 g.add_assign(&g1);
657 g
658 };
659
660 let g00_10 = {
661 let mut g = g00_01;
662 g.add_assign(&g2);
663 g
664 };
665 let g01_10 = {
666 let mut g = g00_10;
667 g.add_assign(&g1);
668 g
669 };
670 let g10_10 = {
671 let mut g = g01_10;
672 g.add_assign(&g1);
673 g
674 };
675 let g11_10 = {
676 let mut g = g10_10;
677 g.add_assign(&g1);
678 g
679 };
680 let g00_11 = {
681 let mut g = g00_10;
682 g.add_assign(&g2);
683 g
684 };
685 let g01_11 = {
686 let mut g = g00_11;
687 g.add_assign(&g1);
688 g
689 };
690 let g10_11 = {
691 let mut g = g01_11;
692 g.add_assign(&g1);
693 g
694 };
695 let g11_11 = {
696 let mut g = g10_11;
697 g.add_assign(&g1);
698 g
699 };
700
701 let mut v = vec![
702 g00_00, g01_00, g10_00, g11_00, g00_01, g01_01, g10_01, g11_01, g00_10, g01_10, g10_10,
703 g11_10, g00_11, g01_11, g10_11, g11_11,
704 ];
705 let v: Vec<_> = G::Group::normalize_batch(v.as_mut_slice());
706 [
707 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],
708 v[14], v[15],
709 ]
710}