From 97d58c7954a98c8b40e224cc4f7970cef9d1122f Mon Sep 17 00:00:00 2001
From: Sean Bowe <ewillbefull@gmail.com>
Date: Tue, 16 May 2017 16:41:31 -0600
Subject: [PATCH] Add batchexp utility function to Engine.

---
 src/curves/bls381/mod.rs | 31 +++++++++++++++++++++++++++++++
 src/curves/mod.rs        |  2 ++
 src/curves/tests/mod.rs  | 40 +++++++++++++++++++++++++++++++++++++++-
 3 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/src/curves/bls381/mod.rs b/src/curves/bls381/mod.rs
index d100da7..91bfc8e 100644
--- a/src/curves/bls381/mod.rs
+++ b/src/curves/bls381/mod.rs
@@ -1025,6 +1025,37 @@ impl Engine for Bls381 {
         f
     }
 
+    fn batchexp<G: Curve<Self>, S: AsRef<[Self::Fr]>>(&self, g: &mut [G::Affine], scalars: S, coeff: Option<&Self::Fr>)
+    {
+        use crossbeam;
+        use num_cpus;
+
+        assert_eq!(g.len(), scalars.as_ref().len());
+
+        let chunk = (g.len() / num_cpus::get()) + 1;
+
+        crossbeam::scope(|scope| {
+            for (g, s) in g.chunks_mut(chunk).zip(scalars.as_ref().chunks(chunk)) {
+                scope.spawn(move || {
+                    let mut table = WindowTable::new();
+
+                    for (g, s) in g.iter_mut().zip(s.iter()) {
+                        let mut s = *s;
+                        match coeff {
+                            Some(coeff) => {
+                                s.mul_assign(self, coeff);
+                            },
+                            _ => {}
+                        };
+                        let mut newg = g.to_jacobian(self);
+                        opt_exp(self, &mut newg, s.into_repr(self), &mut table);
+                        *g = newg.to_affine(self);
+                    }
+                });
+            }
+        });
+    }
+
     fn batch_baseexp<G: Curve<Self>, S: AsRef<[Self::Fr]>>(&self, table: &WindowTable<Self, G, Vec<G>>, s: S) -> Vec<G::Affine>
     {
         use crossbeam;
diff --git a/src/curves/mod.rs b/src/curves/mod.rs
index 60d2a94..7da0ac8 100644
--- a/src/curves/mod.rs
+++ b/src/curves/mod.rs
@@ -44,6 +44,8 @@ pub trait Engine: Sized + Clone + Send + Sync
     /// Perform multi-exponentiation. g and s must have the same length.
     fn multiexp<G: Curve<Self>>(&self, g: &[G::Affine], s: &[Self::Fr]) -> Result<G, ()>;
     fn batch_baseexp<G: Curve<Self>, S: AsRef<[Self::Fr]>>(&self, table: &WindowTable<Self, G, Vec<G>>, scalars: S) -> Vec<G::Affine>;
+
+    fn batchexp<G: Curve<Self>, S: AsRef<[Self::Fr]>>(&self, g: &mut [G::Affine], scalars: S, coeff: Option<&Self::Fr>);
 }
 
 pub trait Group<E: Engine>: Copy + Send + Sync + Sized
diff --git a/src/curves/tests/mod.rs b/src/curves/tests/mod.rs
index 33ba891..2d8da83 100644
--- a/src/curves/tests/mod.rs
+++ b/src/curves/tests/mod.rs
@@ -1,9 +1,44 @@
 use super::{Engine, Curve, CurveAffine, Field, PrimeField};
-use rand;
+use rand::{self, Rng};
 
 mod fields;
 mod curves;
 
+fn test_batchexp<E: Engine, G: Curve<E>>(e: &E) {
+    let rng = &mut rand::thread_rng();
+
+    fn test_batchexp_case<E: Engine, G: Curve<E>, R: Rng>(e: &E, rng: &mut R, amount: usize, coeff: Option<&E::Fr>)
+    {
+        let mut g: Vec<G::Affine> = (0..amount).map(|_| G::random(e, rng).to_affine(e)).collect();
+        let mut s: Vec<E::Fr> = (0..amount).map(|_| E::Fr::random(e, rng)).collect();
+
+        let mut g_batch = g.clone();
+
+        e.batchexp::<G, _>(&mut g_batch, &s, coeff);
+
+        for (g, s) in g.iter_mut().zip(s.iter_mut()) {
+            match coeff {
+                Some(coeff) => {
+                    s.mul_assign(e, &coeff);
+                },
+                _ => {}
+            }
+            *g = g.mul(e, s).to_affine(e);
+        }
+
+        assert_eq!(g_batch, g);
+    }
+
+    for amt in 10..100 {
+        if amt % 2 == 0 {
+            let coeff = &E::Fr::random(e, rng);
+            test_batchexp_case::<E, G, _>(e, rng, amt, Some(coeff));
+        } else {
+            test_batchexp_case::<E, G, _>(e, rng, amt, None);
+        }
+    }
+}
+
 fn test_multiexp<E: Engine, G: Curve<E>>(e: &E) {
     fn naiveexp<E: Engine, G: Curve<E>>(e: &E, g: &[G::Affine], s: &[E::Fr]) -> G
     {
@@ -120,6 +155,9 @@ pub fn test_engine<E: Engine>() {
     test_frobenius(&engine);
     test_multiexp::<E, E::G1>(&engine);
     test_multiexp::<E, E::G2>(&engine);
+
+    test_batchexp::<E, E::G1>(&engine);
+    test_batchexp::<E, E::G2>(&engine);
 }
 
 fn test_frobenius<E: Engine>(e: &E) {