From 589957aec9ebdff12b5d0558a21ed7c0e36a40ef Mon Sep 17 00:00:00 2001
From: Vilius <vilius.bakutis23@gmail.com>
Date: Mon, 20 Jan 2025 21:43:10 +0200
Subject: [PATCH 1/3] optimize compute_cell_and_kzg_proofs

---
 kzg/src/das.rs | 76 +++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 60 insertions(+), 16 deletions(-)

diff --git a/kzg/src/das.rs b/kzg/src/das.rs
index e815f264d..0b4cbb415 100644
--- a/kzg/src/das.rs
+++ b/kzg/src/das.rs
@@ -238,24 +238,68 @@ pub trait DAS<B: EcBackend> {
         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::*;
 
-            reverse_bit_order(cells)?;
-        };
+            if let Some(cells) = cells {
+                let fft_result = fft_settings.fft_fr(&poly_monomial, false)?;
 
-        // 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)?;
+                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)| {
+                        cell.clone_from_slice(chunk);
+                    });
+
+                reverse_bit_order(cells.as_flattened_mut())?;
+            }
+
+            if let Some(proofs) = proofs {
+                let fk20_proofs = compute_fk20_proofs::<FIELD_ELEMENTS_PER_CELL, B>(
+                    cell_size,
+                    &poly_monomial,
+                    ts_size,
+                    fft_settings,
+                    settings,
+                )?;
+
+                proofs
+                    .par_iter_mut()
+                    .zip(fk20_proofs.into_par_iter()) // Consumes fk20_proofs
+                    .for_each(|(proof, result_proof)| {
+                        *proof = result_proof; // Move ownership directly
+                    });
+
+                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::<FIELD_ELEMENTS_PER_CELL, B>(
+                    cell_size,
+                    &poly_monomial,
+                    ts_size,
+                    fft_settings,
+                    settings,
+                )?;
+                proofs.clone_from_slice(&result);
+                reverse_bit_order(proofs)?;
+            }
         }
 
         Ok(())

From f0d175a9cec88a4d9fab57b08a924d0342ea66ed Mon Sep 17 00:00:00 2001
From: Vilius <vilius.bakutis23@gmail.com>
Date: Sun, 26 Jan 2025 17:08:02 +0200
Subject: [PATCH 2/3] try and fix

---
 kzg/src/das.rs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/kzg/src/das.rs b/kzg/src/das.rs
index 0b4cbb415..3b1890618 100644
--- a/kzg/src/das.rs
+++ b/kzg/src/das.rs
@@ -251,14 +251,14 @@ pub trait DAS<B: EcBackend> {
                     .par_chunks(flattened_cells.len() / num_cells)
                     .zip(cells.par_iter_mut())
                     .for_each(|(chunk, cell)| {
-                        cell.clone_from_slice(chunk);
+                        for (dest, src) in cell.iter_mut().zip(chunk.iter()) {
+                            *dest = *src;
+                        }
                     });
-
-                reverse_bit_order(cells.as_flattened_mut())?;
             }
 
             if let Some(proofs) = proofs {
-                let fk20_proofs = compute_fk20_proofs::<FIELD_ELEMENTS_PER_CELL, B>(
+                let fk20_proofs = compute_fk20_proofs::<B>(
                     cell_size,
                     &poly_monomial,
                     ts_size,
@@ -290,7 +290,7 @@ pub trait DAS<B: EcBackend> {
 
             // Compute proofs sequentially
             if let Some(proofs) = proofs {
-                let result = compute_fk20_proofs::<FIELD_ELEMENTS_PER_CELL, B>(
+                let result = compute_fk20_proofs::<B>(
                     cell_size,
                     &poly_monomial,
                     ts_size,

From 7391b8ede36f3ed2aa5207c8f3409ebb30b362fa Mon Sep 17 00:00:00 2001
From: Vilius <vilius.bakutis23@gmail.com>
Date: Sun, 26 Jan 2025 17:19:59 +0200
Subject: [PATCH 3/3] try fix 2

---
 kzg/src/das.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/kzg/src/das.rs b/kzg/src/das.rs
index 3b1890618..d2d423c0e 100644
--- a/kzg/src/das.rs
+++ b/kzg/src/das.rs
@@ -268,12 +268,12 @@ pub trait DAS<B: EcBackend> {
 
                 proofs
                     .par_iter_mut()
-                    .zip(fk20_proofs.into_par_iter()) // Consumes fk20_proofs
+                    .zip(fk20_proofs.into_par_iter())
                     .for_each(|(proof, result_proof)| {
-                        *proof = result_proof; // Move ownership directly
+                        *proof = result_proof;
                     });
 
-                reverse_bit_order(proofs)?;
+                reverse_bit_order(&mut proofs[..])?;
             }
         }