Generalize AirExt::mask_points()

crates/prover/src/core/air/air_ext.rs

@@ -33,23 +33,19 @@ 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 mask = 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);
+            if mask.len() < points.len() {
+                mask.resize(points.len(), vec![]);
+            }
+            mask.as_mut().zip_eq(points).map(|(mask, points)| {
+                mask.extend(points);
+            });
         }
         // Add the composition polynomial mask points.
-        points.push(vec![vec![point]; SECURE_EXTENSION_DEGREE]);
-        points
+        mask.push(vec![vec![point]; SECURE_EXTENSION_DEGREE]);
+        mask
     }
 
     fn eval_composition_polynomial_at_point(