Cache InfoEvaluator result. (#866)

starkware-libs · Nov 5, 2024 · feb141a · feb141a
1 parent 8957dea
commit feb141a
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 7 deletions.
diff --git a/crates/prover/src/constraint_framework/component.rs b/crates/prover/src/constraint_framework/component.rs
@@ -79,15 +79,17 @@ pub trait FrameworkEval {
 pub struct FrameworkComponent<C: FrameworkEval> {
     eval: C,
     trace_locations: TreeVec<TreeSubspan>,
+    info: InfoEvaluator,
 }
 
 impl<E: FrameworkEval> FrameworkComponent<E> {
     pub fn new(location_allocator: &mut TraceLocationAllocator, eval: E) -> Self {
-        let eval_tree_structure = eval.evaluate(InfoEvaluator::default()).mask_offsets;
-        let trace_locations = location_allocator.next_for_structure(&eval_tree_structure);
+        let info = eval.evaluate(InfoEvaluator::default());
+        let trace_locations = location_allocator.next_for_structure(&info.mask_offsets);
         Self {
             eval,
             trace_locations,
+            info,
         }
     }
 
@@ -98,25 +100,26 @@ impl<E: FrameworkEval> FrameworkComponent<E> {
 
 impl<E: FrameworkEval> Component for FrameworkComponent<E> {
     fn n_constraints(&self) -> usize {
-        self.eval.evaluate(InfoEvaluator::default()).n_constraints
+        self.info.n_constraints
     }
 
     fn max_constraint_log_degree_bound(&self) -> u32 {
         self.eval.max_constraint_log_degree_bound()
     }
 
     fn trace_log_degree_bounds(&self) -> TreeVec<ColumnVec<u32>> {
-        let InfoEvaluator { mask_offsets, .. } = self.eval.evaluate(InfoEvaluator::default());
-        mask_offsets.map(|tree_offsets| vec![self.eval.log_size(); tree_offsets.len()])
+        self.info
+            .mask_offsets
+            .as_ref()
+            .map(|tree_offsets| vec![self.eval.log_size(); tree_offsets.len()])
     }
 
     fn mask_points(
         &self,
         point: CirclePoint<SecureField>,
     ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
         let trace_step = CanonicCoset::new(self.eval.log_size()).step();
-        let InfoEvaluator { mask_offsets, .. } = self.eval.evaluate(InfoEvaluator::default());
-        mask_offsets.map_cols(|col_offsets| {
+        self.info.mask_offsets.as_ref().map_cols(|col_offsets| {
             col_offsets
                 .iter()
                 .map(|offset| point + trace_step.mul_signed(*offset).into_ef())

diff --git a/crates/prover/src/core/pcs/utils.rs b/crates/prover/src/core/pcs/utils.rs
@@ -143,6 +143,17 @@ impl<T> TreeVec<ColumnVec<T>> {
     }
 }
 
+impl<T> TreeVec<&ColumnVec<T>> {
+    pub fn map_cols<U, F: FnMut(&T) -> U>(self, mut f: F) -> TreeVec<ColumnVec<U>> {
+        TreeVec(
+            self.0
+                .into_iter()
+                .map(|column| column.iter().map(&mut f).collect())
+                .collect(),
+        )
+    }
+}
+
 impl<'a, T> From<&'a TreeVec<ColumnVec<T>>> for TreeVec<ColumnVec<&'a T>> {
     fn from(val: &'a TreeVec<ColumnVec<T>>) -> Self {
         val.as_cols_ref()