From fde8737967fd0e49299a5f13bf6fa95fea0ca4c4 Mon Sep 17 00:00:00 2001
From: Shahar Papini <spapini@starkware.co>
Date: Sun, 7 Jul 2024 13:47:50 +0300
Subject: [PATCH] is_first constant column

---
 crates/prover/src/builder/constant_cols.rs    | 12 ++++++++++++
 crates/prover/src/builder/mod.rs              |  1 +
 crates/prover/src/core/backend/cpu/mod.rs     |  3 +++
 crates/prover/src/core/backend/mod.rs         |  2 ++
 crates/prover/src/core/backend/simd/column.rs | 12 ++++++++++++
 crates/prover/src/examples/poseidon/mod.rs    | 17 +++++++++++++----
 6 files changed, 43 insertions(+), 4 deletions(-)
 create mode 100644 crates/prover/src/builder/constant_cols.rs

diff --git a/crates/prover/src/builder/constant_cols.rs b/crates/prover/src/builder/constant_cols.rs
new file mode 100644
index 000000000..29e2f9556
--- /dev/null
+++ b/crates/prover/src/builder/constant_cols.rs
@@ -0,0 +1,12 @@
+use num_traits::One;
+
+use crate::core::backend::{Backend, Col, Column};
+use crate::core::fields::m31::BaseField;
+use crate::core::poly::circle::{CanonicCoset, CircleEvaluation};
+use crate::core::poly::BitReversedOrder;
+
+pub fn gen_is_first<B: Backend>(log_size: u32) -> CircleEvaluation<B, BaseField, BitReversedOrder> {
+    let mut col = Col::<B, BaseField>::zeros(1 << log_size);
+    col.set(0, BaseField::one());
+    CircleEvaluation::new(CanonicCoset::new(log_size).circle_domain(), col)
+}
diff --git a/crates/prover/src/builder/mod.rs b/crates/prover/src/builder/mod.rs
index 54a5566c5..e72bb551b 100644
--- a/crates/prover/src/builder/mod.rs
+++ b/crates/prover/src/builder/mod.rs
@@ -1,4 +1,5 @@
 mod assert;
+pub mod constant_cols;
 mod domain;
 mod info;
 mod point;
diff --git a/crates/prover/src/core/backend/cpu/mod.rs b/crates/prover/src/core/backend/cpu/mod.rs
index c033e4cc2..66598cbf8 100644
--- a/crates/prover/src/core/backend/cpu/mod.rs
+++ b/crates/prover/src/core/backend/cpu/mod.rs
@@ -47,6 +47,9 @@ impl<T: Debug + Clone + Default> Column<T> for Vec<T> {
     fn at(&self, index: usize) -> T {
         self[index].clone()
     }
+    fn set(&mut self, index: usize, value: T) {
+        self[index] = value;
+    }
 }
 
 pub type CpuCirclePoly = CirclePoly<CpuBackend>;
diff --git a/crates/prover/src/core/backend/mod.rs b/crates/prover/src/core/backend/mod.rs
index f2d4fd043..c24efeb20 100644
--- a/crates/prover/src/core/backend/mod.rs
+++ b/crates/prover/src/core/backend/mod.rs
@@ -47,4 +47,6 @@ pub trait Column<T>: Clone + Debug + FromIterator<T> {
     }
     /// Retrieves the element at the given index.
     fn at(&self, index: usize) -> T;
+    /// Sets the element at the given index.
+    fn set(&mut self, index: usize, value: T);
 }
diff --git a/crates/prover/src/core/backend/simd/column.rs b/crates/prover/src/core/backend/simd/column.rs
index 5c2fdd380..505bd556c 100644
--- a/crates/prover/src/core/backend/simd/column.rs
+++ b/crates/prover/src/core/backend/simd/column.rs
@@ -72,6 +72,12 @@ impl Column<BaseField> for BaseFieldVec {
     fn at(&self, index: usize) -> BaseField {
         self.data[index / N_LANES].to_array()[index % N_LANES]
     }
+
+    fn set(&mut self, index: usize, value: BaseField) {
+        let mut packed = self.data[index / N_LANES].to_array();
+        packed[index % N_LANES] = value;
+        self.data[index / N_LANES] = PackedBaseField::from_array(packed)
+    }
 }
 
 impl FromIterator<BaseField> for BaseFieldVec {
@@ -124,6 +130,12 @@ impl Column<SecureField> for SecureFieldVec {
     fn at(&self, index: usize) -> SecureField {
         self.data[index / N_LANES].to_array()[index % N_LANES]
     }
+
+    fn set(&mut self, index: usize, value: SecureField) {
+        let mut packed = self.data[index / N_LANES].to_array();
+        packed[index % N_LANES] = value;
+        self.data[index / N_LANES] = PackedSecureField::from_array(packed)
+    }
 }
 
 impl FromIterator<SecureField> for SecureFieldVec {
diff --git a/crates/prover/src/examples/poseidon/mod.rs b/crates/prover/src/examples/poseidon/mod.rs
index d369eb2cb..3e62cd9d5 100644
--- a/crates/prover/src/examples/poseidon/mod.rs
+++ b/crates/prover/src/examples/poseidon/mod.rs
@@ -111,10 +111,10 @@ impl Component for PoseidonComponent {
         &self,
         point: CirclePoint<SecureField>,
     ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
-        TreeVec::new(vec![fixed_mask_points(
-            &vec![vec![0_usize]; N_COLUMNS],
-            point,
-        )])
+        TreeVec::new(vec![
+            fixed_mask_points(&vec![vec![0_usize]; N_COLUMNS], point),
+            vec![vec![point]],
+        ])
     }
 
     fn interaction_element_ids(&self) -> Vec<String> {
@@ -442,6 +442,7 @@ mod tests {
     use tracing::{span, Level};
 
     use super::N_LOG_INSTANCES_PER_ROW;
+    use crate::builder::constant_cols::gen_is_first;
     use crate::core::air::AirExt;
     use crate::core::backend::simd::SimdBackend;
     use crate::core::channel::{Blake2sChannel, Channel};
@@ -531,6 +532,11 @@ mod tests {
         commitment_scheme.commit_on_evals(trace, channel, &twiddles);
         span.exit();
 
+        // Constant trace.
+        let span = span!(Level::INFO, "Constant").entered();
+        commitment_scheme.commit_on_evals(vec![gen_is_first(log_n_rows)], channel, &twiddles);
+        span.exit();
+
         // Prove constraints.
         let component = PoseidonComponent { log_n_rows };
         let air = PoseidonAir { component };
@@ -549,7 +555,10 @@ mod tests {
 
         // Decommit.
         let sizes = air.column_log_sizes();
+        // Trace columns.
         commitment_scheme.commit(proof.commitments[0], &sizes[0], channel);
+        // Constant columns.
+        commitment_scheme.commit(proof.commitments[1], &[log_n_rows], channel);
 
         verify_without_commit(
             &air,