From d273df790a5f4ea136cc29dbd10abd3b6439f71e Mon Sep 17 00:00:00 2001
From: Alon Titelman <alon.titelman@starkware.co>
Date: Mon, 10 Feb 2025 12:50:42 +0200
Subject: [PATCH] Formal `column_size` in expression evaluator.

---
 .../prover/src/constraint_framework/expr/evaluator.rs  | 10 +++-------
 crates/prover/src/examples/state_machine/mod.rs        |  2 +-
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/crates/prover/src/constraint_framework/expr/evaluator.rs b/crates/prover/src/constraint_framework/expr/evaluator.rs
index f44c08fe6..ef0709c91 100644
--- a/crates/prover/src/constraint_framework/expr/evaluator.rs
+++ b/crates/prover/src/constraint_framework/expr/evaluator.rs
@@ -4,8 +4,6 @@ use super::{BaseExpr, ExtExpr};
 use crate::constraint_framework::expr::ColumnExpr;
 use crate::constraint_framework::preprocessed_columns::PreProcessedColumnId;
 use crate::constraint_framework::{EvalAtRow, Relation, RelationEntry, INTERACTION_TRACE_IDX};
-use crate::core::fields::m31::M31;
-use crate::core::fields::FieldExpOps;
 use crate::core::lookups::utils::Fraction;
 
 pub struct FormalLogupAtRow {
@@ -21,6 +19,7 @@ pub struct FormalLogupAtRow {
 impl FormalLogupAtRow {
     pub fn new(interaction: usize, log_size: u32) -> Self {
         let claimed_sum_name = "claimed_sum".to_string();
+        let column_size_name = "column_size".to_string();
 
         Self {
             interaction,
@@ -30,10 +29,7 @@ impl FormalLogupAtRow {
             is_finalized: true,
             is_first: BaseExpr::zero(),
             cumsum_shift: ExtExpr::Param(claimed_sum_name)
-                * BaseExpr::Inv(Box::new(BaseExpr::pow(
-                    &BaseExpr::Const(M31(2)),
-                    log_size as u128,
-                ))),
+                * BaseExpr::Inv(Box::new(BaseExpr::Param(column_size_name))),
             log_size,
         }
     }
@@ -207,7 +203,7 @@ mod tests {
 \
         let constraint_1 = (QM31Impl::from_partial_evals([trace_2_column_3_offset_0, trace_2_column_4_offset_0, trace_2_column_5_offset_0, trace_2_column_6_offset_0]) \
             - (QM31Impl::from_partial_evals([trace_2_column_3_offset_neg_1, trace_2_column_4_offset_neg_1, trace_2_column_5_offset_neg_1, trace_2_column_6_offset_neg_1])) \
-                + (claimed_sum) * (qm31(32768, 0, 0, 0))) \
+                + (claimed_sum) * (1 / (column_size))) \
             * (intermediate1) \
             - (qm31(1, 0, 0, 0));"
             .to_string();
diff --git a/crates/prover/src/examples/state_machine/mod.rs b/crates/prover/src/examples/state_machine/mod.rs
index 4f6033cb5..3549f4103 100644
--- a/crates/prover/src/examples/state_machine/mod.rs
+++ b/crates/prover/src/examples/state_machine/mod.rs
@@ -344,7 +344,7 @@ mod tests {
 \
         let constraint_0 = (QM31Impl::from_partial_evals([trace_2_column_2_offset_0, trace_2_column_3_offset_0, trace_2_column_4_offset_0, trace_2_column_5_offset_0]) \
             - (QM31Impl::from_partial_evals([trace_2_column_2_offset_neg_1, trace_2_column_3_offset_neg_1, trace_2_column_4_offset_neg_1, trace_2_column_5_offset_neg_1])) \
-                + (claimed_sum) * (qm31(8388608, 0, 0, 0))\
+                + (claimed_sum) * (1 / (column_size))\
             ) \
             * ((intermediate0) * (intermediate1)) \
             - (intermediate1 - (intermediate0));"