Modify mask_points for interaction.

starkware-libs · May 28, 2024 · ae9d116 · ae9d116
1 parent 482feb8
commit ae9d116
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 45 deletions.
diff --git a/crates/prover/src/core/air/air_ext.rs b/crates/prover/src/core/air/air_ext.rs
@@ -28,13 +28,22 @@ pub trait AirExt: Air {
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
-    ) -> ComponentVec<Vec<CirclePoint<SecureField>>> {
-        let mut component_points = ComponentVec(Vec::new());
+    ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
+        let mut trace_component_points = vec![];
+        let mut interaction_component_points = vec![];
         for component in self.components() {
             let points = component.mask_points(point);
-            component_points.push(points);
+            trace_component_points.extend(points[0].clone());
+            interaction_component_points.extend(points[1].clone());
         }
-        component_points
+        let mut points = TreeVec::new(vec![trace_component_points]);
+        if !interaction_component_points
+            .iter()
+            .all(|column| column.is_empty())
+        {
+            points.push(interaction_component_points);
+        }
+        points
     }
 
     fn interaction_elements(&self, channel: &mut Blake2sChannel) -> InteractionElements {

diff --git a/crates/prover/src/core/air/mod.rs b/crates/prover/src/core/air/mod.rs
@@ -40,7 +40,7 @@ pub trait Component {
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
-    ) -> ColumnVec<Vec<CirclePoint<SecureField>>>;
+    ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>>;
 
     /// Returns the ids of the interaction elements used by the component.
     fn interaction_element_ids(&self) -> Vec<String>;

diff --git a/crates/prover/src/core/mod.rs b/crates/prover/src/core/mod.rs
@@ -26,7 +26,7 @@ pub mod vcs;
 pub type ColumnVec<T> = Vec<T>;
 
 /// A vector of [ColumnVec]s. Each [ColumnVec] relates (by index) to a component in the air.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct ComponentVec<T>(pub Vec<ColumnVec<T>>);
 
 impl<T> ComponentVec<T> {

diff --git a/crates/prover/src/core/prover/mod.rs b/crates/prover/src/core/prover/mod.rs
@@ -114,19 +114,9 @@ pub fn generate_proof<B: Backend + MerkleOps<MerkleHasher>>(
     let oods_point = CirclePoint::<SecureField>::get_random_point(channel);
 
     // Get mask sample points relative to oods point.
-    let sample_points = air.mask_points(oods_point);
-
-    // TODO(spapini): Change when we support multiple interactions.
-    // First tree - trace.
-    let mut sample_points = TreeVec::new(vec![sample_points.flatten()]);
-    if commitment_scheme.trees.len() > 2 {
-        // Second tree - interaction trace.
-        sample_points.push(vec![
-            vec![oods_point];
-            commitment_scheme.trees[1].polynomials.len()
-        ]);
-    }
-    // Final tree - composition polynomial.
+    let mut sample_points = air.mask_points(oods_point);
+
+    // Get composition polynomial sample points.
     sample_points.push(vec![vec![oods_point]; 4]);
 
     // Prove the trace and composition OODS values, and retrieve them.
@@ -226,17 +216,9 @@ pub fn verify(
     let oods_point = CirclePoint::<SecureField>::get_random_point(channel);
 
     // Get mask sample points relative to oods point.
-    let trace_sample_points = air.mask_points(oods_point);
+    let mut sample_points = air.mask_points(oods_point);
 
-    // TODO(spapini): Change when we support multiple interactions.
-    // First tree - trace.
-    let mut sample_points = TreeVec::new(vec![trace_sample_points.flatten()]);
-    if proof.commitments.len() > 2 {
-        // Second tree - interaction trace.
-        // TODO(AlonH): Get the number of interaction traces from the air.
-        sample_points.push(vec![vec![oods_point]; 1]);
-    }
-    // Final tree - composition polynomial.
+    // Get composition polynomial sample points.
     sample_points.push(vec![vec![oods_point]; 4]);
 
     // TODO(spapini): Save clone.
@@ -277,7 +259,7 @@ fn sampled_values_to_mask(
     air.components().iter().for_each(|c| {
         trace_oods_values.push(
             flat_trace_values
-                .take(c.mask_points(CirclePoint::zero()).len())
+                .take(c.mask_points(CirclePoint::zero())[0].len())
                 .cloned()
                 .collect_vec(),
         )
@@ -292,9 +274,13 @@ fn sampled_values_to_mask(
         air.components()
             .iter()
             .zip_eq(&mut trace_oods_values)
-            .for_each(|(_component, values)| {
-                // TODO(AlonH): Implement n_interaction_columns() for component.
-                values.extend(interaction_values.take(1).cloned().collect_vec())
+            .for_each(|(component, values)| {
+                values.extend(
+                    interaction_values
+                        .take(component.mask_points(CirclePoint::zero())[1].len())
+                        .cloned()
+                        .collect_vec(),
+                )
             });
     }
 
@@ -410,8 +396,8 @@ mod tests {
         fn mask_points(
             &self,
             point: CirclePoint<SecureField>,
-        ) -> crate::core::ColumnVec<Vec<CirclePoint<SecureField>>> {
-            vec![vec![point]]
+        ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
+            TreeVec::new(vec![vec![vec![point]], vec![]])
         }
 
         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
@@ -88,12 +88,15 @@ impl Component for FibonacciComponent {
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
-    ) -> ColumnVec<Vec<CirclePoint<SecureField>>> {
-        shifted_mask_points(
-            &vec![vec![0, 1, 2]],
-            &[CanonicCoset::new(self.log_size)],
-            point,
-        )
+    ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
+        TreeVec::new(vec![
+            shifted_mask_points(
+                &vec![vec![0, 1, 2]],
+                &[CanonicCoset::new(self.log_size)],
+                point,
+            ),
+            vec![],
+        ])
     }
 
     fn interaction_element_ids(&self) -> Vec<String> {

diff --git a/crates/prover/src/examples/wide_fibonacci/component.rs b/crates/prover/src/examples/wide_fibonacci/component.rs
@@ -76,8 +76,11 @@ impl Component for WideFibComponent {
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
-    ) -> ColumnVec<Vec<CirclePoint<SecureField>>> {
-        fixed_mask_points(&vec![vec![0_usize]; self.n_columns()], point)
+    ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
+        TreeVec::new(vec![
+            fixed_mask_points(&vec![vec![0_usize]; self.n_columns()], point),
+            vec![vec![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
@@ -73,8 +73,11 @@ impl Component for SimdWideFibComponent {
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
-    ) -> ColumnVec<Vec<CirclePoint<SecureField>>> {
-        fixed_mask_points(&vec![vec![0_usize]; self.n_columns()], point)
+    ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
+        TreeVec::new(vec![
+            fixed_mask_points(&vec![vec![0_usize]; self.n_columns()], point),
+            vec![],
+        ])
     }
 
     fn interaction_element_ids(&self) -> Vec<String> {