fix: thread_seeds as HashMap

privacy-scaling-explorations · Sep 4, 2023 · 8a25413 · 8a25413
1 parent f51e8c3
commit 8a25413
Showing 1 changed file with 18 additions and 7 deletions.
diff --git a/halo2_proofs/src/plonk/vanishing/prover.rs b/halo2_proofs/src/plonk/vanishing/prover.rs
@@ -1,4 +1,4 @@
-use std::iter;
+use std::{collections::HashMap, iter};
 
 use ff::Field;
 use group::Curve;
@@ -51,16 +51,27 @@ impl<C: CurveAffine> Argument<C> {
         let n = 1usize << domain.k() as usize;
         let mut rand_vec = vec![C::Scalar::ZERO; n];
 
-        let thread_seeds: Vec<ChaCha20Rng> = (0..current_num_threads())
-            .map(|_| {
+        let num_threads = current_num_threads();
+        let chunk_size = n / num_threads;
+        let thread_seeds = (0..)
+            .step_by(chunk_size + 1)
+            .take(n % num_threads)
+            .chain(
+                (chunk_size != 0)
+                    .then(|| ((n % num_threads) * (chunk_size + 1)..).step_by(chunk_size))
+                    .into_iter()
+                    .flatten(),
+            )
+            .take(num_threads)
+            .zip(iter::repeat_with(|| {
                 let mut seed = [0u8; 32];
                 rng.fill_bytes(&mut seed);
                 ChaCha20Rng::from_seed(seed)
-            })
-            .collect();
+            }))
+            .collect::<HashMap<_, _>>();
 
-        parallelize(&mut rand_vec, |chunk, idx| {
-            let mut rng = thread_seeds[idx].clone();
+        parallelize(&mut rand_vec, |chunk, offset| {
+            let mut rng = thread_seeds[&offset].clone();
             chunk
                 .iter_mut()
                 .for_each(|v| *v = C::Scalar::random(&mut rng));