From f1b9436deed3b6cc71ad1750f787ed7ef96a89b6 Mon Sep 17 00:00:00 2001
From: Shahar Papini <spapini@starkware.co>
Date: Tue, 9 Jul 2024 07:57:43 +0300
Subject: [PATCH] Eval framework

---
 crates/prover/src/constraint_framework/mod.rs | 64 +++++++++++++++++++
 crates/prover/src/core/backend/simd/m31.rs    |  6 ++
 crates/prover/src/lib.rs                      |  1 +
 3 files changed, 71 insertions(+)
 create mode 100644 crates/prover/src/constraint_framework/mod.rs

diff --git a/crates/prover/src/constraint_framework/mod.rs b/crates/prover/src/constraint_framework/mod.rs
new file mode 100644
index 000000000..8ca801aec
--- /dev/null
+++ b/crates/prover/src/constraint_framework/mod.rs
@@ -0,0 +1,64 @@
+/// ! This module contains helpers to express and use constraints for components.
+use std::fmt::Debug;
+use std::ops::{Add, AddAssign, Mul, Sub};
+
+use num_traits::{One, Zero};
+
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::qm31::SecureField;
+use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE;
+use crate::core::fields::FieldExpOps;
+
+/// A trait for evaluating expressions at some point or row.
+pub trait EvalAtRow {
+    // TODO(spapini): Use a better trait for these, like 'Algebra' or something.
+    /// The field type holding values of columns for the component. These are the inputs to the
+    /// constraints. It might be [BaseField] packed types, or even [SecureField], when evaluating
+    /// the columns out of domain.
+    type F: FieldExpOps
+        + Copy
+        + Debug
+        + AddAssign<Self::F>
+        + AddAssign<BaseField>
+        + Add<Self::F, Output = Self::F>
+        + Sub<Self::F, Output = Self::F>
+        + Mul<BaseField, Output = Self::F>
+        + Add<SecureField, Output = Self::EF>
+        + Mul<SecureField, Output = Self::EF>
+        + From<BaseField>;
+
+    /// A field type representing the closure of `F` with multiplying by [SecureField]. Constraints
+    /// usually get multiplied by [SecureField] values for security.
+    type EF: One
+        + Copy
+        + Debug
+        + Zero
+        + Add<SecureField, Output = Self::EF>
+        + Sub<SecureField, Output = Self::EF>
+        + Mul<SecureField, Output = Self::EF>
+        + Add<Self::F, Output = Self::EF>
+        + Mul<Self::F, Output = Self::EF>
+        + Sub<Self::EF, Output = Self::EF>
+        + Mul<Self::EF, Output = Self::EF>;
+
+    /// Returns the next mask value for the first interaction at offset 0.
+    fn next_trace_mask(&mut self) -> Self::F {
+        let [mask_item] = self.next_interaction_mask(0, [0]);
+        mask_item
+    }
+
+    /// Returns the mask values of the given offsets for the next column in the interaction.
+    fn next_interaction_mask<const N: usize>(
+        &mut self,
+        interaction: usize,
+        offsets: [isize; N],
+    ) -> [Self::F; N];
+
+    /// Adds a constraint to the component.
+    fn add_constraint<G>(&mut self, constraint: G)
+    where
+        Self::EF: Mul<G, Output = Self::EF>;
+
+    /// Combines 4 base field values into a single extension field value.
+    fn combine_ef(values: [Self::F; SECURE_EXTENSION_DEGREE]) -> Self::EF;
+}
diff --git a/crates/prover/src/core/backend/simd/m31.rs b/crates/prover/src/core/backend/simd/m31.rs
index 6ea2a7d65..efb664d23 100644
--- a/crates/prover/src/core/backend/simd/m31.rs
+++ b/crates/prover/src/core/backend/simd/m31.rs
@@ -253,6 +253,12 @@ impl From<[BaseField; N_LANES]> for PackedM31 {
     }
 }
 
+impl From<BaseField> for PackedM31 {
+    fn from(v: BaseField) -> Self {
+        Self::broadcast(v)
+    }
+}
+
 impl Distribution<PackedM31> for Standard {
     fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> PackedM31 {
         PackedM31::from_array(rng.gen())
diff --git a/crates/prover/src/lib.rs b/crates/prover/src/lib.rs
index 36ebc8b35..6c9ddaa02 100644
--- a/crates/prover/src/lib.rs
+++ b/crates/prover/src/lib.rs
@@ -12,6 +12,7 @@
     assert_matches,
     portable_simd
 )]
+pub mod constraint_framework;
 pub mod core;
 pub mod examples;
 pub mod math;