From bc5a93f785b21dc1851019c3b0dd1fac610a9434 Mon Sep 17 00:00:00 2001
From: Gali Michlevich <gali@starkware.co>
Date: Tue, 21 Jan 2025 11:45:34 +0200
Subject: [PATCH] Add Bitwise Xor to Preprocessed Column Enum

---
 .../prover/src/cairo_air/preprocessed.rs      | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 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
index 81ae7ee3..fea36e89 100644
--- a/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed.rs
+++ b/stwo_cairo_prover/crates/prover/src/cairo_air/preprocessed.rs
@@ -35,31 +35,45 @@ const fn preprocessed_log_sizes() -> [u32; N_PREPROCESSED_COLUMN_SIZES] {
 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) => {
+                let col: Col<SimdBackend, BaseField> = Col::<SimdBackend, BaseField>::from_iter(
+                    (0..((1 << column.log_size()) / N_LANES))
+                        .flat_map(|i| column.packed_at(i).to_array()),
+                );
+                CircleEvaluation::new(CanonicCoset::new(column.log_size()).circle_domain(), col)
+            }
         }
     }
 
     pub fn packed_at(&self, vec_row: usize) -> PackedM31 {
+        assert!(vec_row < (1 << self.log_size()) / N_LANES);
         match self {
             PreProcessedColumn::IsFirst(column) => column.packed_at(vec_row),
             PreProcessedColumn::Seq(column) => column.packed_at(vec_row),
+            PreProcessedColumn::BitwiseXor(column) => column.packed_at(vec_row),
         }
     }
 }
@@ -85,7 +99,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) }
     }
@@ -155,6 +168,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)]