From f1cc284d42ca3e3fc3ee6b966292f11d4859954e Mon Sep 17 00:00:00 2001
From: Emanuele Cesena <ecesena@jumpcrypto.com>
Date: Wed, 10 May 2023 12:28:44 -0500
Subject: [PATCH] Optimize transpose

---
 plonky2/src/util/mod.rs | 36 ++++++------------------------------
 1 file changed, 6 insertions(+), 30 deletions(-)

diff --git a/plonky2/src/util/mod.rs b/plonky2/src/util/mod.rs
index dfc32fd8..d08f194d 100644
--- a/plonky2/src/util/mod.rs
+++ b/plonky2/src/util/mod.rs
@@ -1,6 +1,6 @@
-use alloc::vec;
 use alloc::vec::Vec;
 
+use plonky2_maybe_rayon::*;
 #[doc(inline)]
 pub use plonky2_util::*;
 
@@ -20,35 +20,11 @@ pub(crate) fn transpose_poly_values<F: Field>(polys: Vec<PolynomialValues<F>>) -
 }
 
 pub fn transpose<F: Field>(matrix: &[Vec<F>]) -> Vec<Vec<F>> {
-    let l = matrix.len();
-    let w = matrix[0].len();
-
-    let mut transposed = vec![vec![]; w];
-    for i in 0..w {
-        transposed[i].reserve_exact(l);
-        unsafe {
-            // After .reserve_exact(l), transposed[i] will have capacity at least l. Hence, set_len
-            // will not cause the buffer to overrun.
-            transposed[i].set_len(l);
-        }
-    }
-
-    // Optimization: ensure the larger loop is outside.
-    if w >= l {
-        for i in 0..w {
-            for j in 0..l {
-                transposed[i][j] = matrix[j][i];
-            }
-        }
-    } else {
-        for j in 0..l {
-            for i in 0..w {
-                transposed[i][j] = matrix[j][i];
-            }
-        }
-    }
-
-    transposed
+    let len = matrix[0].len();
+    (0..len)
+        .into_par_iter()
+        .map(|i| matrix.iter().map(|row| row[i]).collect())
+        .collect()
 }
 
 pub(crate) fn reverse_bits(n: usize, num_bits: usize) -> usize {