From 3d8e9031088a1e7908e50f7d2c3113f126e2e4a4 Mon Sep 17 00:00:00 2001
From: chriseth <chriseth.github@gmail.com>
Date: Wed, 15 Jan 2025 17:00:52 +0000
Subject: [PATCH] Evaluaton of fixed columns that are constant outside the
 first and last row.

---
 executor-utils/src/lib.rs                     | 31 +++++++++++++-
 .../src/witgen/jit/single_step_processor.rs   | 42 +++++++++++++------
 executor/src/witgen/mod.rs                    |  6 ++-
 3 files changed, 64 insertions(+), 15 deletions(-)

diff --git a/executor-utils/src/lib.rs b/executor-utils/src/lib.rs
index d9d6b4e307..f81e81a844 100644
--- a/executor-utils/src/lib.rs
+++ b/executor-utils/src/lib.rs
@@ -1,3 +1,4 @@
+use itertools::Itertools;
 use powdr_ast::analyzed::Analyzed;
 use serde::{Deserialize, Serialize};
 use std::collections::{BTreeMap, BTreeSet};
@@ -41,6 +42,9 @@ impl<T: FieldElement> WitgenCallback<T> {
 #[derive(Serialize, Deserialize)]
 pub struct VariablySizedColumn<F> {
     column_by_size: BTreeMap<DegreeType, Vec<F>>,
+    /// If this is Some(x), then all sizes of this column have this value
+    /// in all rows except the first and the last.
+    constant_inner_value: Option<F>,
 }
 
 #[derive(Debug)]
@@ -66,23 +70,46 @@ impl<F> VariablySizedColumn<F> {
             .get(&size)
             .map(|column| column.as_slice())
     }
+
+    /// If this returns Some(x), then all sizes of this column have this value
+    /// in all rows except the first and the last.
+    pub fn has_constant_inner_value(&self) -> &Option<F> {
+        &self.constant_inner_value
+    }
 }
 
-impl<F> From<Vec<F>> for VariablySizedColumn<F> {
+impl<F: PartialEq + Copy> From<Vec<F>> for VariablySizedColumn<F> {
     fn from(column: Vec<F>) -> Self {
+        let constant_inner_value = constant_inner_value(&column);
         VariablySizedColumn {
             column_by_size: [(column.len() as DegreeType, column)].into_iter().collect(),
+            constant_inner_value,
         }
     }
 }
 
-impl<F> From<Vec<Vec<F>>> for VariablySizedColumn<F> {
+impl<F: PartialEq + Copy> From<Vec<Vec<F>>> for VariablySizedColumn<F> {
     fn from(columns: Vec<Vec<F>>) -> Self {
+        let constant_inner_values = columns
+            .iter()
+            .map(|v| constant_inner_value(v))
+            .collect::<Option<Vec<_>>>();
+        let constant_inner_value =
+            constant_inner_values.and_then(|v| v.iter().all_equal_value().ok().cloned());
         VariablySizedColumn {
             column_by_size: columns
                 .into_iter()
                 .map(|column| (column.len() as DegreeType, column))
                 .collect(),
+            constant_inner_value,
         }
     }
 }
+
+fn constant_inner_value<F: PartialEq + Copy>(column: &[F]) -> Option<F> {
+    column[1..column.len() - 1]
+        .iter()
+        .all_equal_value()
+        .ok()
+        .cloned()
+}
diff --git a/executor/src/witgen/jit/single_step_processor.rs b/executor/src/witgen/jit/single_step_processor.rs
index 8acbe0680c..ab0bff5912 100644
--- a/executor/src/witgen/jit/single_step_processor.rs
+++ b/executor/src/witgen/jit/single_step_processor.rs
@@ -50,11 +50,12 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> {
             (id, row_offset)
         });
         let block_size = 1;
-        let witgen = WitgenInference::new(self.fixed_data, NoEval, known_variables);
+
+        let witgen = WitgenInference::new(self.fixed_data, self, known_variables);
 
         Processor::new(
             self.fixed_data,
-            NoEval,
+            self,
             identities,
             block_size,
             false,
@@ -74,16 +75,11 @@ impl<'a, T: FieldElement> SingleStepProcessor<'a, T> {
     }
 }
 
-#[derive(Clone)]
-pub struct NoEval;
-
-impl<T: FieldElement> FixedEvaluator<T> for NoEval {
-    fn evaluate(&self, _var: &AlgebraicReference, _row_offset: i32) -> Option<T> {
-        // We can only return something here if the fixed column is constant
-        // in the region we are considering.
-        // This might be the case if we know we are not evaluating the first or the last
-        // row, but this is not yet implemented.
-        None
+/// Evaluator for fixed columns which are constant except for the first and last row.
+impl<T: FieldElement> FixedEvaluator<T> for &SingleStepProcessor<'_, T> {
+    fn evaluate(&self, var: &AlgebraicReference, _row_offset: i32) -> Option<T> {
+        assert!(var.is_fixed());
+        self.fixed_data.fixed_cols[&var.poly_id].has_constant_inner_value()
     }
 }
 
@@ -148,6 +144,25 @@ mod test {
         );
     }
 
+    #[test]
+    fn fib_with_boundary_conditions() {
+        let input = "
+namespace M(256);
+    col fixed FIRST = [1] + [0]*;
+    col fixed LAST = [0]* + [1];
+    let X;
+    let Y;
+    FIRST * (X - 1) = 0;
+    FIRST * (Y - 1) = 0;
+    (X' - Y) * (1 - LAST) = 0;
+    (Y' - (X + Y)) * (1 - LAST) = 0;";
+        let code = generate_single_step(input, "M").unwrap();
+        assert_eq!(
+            format_code(&code),
+            "M::X[1] = M::Y[0];\nM::Y[1] = (M::X[0] + M::Y[0]);"
+        );
+    }
+
     #[test]
     fn no_progress() {
         let input = "namespace M(256); let X; let Y; X' = X;";
@@ -241,4 +256,7 @@ machine_call(2, [Known(call_var(2, 1, 0)), Known(call_var(2, 1, 1)), Unknown(cal
 VM::instr_mul[1] = 1;"
         );
     }
+
+    #[test]
+    fn eval_fixed_on_inner_rows() {}
 }
diff --git a/executor/src/witgen/mod.rs b/executor/src/witgen/mod.rs
index 5faa25d122..6f3c1a1239 100644
--- a/executor/src/witgen/mod.rs
+++ b/executor/src/witgen/mod.rs
@@ -568,7 +568,7 @@ pub struct FixedColumn<'a, T> {
     values: &'a VariablySizedColumn<T>,
 }
 
-impl<'a, T> FixedColumn<'a, T> {
+impl<'a, T: FieldElement> FixedColumn<'a, T> {
     pub fn new(name: &'a str, values: &'a VariablySizedColumn<T>) -> FixedColumn<'a, T> {
         let name = name.to_string();
         FixedColumn { name, values }
@@ -589,6 +589,10 @@ impl<'a, T> FixedColumn<'a, T> {
         let max_size = self.values.available_sizes().into_iter().max().unwrap() as DegreeType;
         self.values(max_size)
     }
+
+    pub fn has_constant_inner_value(&self) -> Option<T> {
+        *self.values.has_constant_inner_value()
+    }
 }
 
 #[derive(Debug)]