From 9f74d8693e28f89c1727634eec6afa76bcf97955 Mon Sep 17 00:00:00 2001
From: Shahar Samocha <shahars@starkware.co>
Date: Thu, 4 Apr 2024 14:26:54 +0300
Subject: [PATCH] Component mask_points api

---
 src/core/air/air_ext.rs                        | 16 ++++++----------
 src/core/air/mod.rs                            | 14 ++++++--------
 src/core/prover/mod.rs                         |  7 +++++++
 src/examples/fibonacci/component.rs            |  8 ++++++++
 src/examples/wide_fibonacci/constraint_eval.rs | 10 ++++++++++
 5 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/src/core/air/air_ext.rs b/src/core/air/air_ext.rs
index 9047107ff..03347d255 100644
--- a/src/core/air/air_ext.rs
+++ b/src/core/air/air_ext.rs
@@ -66,16 +66,12 @@ pub trait AirExt<B: Backend>: Air<B> {
         &self,
         point: CirclePoint<SecureField>,
     ) -> ComponentVec<Vec<CirclePoint<SecureField>>> {
-        let mut points = ComponentVec(Vec::new());
-        self.components().iter().for_each(|component| {
-            let domains = component
-                .trace_log_degree_bounds()
-                .iter()
-                .map(|&log_size| CanonicCoset::new(log_size))
-                .collect_vec();
-            points.push(component.mask().to_points(&domains, point));
-        });
-        points
+        let mut component_points = ComponentVec(Vec::new());
+        for component in self.components() {
+            let points = component.mask_points(point);
+            component_points.push(points);
+        }
+        component_points
     }
 
     fn eval_composition_polynomial_at_point(
diff --git a/src/core/air/mod.rs b/src/core/air/mod.rs
index a052c3c07..6fda0e7ed 100644
--- a/src/core/air/mod.rs
+++ b/src/core/air/mod.rs
@@ -1,8 +1,6 @@
 use std::iter::zip;
 use std::ops::Deref;
 
-use itertools::Itertools;
-
 use self::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator};
 use super::backend::Backend;
 use super::circle::CirclePoint;
@@ -76,6 +74,11 @@ pub trait Component<B: Backend> {
 
     fn mask(&self) -> Mask;
 
+    fn mask_points(
+        &self,
+        point: CirclePoint<SecureField>,
+    ) -> ColumnVec<Vec<CirclePoint<SecureField>>>;
+
     /// Calculates the mask points and evaluates them at each column.
     /// The mask values are used to evaluate the composition polynomial at a certain point.
     /// Returns two vectors with an entry for each column. Each entry holds the points/values
@@ -88,12 +91,7 @@ pub trait Component<B: Backend> {
         ColumnVec<Vec<CirclePoint<SecureField>>>,
         ColumnVec<Vec<SecureField>>,
     ) {
-        let domains = trace
-            .columns
-            .iter()
-            .map(|col| CanonicCoset::new(col.log_size()))
-            .collect_vec();
-        let points = self.mask().to_points(&domains, point);
+        let points = self.mask_points(point);
         let values = zip(&points, &trace.columns)
             .map(|(col_points, col)| {
                 col_points
diff --git a/src/core/prover/mod.rs b/src/core/prover/mod.rs
index 9222a4fe7..72cb8d067 100644
--- a/src/core/prover/mod.rs
+++ b/src/core/prover/mod.rs
@@ -284,6 +284,13 @@ mod tests {
             Mask(vec![vec![0]])
         }
 
+        fn mask_points(
+            &self,
+            point: CirclePoint<SecureField>,
+        ) -> crate::core::ColumnVec<Vec<CirclePoint<SecureField>>> {
+            vec![vec![point]]
+        }
+
         fn evaluate_constraint_quotients_at_point(
             &self,
             _point: CirclePoint<SecureField>,
diff --git a/src/examples/fibonacci/component.rs b/src/examples/fibonacci/component.rs
index a8e88db04..a8e08ba2f 100644
--- a/src/examples/fibonacci/component.rs
+++ b/src/examples/fibonacci/component.rs
@@ -120,6 +120,14 @@ impl Component<CPUBackend> for FibonacciComponent {
         Mask(vec![vec![0, 1, 2]])
     }
 
+    fn mask_points(
+        &self,
+        point: CirclePoint<SecureField>,
+    ) -> ColumnVec<Vec<CirclePoint<SecureField>>> {
+        self.mask()
+            .to_points(&[CanonicCoset::new(self.log_size)], point)
+    }
+
     fn evaluate_constraint_quotients_at_point(
         &self,
         point: CirclePoint<SecureField>,
diff --git a/src/examples/wide_fibonacci/constraint_eval.rs b/src/examples/wide_fibonacci/constraint_eval.rs
index 261b6437d..8405f02c6 100644
--- a/src/examples/wide_fibonacci/constraint_eval.rs
+++ b/src/examples/wide_fibonacci/constraint_eval.rs
@@ -30,6 +30,16 @@ impl Component<CPUBackend> for WideFibComponent {
         Mask(vec![vec![0_usize]; 256])
     }
 
+    fn mask_points(
+        &self,
+        point: CirclePoint<SecureField>,
+    ) -> ColumnVec<Vec<CirclePoint<SecureField>>> {
+        self.mask()
+            .iter()
+            .map(|col| col.iter().map(|_| point).collect())
+            .collect()
+    }
+
     // TODO(ShaharS), precompute random coeff powers.
     // TODO(ShaharS), use intermidiate value to save the computation of the squares.
     fn evaluate_constraint_quotients_on_domain(