parralelize ret (#234)

stwo_cairo_prover/crates/prover/Cargo.toml

@@ -3,6 +3,9 @@ name = "stwo_cairo_prover"
 version = "0.1.0"
 edition = "2021"
 
+[features]
+parallel = ["rayon"]
+
 [dependencies]
 bytemuck.workspace = true
 cairo-lang-casm.workspace = true
@@ -17,6 +20,7 @@ thiserror.workspace = true
 tracing.workspace = true
 paste.workspace = true
 prover_types = { path = "../prover_types" }
+rayon = { version = "1.10.0", optional = true }
 
 [dev-dependencies]
 cairo-lang-casm.workspace = true

stwo_cairo_prover/crates/prover/src/components/mod.rs

@@ -1,5 +1,10 @@
+use std::array;
+use std::cell::UnsafeCell;
+
 use prover_types::simd::N_LANES;
+use stwo_prover::core::backend::simd::column::BaseColumn;
 use stwo_prover::core::backend::simd::conversion::Pack;
+use stwo_prover::core::backend::Column;
 
 pub mod add_ap_opcode_is_imm_f_op1_base_fp_f;
 pub mod add_ap_opcode_is_imm_f_op1_base_fp_t;
@@ -47,3 +52,23 @@ pub fn pack_values<T: Pack>(values: &[T]) -> Vec<T::SimdType> {
         .map(|c| T::pack(*c))
         .collect()
 }
+
+// TODO(Ohad): move to another file.
+struct ThreadSafeTrace<const N: usize> {
+    data: [UnsafeCell<BaseColumn>; N],
+}
+
+unsafe impl<const N: usize> Send for ThreadSafeTrace<N> {}
+unsafe impl<const N: usize> Sync for ThreadSafeTrace<N> {}
+
+impl<const N: usize> ThreadSafeTrace<N> {
+    fn new(rows: usize) -> Self {
+        Self {
+            data: array::from_fn(|_| UnsafeCell::new(unsafe { BaseColumn::uninitialized(rows) })),
+        }
+    }
+
+    pub fn inner(self) -> [BaseColumn; N] {
+        self.data.map(|c| c.into_inner())
+    }
+}