Add Bitwise Xor to Preprocessed Column Enum

starkware-libs · Jan 30, 2025 · cb12a09 · cb12a09
1 parent b8dd76c
commit cb12a09
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed.rs b/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed.rs
@@ -32,27 +32,34 @@ const fn preprocessed_log_sizes() -> [u32; N_PREPROCESSED_COLUMN_SIZES] {
     arr
 }
 
+// TODO(Gali): Convert to dyn trait.
 pub enum PreProcessedColumn {
     IsFirst(IsFirst),
     Seq(Seq),
+    BitwiseXor(BitwiseXor),
 }
 impl PreProcessedColumn {
     pub fn log_size(&self) -> u32 {
         match self {
             PreProcessedColumn::IsFirst(column) => column.log_size,
             PreProcessedColumn::Seq(column) => column.log_size,
+            PreProcessedColumn::BitwiseXor(column) => column.log_size(),
         }
     }
+
     pub fn id(&self) -> PreProcessedColumnId {
         match self {
             PreProcessedColumn::IsFirst(column) => column.id(),
             PreProcessedColumn::Seq(column) => column.id(),
+            PreProcessedColumn::BitwiseXor(column) => column.id(),
         }
     }
+
     pub fn gen_column_simd(&self) -> CircleEvaluation<SimdBackend, BaseField, BitReversedOrder> {
         match self {
             PreProcessedColumn::IsFirst(column) => column.gen_column_simd(),
             PreProcessedColumn::Seq(column) => column.gen_column_simd(),
+            PreProcessedColumn::BitwiseXor(column) => column.gen_column_simd(),
         }
     }
 }
@@ -78,7 +85,6 @@ impl Seq {
     }
 
     pub fn packed_at(&self, vec_row: usize) -> PackedM31 {
-        assert!(vec_row < (1 << self.log_size) / N_LANES);
         PackedM31::broadcast(M31::from(vec_row * N_LANES))
             + unsafe { PackedM31::from_simd_unchecked(SIMD_ENUMERATION_0) }
     }
@@ -140,6 +146,13 @@ impl BitwiseXor {
         };
         unsafe { PackedM31::from_simd_unchecked(simd) }
     }
+
+    pub fn gen_column_simd(&self) -> CircleEvaluation<SimdBackend, BaseField, BitReversedOrder> {
+        let col: Col<SimdBackend, BaseField> = (0..((1 << self.log_size()) / N_LANES))
+            .flat_map(|i| self.packed_at(i).to_array())
+            .collect();
+        CircleEvaluation::new(CanonicCoset::new(self.log_size()).circle_domain(), col)
+    }
 }
 
 #[cfg(test)]