fix the parallel

arkworks-rs · May 29, 2024 · 716b728 · 716b728
1 parent 3a265f9
commit 716b728
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 13 deletions.
diff --git a/ec/src/models/short_weierstrass/affine.rs b/ec/src/models/short_weierstrass/affine.rs
@@ -1,5 +1,6 @@
 use ark_serialize::{
-    CanonicalDeserialize, CanonicalSerialize, Compress, SerializationError, Valid, Validate,
+    batch_check_helper, CanonicalDeserialize, CanonicalSerialize, Compress, SerializationError,
+    Valid, Validate,
 };
 use ark_std::{
     borrow::Borrow,
@@ -382,6 +383,13 @@ impl<P: SWCurveConfig> Valid for Affine<P> {
             Err(SerializationError::InvalidData)
         }
     }
+
+    fn batch_check<'a>(batch: impl Iterator<Item = &'a Self>) -> Result<(), SerializationError>
+    where
+        Self: 'a,
+    {
+        batch_check_helper(batch)
+    }
 }
 
 impl<P: SWCurveConfig> CanonicalDeserialize for Affine<P> {

diff --git a/ec/src/models/twisted_edwards/affine.rs b/ec/src/models/twisted_edwards/affine.rs
@@ -1,5 +1,6 @@
 use ark_serialize::{
-    CanonicalDeserialize, CanonicalSerialize, Compress, SerializationError, Valid, Validate,
+    batch_check_helper, CanonicalDeserialize, CanonicalSerialize, Compress, SerializationError,
+    Valid, Validate,
 };
 use ark_std::{
     borrow::Borrow,
@@ -339,6 +340,13 @@ impl<P: TECurveConfig> Valid for Affine<P> {
             Err(SerializationError::InvalidData)
         }
     }
+
+    fn batch_check<'a>(batch: impl Iterator<Item = &'a Self>) -> Result<(), SerializationError>
+    where
+        Self: 'a,
+    {
+        batch_check_helper(batch)
+    }
 }
 
 impl<P: TECurveConfig> CanonicalDeserialize for Affine<P> {

diff --git a/ec/src/pairing.rs b/ec/src/pairing.rs
@@ -1,6 +1,7 @@
 use ark_ff::{AdditiveGroup, CyclotomicMultSubgroup, Field, One, PrimeField};
 use ark_serialize::{
-    CanonicalDeserialize, CanonicalSerialize, Compress, SerializationError, Valid, Validate,
+    batch_check_helper, CanonicalDeserialize, CanonicalSerialize, Compress, SerializationError,
+    Valid, Validate,
 };
 use ark_std::{
     borrow::Borrow,
@@ -165,6 +166,13 @@ impl<P: Pairing> Valid for PairingOutput<P> {
             Err(SerializationError::InvalidData)
         }
     }
+
+    fn batch_check<'a>(batch: impl Iterator<Item = &'a Self>) -> Result<(), SerializationError>
+    where
+        Self: 'a,
+    {
+        batch_check_helper(batch)
+    }
 }
 
 impl<P: Pairing> CanonicalDeserialize for PairingOutput<P> {

diff --git a/serialize/src/lib.rs b/serialize/src/lib.rs
@@ -69,16 +69,8 @@ pub trait Valid: Sized {
     where
         Self: 'a,
     {
-        #[cfg(feature = "parallel")]
-        {
-            use rayon::{iter::ParallelBridge, prelude::ParallelIterator};
-            batch.par_bridge().try_for_each(|e| e.check())?;
-        }
-        #[cfg(not(feature = "parallel"))]
-        {
-            for item in batch {
-                item.check()?;
-            }
+        for item in batch {
+            item.check()?;
         }
         Ok(())
     }
@@ -241,3 +233,28 @@ pub fn buffer_bit_byte_size(modulus_bits: usize) -> (usize, usize) {
 pub const fn buffer_byte_size(modulus_bits: usize) -> usize {
     (modulus_bits + 7) / 8
 }
+
+pub fn batch_check_helper<'a, T>(
+    batch: impl Iterator<Item = &'a T>,
+) -> Result<(), SerializationError>
+where
+    T: 'a + Valid + Send + Clone,
+{
+    #[cfg(feature = "parallel")]
+    {
+        use ark_std::vec::Vec;
+        use rayon::prelude::{IntoParallelIterator, ParallelIterator};
+        batch
+            .map(|x| (*x).clone())
+            .collect::<Vec<T>>()
+            .into_par_iter()
+            .try_for_each(|e| e.check())?;
+    }
+    #[cfg(not(feature = "parallel"))]
+    {
+        for item in batch {
+            item.check()?;
+        }
+    }
+    Ok(())
+}