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

---
 .../prover/src/constraint_framework/info.rs   | 46 +++++++++++++++++++
 crates/prover/src/constraint_framework/mod.rs |  2 +
 2 files changed, 48 insertions(+)
 create mode 100644 crates/prover/src/constraint_framework/info.rs

diff --git a/crates/prover/src/constraint_framework/info.rs b/crates/prover/src/constraint_framework/info.rs
new file mode 100644
index 000000000..a5d62aac5
--- /dev/null
+++ b/crates/prover/src/constraint_framework/info.rs
@@ -0,0 +1,46 @@
+use std::ops::Mul;
+
+use num_traits::One;
+
+use super::EvalAtRow;
+use crate::core::fields::m31::BaseField;
+use crate::core::fields::qm31::SecureField;
+use crate::core::pcs::TreeVec;
+
+/// Collects information about the constraints.
+/// This includes mask offsets and columns at each interaction, and the number of constraints.
+#[derive(Default)]
+pub struct InfoEvaluator {
+    pub mask_offsets: TreeVec<Vec<Vec<isize>>>,
+    pub n_constraints: usize,
+}
+impl InfoEvaluator {
+    pub fn new() -> Self {
+        Self::default()
+    }
+}
+impl EvalAtRow for InfoEvaluator {
+    type F = BaseField;
+    type EF = SecureField;
+    fn next_interaction_mask<const N: usize>(
+        &mut self,
+        interaction: usize,
+        offsets: [isize; N],
+    ) -> [Self::F; N] {
+        if self.mask_offsets.len() <= interaction {
+            self.mask_offsets.resize(interaction + 1, vec![]);
+        }
+        self.mask_offsets[interaction].push(offsets.into_iter().collect());
+        [BaseField::one(); N]
+    }
+    fn add_constraint<G>(&mut self, _constraint: G)
+    where
+        Self::EF: Mul<G, Output = Self::EF>,
+    {
+        self.n_constraints += 1;
+    }
+
+    fn combine_ef(_values: [Self::F; 4]) -> Self::EF {
+        SecureField::one()
+    }
+}
diff --git a/crates/prover/src/constraint_framework/mod.rs b/crates/prover/src/constraint_framework/mod.rs
index f8b6197dd..c231da90b 100644
--- a/crates/prover/src/constraint_framework/mod.rs
+++ b/crates/prover/src/constraint_framework/mod.rs
@@ -1,5 +1,6 @@
 /// ! This module contains helpers to express and use constraints for components.
 mod assert;
+mod info;
 mod point;
 mod simd_domain;
 
@@ -7,6 +8,7 @@ use std::fmt::Debug;
 use std::ops::{Add, AddAssign, Mul, Sub};
 
 pub use assert::{assert_constraints, AssertEvaluator};
+pub use info::InfoEvaluator;
 use num_traits::{One, Zero};
 pub use point::PointEvaluator;
 pub use simd_domain::SimdDomainEvaluator;