From 55d0ff0072d41460fc3901fc80d263daaf071daa 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/builder/mod.rs           | 58 ++++++++++++++++++++++
 crates/prover/src/core/backend/simd/m31.rs |  6 +++
 crates/prover/src/lib.rs                   |  1 +
 3 files changed, 65 insertions(+)
 create mode 100644 crates/prover/src/builder/mod.rs

diff --git a/crates/prover/src/builder/mod.rs b/crates/prover/src/builder/mod.rs
new file mode 100644
index 000000000..19d747fe3
--- /dev/null
+++ b/crates/prover/src/builder/mod.rs
@@ -0,0 +1,58 @@
+/// ! 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::FieldExpOps;
+
+/// A trait for evaluating expressions at some point or row.
+pub trait EvalAtRow {
+    /// The base field type.
+    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>;
+
+    /// The extension field type.
+    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.
+    fn next_mask(&mut self) -> Self::F {
+        self.next_interaction_mask(0, [0])[0]
+    }
+
+    /// Returns the next mask for the given 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; 4]) -> 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 79887aae3..01affbf1a 100644
--- a/crates/prover/src/lib.rs
+++ b/crates/prover/src/lib.rs
@@ -12,6 +12,7 @@
     assert_matches,
     portable_simd
 )]
+pub mod builder;
 pub mod core;
 pub mod examples;
 pub mod hash_functions;