transpose_slice_to_basis_coefficients

Author: TomWambsgans

crates/lean_prover/src/prove_execution.rs

@@ -104,15 +104,8 @@ pub fn prove_execution(
     let dot_product_lengths: Vec<PF<EF>> = field_slice_as_base(&dot_product_columns[1]).unwrap();
 
     let dot_product_computations: &[EF] = &dot_product_columns[8];
-    let mut dot_product_computations_base = (0..DIMENSION).map(|_| Vec::new()).collect::<Vec<_>>();
-    for row in dot_product_computations {
-        for (j, coeff) in <EF as BasedVectorSpace<PF<EF>>>::as_basis_coefficients_slice(row)
-            .iter()
-            .enumerate()
-        {
-            dot_product_computations_base[j].push(*coeff);
-        }
-    }
+    let dot_product_computations_base =
+        transpose_slice_to_basis_coefficients::<F, EF>(dot_product_computations);
 
     let n_rows_table_dot_products = dot_product_columns[0].len() - dot_product_padding_len;
     let log_n_rows_dot_product_table = log2_strict_usize(dot_product_columns[0].len());

crates/utils/src/misc.rs

@@ -68,6 +68,22 @@ pub fn field_slice_as_base<F: Field, EF: ExtensionField<F>>(slice: &[EF]) -> Opt
     slice.par_iter().map(|x| x.as_base()).collect()
 }
 
+pub fn transpose_slice_to_basis_coefficients<F: Field, EF: ExtensionField<F>>(
+    slice: &[EF],
+) -> Vec<Vec<F>> {
+    let res = vec![F::zero_vec(slice.len()); EF::DIMENSION];
+    slice.par_iter().enumerate().for_each(|(i, row)| {
+        let coeffs = EF::as_basis_coefficients_slice(row);
+        unsafe {
+            for (j, &coeff) in coeffs.iter().enumerate() {
+                let raw_ptr = res[j].as_ptr() as *mut F;
+                *raw_ptr.add(i) = coeff;
+            }
+        }
+    });
+    res
+}
+
 pub fn dot_product_with_base<EF: ExtensionField<PF<EF>>>(slice: &[EF]) -> EF {
     assert_eq!(slice.len(), <EF as BasedVectorSpace<PF<EF>>>::DIMENSION);
     (0..EF::DIMENSION)