Generalize AirExt::mask_points()

starkware-libs · Jul 8, 2024 · ba3839f · ba3839f
1 parent 5ab9ac2
commit ba3839f
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 45 deletions.
diff --git a/crates/prover/src/core/air/air_ext.rs b/crates/prover/src/core/air/air_ext.rs
@@ -33,23 +33,21 @@ pub trait AirExt: Air {
         &self,
         point: CirclePoint<SecureField>,
     ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
-        let mut trace_component_points = vec![];
-        let mut interaction_component_points = vec![];
+        let mut air_points = TreeVec::default();
         for component in self.components() {
-            let points = component.mask_points(point);
-            trace_component_points.extend(points[0].clone());
-            interaction_component_points.extend(points[1].clone());
-        }
-        let mut points = TreeVec::new(vec![trace_component_points]);
-        if !interaction_component_points
-            .iter()
-            .all(|column| column.is_empty())
-        {
-            points.push(interaction_component_points);
+            let component_points = component.mask_points(point);
+            if air_points.len() < component_points.len() {
+                air_points.resize(component_points.len(), vec![]);
+            }
+            air_points.as_mut().zip_eq(component_points).map(
+                |(air_tree_points, component_tree_points)| {
+                    air_tree_points.extend(component_tree_points);
+                },
+            );
         }
         // Add the composition polynomial mask points.
-        points.push(vec![vec![point]; SECURE_EXTENSION_DEGREE]);
-        points
+        air_points.push(vec![vec![point]; SECURE_EXTENSION_DEGREE]);
+        air_points
     }
 
     fn eval_composition_polynomial_at_point(
@@ -72,18 +70,19 @@ pub trait AirExt: Air {
     }
 
     fn column_log_sizes(&self) -> TreeVec<ColumnVec<u32>> {
-        let mut trace_tree = vec![];
-        let mut interaction_tree = vec![];
+        let mut air_sizes = TreeVec::default();
         self.components().iter().for_each(|component| {
-            let bounds = component.trace_log_degree_bounds();
-            trace_tree.extend(bounds[0].clone());
-            interaction_tree.extend(bounds[1].clone());
+            let component_sizes = component.trace_log_degree_bounds();
+            if air_sizes.len() < component_sizes.len() {
+                air_sizes.resize(component_sizes.len(), vec![]);
+            }
+            air_sizes.as_mut().zip_eq(component_sizes).map(
+                |(air_tree_sizes, component_tree_sizes)| {
+                    air_tree_sizes.extend(component_tree_sizes)
+                },
+            );
         });
-        let mut sizes = TreeVec::new(vec![trace_tree]);
-        if !interaction_tree.is_empty() {
-            sizes.push(interaction_tree);
-        }
-        sizes
+        air_sizes
     }
 
     fn component_traces<'a, B: Backend + MerkleOps<Blake2sMerkleHasher>>(

diff --git a/crates/prover/src/core/prover/mod.rs b/crates/prover/src/core/prover/mod.rs
@@ -390,14 +390,14 @@ mod tests {
         }
 
         fn trace_log_degree_bounds(&self) -> TreeVec<ColumnVec<u32>> {
-            TreeVec::new(vec![vec![self.log_size], vec![]])
+            TreeVec::new(vec![vec![self.log_size]])
         }
 
         fn mask_points(
             &self,
             point: CirclePoint<SecureField>,
         ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
-            TreeVec::new(vec![vec![vec![point]], vec![]])
+            TreeVec::new(vec![vec![vec![point]]])
         }
 
         fn interaction_element_ids(&self) -> Vec<String> {

diff --git a/crates/prover/src/examples/fibonacci/component.rs b/crates/prover/src/examples/fibonacci/component.rs
@@ -86,21 +86,18 @@ impl Component for FibonacciComponent {
     }
 
     fn trace_log_degree_bounds(&self) -> TreeVec<ColumnVec<u32>> {
-        TreeVec::new(vec![vec![self.log_size], vec![]])
+        TreeVec::new(vec![vec![self.log_size]])
     }
 
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
     ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
-        TreeVec::new(vec![
-            shifted_mask_points(
-                &vec![vec![0, 1, 2]],
-                &[CanonicCoset::new(self.log_size)],
-                point,
-            ),
-            vec![],
-        ])
+        TreeVec::new(vec![shifted_mask_points(
+            &vec![vec![0, 1, 2]],
+            &[CanonicCoset::new(self.log_size)],
+            point,
+        )])
     }
 
     fn interaction_element_ids(&self) -> Vec<String> {

diff --git a/crates/prover/src/examples/poseidon/mod.rs b/crates/prover/src/examples/poseidon/mod.rs
@@ -101,17 +101,17 @@ impl Component for PoseidonComponent {
     }
 
     fn trace_log_degree_bounds(&self) -> TreeVec<ColumnVec<u32>> {
-        TreeVec::new(vec![vec![self.log_column_size(); N_COLUMNS], vec![]])
+        TreeVec::new(vec![vec![self.log_column_size(); N_COLUMNS]])
     }
 
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
     ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
-        TreeVec::new(vec![
-            fixed_mask_points(&vec![vec![0_usize]; N_COLUMNS], point),
-            vec![],
-        ])
+        TreeVec::new(vec![fixed_mask_points(
+            &vec![vec![0_usize]; N_COLUMNS],
+            point,
+        )])
     }
 
     fn interaction_element_ids(&self) -> Vec<String> {

diff --git a/crates/prover/src/examples/wide_fibonacci/simd.rs b/crates/prover/src/examples/wide_fibonacci/simd.rs
@@ -93,17 +93,17 @@ impl Component for SimdWideFibComponent {
     }
 
     fn trace_log_degree_bounds(&self) -> TreeVec<ColumnVec<u32>> {
-        TreeVec::new(vec![vec![self.log_column_size(); self.n_columns()], vec![]])
+        TreeVec::new(vec![vec![self.log_column_size(); self.n_columns()]])
     }
 
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
     ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
-        TreeVec::new(vec![
-            fixed_mask_points(&vec![vec![0_usize]; self.n_columns()], point),
-            vec![],
-        ])
+        TreeVec::new(vec![fixed_mask_points(
+            &vec![vec![0_usize]; self.n_columns()],
+            point,
+        )])
     }
 
     fn interaction_element_ids(&self) -> Vec<String> {