Feat/batch addition

jolt-optimizations/Cargo.toml

@@ -52,5 +52,9 @@ harness = false
 name = "vector_scalar_mul_add_gamma_g2"
 harness = false
 
+[[bench]]
+name = "batch_addition"
+harness = false
+
 [[example]]
 name = "memory_test"

jolt-optimizations/benches/batch_addition.rs

@@ -0,0 +1,42 @@
+use ark_bn254::G1Affine;
+use ark_ec::{AffineRepr, CurveGroup};
+use ark_std::rand::RngCore;
+use ark_std::UniformRand;
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
+use jolt_optimizations::batch_g1_additions;
+use rayon::prelude::*;
+
+fn naive_parallel_sum(bases: &[G1Affine], indices: &[usize]) -> G1Affine {
+    indices.par_iter().map(|&idx| bases[idx]).reduce(
+        || G1Affine::zero(),
+        |acc, point| (acc + point).into_affine(),
+    )
+}
+
+fn bench_batch_addition(c: &mut Criterion) {
+    let mut group = c.benchmark_group("batch_g1_addition");
+    let mut rng = ark_std::test_rng();
+
+    // Test different sizes
+    for size in [1 << 20].iter() {
+        let bases: Vec<G1Affine> = (0..*size).map(|_| G1Affine::rand(&mut rng)).collect();
+
+        // Use half the points
+        let indices: Vec<usize> = (0..size / 2)
+            .map(|_| (rng.next_u64() as usize) % size)
+            .collect();
+
+        group.bench_with_input(BenchmarkId::new("batch_optimized", size), size, |b, _| {
+            b.iter(|| black_box(batch_g1_additions(&bases, &indices)));
+        });
+
+        group.bench_with_input(BenchmarkId::new("naive_parallel", size), size, |b, _| {
+            b.iter(|| black_box(naive_parallel_sum(&bases, &indices)));
+        });
+    }
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_batch_addition);
+criterion_main!(benches);

jolt-optimizations/src/batch_addition.rs

@@ -0,0 +1,148 @@
+//! Batch affine point addition for G1
+//!
+//! Implements efficient batch addition of affine elliptic curve points
+//! using Montgomery's batch inversion trick to minimize field inversions.
+
+use ark_bn254::G1Affine;
+use ark_ec::AffineRepr;
+use rayon::prelude::*;
+
+/// Performs batch addition of G1 affine points.
+///
+/// Given a slice of base points and indices, computes the sum of all points
+/// at the specified indices: bases[indices[0]] + bases[indices[1]] + ... + bases[indices[n-1]]
+///
+/// Uses batch inversion to compute all divisions efficiently:
+/// - Standard addition: 1 inversion per addition
+/// - Batch addition: 1 inversion + 3n multiplications for n additions
+///
+/// # Arguments
+/// * `bases` - Slice of G1 affine points to select from
+/// * `indices` - Slice of indices specifying which points to sum
+///
+/// # Returns
+/// The sum of all selected points as a single G1Affine point
+pub fn batch_g1_additions(bases: &[G1Affine], indices: &[usize]) -> G1Affine {
+    if indices.is_empty() {
+        return G1Affine::zero();
+    }
+
+    if indices.len() == 1 {
+        return bases[indices[0]];
+    }
+
+    let mut points: Vec<G1Affine> = Vec::with_capacity(indices.len());
+    points.extend(indices.iter().map(|&i| bases[i]));
+
+    while points.len() > 1 {
+        let current_len = points.len();
+        let pairs_count = current_len / 2;
+        let has_odd = current_len % 2 == 1;
+
+        // Collect denominators in parallel
+        let denominators: Vec<_> = (0..pairs_count)
+            .into_par_iter()
+            .map(|i| {
+                let p1 = points[i * 2];
+                let p2 = points[i * 2 + 1];
+                p2.x - p1.x
+            })
+            .collect();
+
+        // Batch invert all denominators
+        let mut inverses = denominators;
+        ark_ff::fields::batch_inversion(&mut inverses);
+
+        // Apply all additions in parallel
+        let mut new_points: Vec<G1Affine> = (0..pairs_count)
+            .into_par_iter()
+            .zip(inverses.par_iter())
+            .map(|(i, inv)| {
+                let p1 = points[i * 2];
+                let p2 = points[i * 2 + 1];
+                let lambda = (p2.y - p1.y) * inv;
+                let x3 = lambda * lambda - p1.x - p2.x;
+                let y3 = lambda * (p1.x - x3) - p1.y;
+                G1Affine::new(x3, y3)
+            })
+            .collect();
+
+        // Handle odd element
+        if has_odd {
+            new_points.push(points[current_len - 1]);
+        }
+
+        points = new_points;
+    }
+
+    points[0]
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ark_ec::CurveGroup;
+    use ark_std::rand::RngCore;
+    use ark_std::UniformRand;
+
+    #[test]
+    fn test_batch_addition_correctness() {
+        let mut rng = ark_std::test_rng();
+
+        let bases: Vec<G1Affine> = (0..10).map(|_| G1Affine::rand(&mut rng)).collect();
+
+        let indices = vec![2, 3, 4, 5, 6, 7];
+
+        let batch_result = batch_g1_additions(&bases, &indices);
+
+        let mut expected = G1Affine::zero();
+        for &idx in &indices {
+            expected = (expected + bases[idx]).into_affine();
+        }
+
+        assert_eq!(batch_result, expected, "Batch addition mismatch");
+    }
+
+    #[test]
+    fn test_empty_indices() {
+        let bases: Vec<G1Affine> = vec![G1Affine::generator(); 5];
+        let result = batch_g1_additions(&bases, &[]);
+        assert_eq!(result, G1Affine::zero());
+    }
+
+    #[test]
+    fn test_single_index() {
+        let mut rng = ark_std::test_rng();
+        let bases: Vec<G1Affine> = (0..5).map(|_| G1Affine::rand(&mut rng)).collect();
+
+        let result = batch_g1_additions(&bases, &[2]);
+        assert_eq!(result, bases[2]);
+    }
+
+    #[test]
+    fn test_stress_test_correctness() {
+        let mut rng = ark_std::test_rng();
+
+        let base_size = 10000;
+        let indices_size = 5000;
+
+        let bases: Vec<G1Affine> = (0..base_size).map(|_| G1Affine::rand(&mut rng)).collect();
+
+        let indices: Vec<usize> = (0..indices_size)
+            .map(|_| (rng.next_u64() as usize) % base_size)
+            .collect();
+
+        let batch_result = batch_g1_additions(&bases, &indices);
+
+        // Compute expected result using naive sequential addition
+        let mut expected = G1Affine::zero();
+        for &idx in &indices {
+            expected = (expected + bases[idx]).into_affine();
+        }
+
+        assert_eq!(
+            batch_result, expected,
+            "Stress test failed: batch result doesn't match expected sum"
+        );
+    }
+}

jolt-optimizations/src/lib.rs

@@ -8,6 +8,7 @@
 //! Also provides BN254 G1 equivalents.
 //! Uses Strauss-shamir batched scalar multiplication to maximally take advantage of GLV.
 
+pub mod batch_addition;
 pub mod constants;
 pub mod decomp_2d;
 pub mod decomp_4d;
@@ -51,3 +52,5 @@ pub use dory_g2::{
     vector_add_scalar_mul_g2_online, vector_add_scalar_mul_g2_precomputed,
     vector_add_scalar_mul_g2_windowed2_signed, vector_scalar_mul_add_gamma_g2_online,
 };
+
+pub use batch_addition::batch_g1_additions;