optimize compute_cell_and_kzg_proofs

kzg/src/das.rs

@@ -226,7 +226,7 @@
             return Err("Both cells & proofs cannot be none".to_string());
         }
 
         let settings = self.kzg_settings();
         let ts_size = settings.get_g1_monomial().len();
         let cell_size = settings.get_cell_size();
 
@@ -238,24 +238,68 @@
         let fft_settings = self.kzg_settings().get_fft_settings();
         poly_lagrange_to_monomial::<B>(&mut poly_monomial[..ts_size], fft_settings)?;
 
-        // compute cells
-        if let Some(cells) = cells {
-            cells.clone_from_slice(&fft_settings.fft_fr(&poly_monomial, false)?);
+        #[cfg(feature = "parallel")]
+        {
+            use rayon::prelude::*;
+
+            if let Some(cells) = cells {
+                let fft_result = fft_settings.fft_fr(&poly_monomial, false)?;
+
+                let flattened_cells = fft_result.as_slice().to_vec();
+                let num_cells = cells.len();
+                flattened_cells
+                    .par_chunks(flattened_cells.len() / num_cells)
+                    .zip(cells.par_iter_mut())
+                    .for_each(|(chunk, cell)| {
+                        for (dest, src) in cell.iter_mut().zip(chunk.iter()) {
+                            *dest = *src;
+                        }
+                    });
+            }
 
-            reverse_bit_order(cells)?;
-        };
+            if let Some(proofs) = proofs {
+                let fk20_proofs = compute_fk20_proofs::<B>(
+                    cell_size,
+                    &poly_monomial,
+                    ts_size,
+                    fft_settings,
+                    settings,
+                )?;
+
+                proofs
+                    .par_iter_mut()
+                    .zip(fk20_proofs.into_par_iter())
+                    .for_each(|(proof, result_proof)| {
+                        *proof = result_proof;
+                    });
+
+                reverse_bit_order(&mut proofs[..])?;
+            }
+        }
 
-        // compute proofs
-        if let Some(proofs) = proofs {
-            let result = compute_fk20_proofs::<B>(
-                cell_size,
-                &poly_monomial,
-                ts_size,
-                fft_settings,
-                settings,
-            )?;
-            proofs.clone_from_slice(&result);
-            reverse_bit_order(proofs)?;
+        #[cfg(not(feature = "parallel"))]
+        {
+            // Compute cells sequentially
+            if let Some(cells) = cells {
+                cells
+                    .as_flattened_mut()
+                    .clone_from_slice(&fft_settings.fft_fr(&poly_monomial, false)?);
+
+                reverse_bit_order(cells.as_flattened_mut())?;
+            }
+
+            // Compute proofs sequentially
+            if let Some(proofs) = proofs {
+                let result = compute_fk20_proofs::<B>(
+                    cell_size,
+                    &poly_monomial,
+                    ts_size,
+                    fft_settings,
+                    settings,
+                )?;
+                proofs.clone_from_slice(&result);
+                reverse_bit_order(proofs)?;
+            }
         }
 
         Ok(())