add lagrange col idx to LagrangeKernelConstraints

facebook · Mar 20, 2024 · 8c5758a · 8c5758a
1 parent 0eb5c3c
commit 8c5758a
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 21 deletions.
diff --git a/air/src/air/lagrange/mod.rs b/air/src/air/lagrange/mod.rs
@@ -14,6 +14,7 @@ use crate::{AirContext, LagrangeConstraintsCompositionCoefficients};
 pub struct LagrangeKernelConstraints<E: FieldElement> {
     pub transition: LagrangeKernelTransitionConstraints<E>,
     pub boundary: LagrangeKernelBoundaryConstraint<E>,
+    pub lagrange_kernel_col_idx: usize,
 }
 
 impl<E: FieldElement> LagrangeKernelConstraints<E> {
@@ -22,6 +23,7 @@ impl<E: FieldElement> LagrangeKernelConstraints<E> {
         context: &AirContext<E::BaseField>,
         lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients<E>,
         lagrange_kernel_rand_elements: &[E],
+        lagrange_kernel_col_idx: usize,
     ) -> Self {
         Self {
             transition: LagrangeKernelTransitionConstraints::new(
@@ -32,6 +34,7 @@ impl<E: FieldElement> LagrangeKernelConstraints<E> {
                 lagrange_composition_coefficients.boundary,
                 lagrange_kernel_rand_elements,
             ),
+            lagrange_kernel_col_idx,
         }
     }
 }
diff --git a/air/src/air/mod.rs b/air/src/air/mod.rs
@@ -299,15 +299,14 @@ pub trait Air: Send + Sync {
         lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients<E>,
         lagrange_kernel_rand_elements: &[E],
     ) -> Option<LagrangeKernelConstraints<E>> {
-        if self.context().has_lagrange_kernel_aux_column() {
-            Some(LagrangeKernelConstraints::new(
+        self.context().lagrange_kernel_aux_column_idx().map(|col_idx| {
+            LagrangeKernelConstraints::new(
                 self.context(),
                 lagrange_composition_coefficients,
                 lagrange_kernel_rand_elements,
-            ))
-        } else {
-            None
-        }
+                col_idx,
+            )
+        })
     }
 
     /// Returns values for all periodic columns used in the computation.

diff --git a/prover/src/constraints/evaluator/default.rs b/prover/src/constraints/evaluator/default.rs
@@ -300,19 +300,9 @@ where
         trace: &T,
         domain: &StarkDomain<A::BaseField>,
     ) -> Option<Vec<E>> {
-        self.lagrange_constraints_evaluator.as_ref().map(|evaluator| {
-            let lagrange_kernel_aux_column_idx = self
-                .air
-                .context()
-                .lagrange_kernel_aux_column_idx()
-                .expect("expected Lagrange kernel aux column index to be present");
-
-            evaluator.evaluate_lagrange_kernel_constraints(
-                trace,
-                lagrange_kernel_aux_column_idx,
-                domain,
-            )
-        })
+        self.lagrange_constraints_evaluator
+            .as_ref()
+            .map(|evaluator| evaluator.evaluate_lagrange_kernel_constraints(trace, domain))
     }
 
     // TRANSITION CONSTRAINT EVALUATORS

diff --git a/prover/src/constraints/evaluator/lagrange.rs b/prover/src/constraints/evaluator/lagrange.rs
@@ -46,7 +46,6 @@ impl<E: FieldElement> LagrangeKernelConstraintsBatchEvaluator<E> {
     pub fn evaluate_lagrange_kernel_constraints<T>(
         &self,
         trace: &T,
-        lagrange_kernel_column_idx: usize,
         domain: &StarkDomain<E::BaseField>,
     ) -> Vec<E>
     where
@@ -68,7 +67,7 @@ impl<E: FieldElement> LagrangeKernelConstraintsBatchEvaluator<E> {
             // compute Lagrange kernel frame
             trace.read_lagrange_kernel_frame_into(
                 step << lde_shift,
-                lagrange_kernel_column_idx,
+                self.lagrange_kernel_constraints.lagrange_kernel_col_idx,
                 &mut frame,
             );