Remove trace_aux_segment_has_only_lagrange_kernel_column

air/src/air/context.rs

@@ -3,10 +3,7 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 
-use crate::{
-    air::{trace_aux_segment_has_only_lagrange_kernel_column, TransitionConstraintDegree},
-    ProofOptions, TraceInfo,
-};
+use crate::{air::TransitionConstraintDegree, ProofOptions, TraceInfo};
 use core::cmp;
 use math::StarkField;
 use utils::collections::*;
@@ -108,19 +105,14 @@ impl<B: StarkField> AirContext<B> {
         if trace_info.is_multi_segment() {
             // If the only auxiliary column is the Lagrange kernel one, then we don't require any
             // other boundary/transition constraints
-            if !trace_aux_segment_has_only_lagrange_kernel_column(
-                lagrange_kernel_aux_column_idx,
-                &trace_info,
-            ) {
-                assert!(
+            assert!(
                 !aux_transition_constraint_degrees.is_empty(),
                 "at least one transition constraint degree must be specified for auxiliary trace segments"
                 );
-                assert!(
-                    num_aux_assertions > 0,
-                    "at least one assertion must be specified against auxiliary trace segments"
-                );
-            }
+            assert!(
+                num_aux_assertions > 0,
+                "at least one assertion must be specified against auxiliary trace segments"
+            );
         } else {
             assert!(
                 aux_transition_constraint_degrees.is_empty(),

air/src/air/mod.rs

@@ -583,18 +583,3 @@ pub trait Air: Send + Sync {
         })
     }
 }
-
-// UTILS
-// ================================================================================================
-
-/// Returns true if there is only one auxiliary column, and that column is the Lagrange kernel column
-pub fn trace_aux_segment_has_only_lagrange_kernel_column(
-    lagrange_kernel_aux_column_idx: Option<usize>,
-    trace_info: &TraceInfo,
-) -> bool {
-    if lagrange_kernel_aux_column_idx.is_some() {
-        trace_info.aux_trace_width() == 1
-    } else {
-        false
-    }
-}

air/src/lib.rs

@@ -43,10 +43,9 @@ pub use options::{FieldExtension, ProofOptions};
 
 mod air;
 pub use air::{
-    trace_aux_segment_has_only_lagrange_kernel_column, Air, AirContext, Assertion,
-    AuxTraceRandElements, BoundaryConstraint, BoundaryConstraintGroup, BoundaryConstraints,
-    ConstraintCompositionCoefficients, ConstraintDivisor, DeepCompositionCoefficients,
-    EvaluationFrame, LagrangeKernelBoundaryConstraint, LagrangeKernelEvaluationFrame,
-    LagrangeKernelTransitionConstraints, TraceInfo, TransitionConstraintDegree,
-    TransitionConstraints,
+    Air, AirContext, Assertion, AuxTraceRandElements, BoundaryConstraint, BoundaryConstraintGroup,
+    BoundaryConstraints, ConstraintCompositionCoefficients, ConstraintDivisor,
+    DeepCompositionCoefficients, EvaluationFrame, LagrangeKernelBoundaryConstraint,
+    LagrangeKernelEvaluationFrame, LagrangeKernelTransitionConstraints, TraceInfo,
+    TransitionConstraintDegree, TransitionConstraints,
 };

prover/src/constraints/evaluator/default.rs

@@ -8,9 +8,8 @@ use super::{
     ConstraintEvaluationTable, ConstraintEvaluator, PeriodicValueTable, StarkDomain, TraceLde,
 };
 use air::{
-    trace_aux_segment_has_only_lagrange_kernel_column, Air, AuxTraceRandElements,
-    ConstraintCompositionCoefficients, EvaluationFrame, LagrangeKernelBoundaryConstraint,
-    LagrangeKernelTransitionConstraints, TransitionConstraints,
+    Air, AuxTraceRandElements, ConstraintCompositionCoefficients, EvaluationFrame,
+    LagrangeKernelBoundaryConstraint, LagrangeKernelTransitionConstraints, TransitionConstraints,
 };
 use math::FieldElement;
 use utils::{collections::*, iter_mut};
@@ -281,18 +280,8 @@ where
             // can just add up the results of evaluating main and auxiliary constraints.
             evaluations[0] = self.evaluate_main_transition(&main_frame, step, &mut tm_evaluations);
 
-            // Make sure to only evaluate the aux transition if the user actually defined one
-            if !trace_aux_segment_has_only_lagrange_kernel_column(
-                self.air.context().lagrange_kernel_aux_column_idx(),
-                self.air.trace_info(),
-            ) {
-                evaluations[0] += self.evaluate_aux_transition(
-                    &main_frame,
-                    &aux_frame,
-                    step,
-                    &mut ta_evaluations,
-                );
-            }
+            evaluations[0] +=
+                self.evaluate_aux_transition(&main_frame, &aux_frame, step, &mut ta_evaluations);
 
             // when in debug mode, save transition constraint evaluations
             #[cfg(debug_assertions)]

prover/src/trace/mod.rs

@@ -5,8 +5,7 @@
 
 use super::{matrix::MultiColumnIter, ColMatrix};
 use air::{
-    trace_aux_segment_has_only_lagrange_kernel_column, Air, AuxTraceRandElements, EvaluationFrame,
-    LagrangeKernelBoundaryConstraint, TraceInfo,
+    Air, AuxTraceRandElements, EvaluationFrame, LagrangeKernelBoundaryConstraint, TraceInfo,
 };
 use math::{polynom, FieldElement, StarkField};
 
@@ -180,11 +179,7 @@ pub trait Trace: Sized {
         // initialize buffers to hold evaluation frames and results of constraint evaluations
         let mut x = Self::BaseField::ONE;
         let mut main_frame = EvaluationFrame::new(self.main_trace_width());
-        let mut aux_frame = if air.trace_info().is_multi_segment()
-            && !trace_aux_segment_has_only_lagrange_kernel_column(
-                air.context().lagrange_kernel_aux_column_idx(),
-                air.trace_info(),
-            ) {
+        let mut aux_frame = if air.trace_info().is_multi_segment() {
             Some(EvaluationFrame::<E>::new(self.aux_trace_width()))
         } else {
             None

winterfell/src/tests.rs

@@ -1,7 +1,7 @@
 use super::*;
 use prover::{
     crypto::{hashers::Blake3_256, DefaultRandomCoin},
-    math::{fields::f64::BaseElement, FieldElement},
+    math::{fields::f64::BaseElement, ExtensionOf, FieldElement},
     matrix::ColMatrix,
 };
 
@@ -49,7 +49,7 @@ impl LagrangeMockTrace {
 
         Self {
             main_trace: ColMatrix::new(vec![col]),
-            info: TraceInfo::new_multi_segment(1, [1], [3], Self::TRACE_LENGTH, vec![]),
+            info: TraceInfo::new_multi_segment(1, [2], [3], Self::TRACE_LENGTH, vec![]),
         }
     }
 }
@@ -79,7 +79,7 @@ impl Trace for LagrangeMockTrace {
         let r1 = lagrange_rand_elements[1];
         let r2 = lagrange_rand_elements[2];
 
-        let col = vec![
+        let lagrange_col = vec![
             (E::ONE - r2) * (E::ONE - r1) * (E::ONE - r0),
             (E::ONE - r2) * (E::ONE - r1) * r0,
             (E::ONE - r2) * r1 * (E::ONE - r0),
@@ -90,7 +90,9 @@ impl Trace for LagrangeMockTrace {
             r2 * r1 * r0,
         ];
 
-        Some(ColMatrix::new(vec![col]))
+        let dummy_col = vec![E::ZERO; 8];
+
+        Some(ColMatrix::new(vec![lagrange_col, dummy_col]))
     }
 
     fn read_main_frame(&self, row_idx: usize, frame: &mut EvaluationFrame<BaseElement>) {
@@ -120,9 +122,9 @@ impl Air for LagrangeKernelMockAir {
             context: AirContext::new_multi_segment(
                 trace_info,
                 vec![TransitionConstraintDegree::new(1)],
-                Vec::new(),
+                vec![TransitionConstraintDegree::new(1)],
+                1,
                 1,
-                0,
                 Some(0),
                 options,
             ),
@@ -149,6 +151,27 @@ impl Air for LagrangeKernelMockAir {
     fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
         vec![Assertion::single(0, 0, BaseElement::ZERO)]
     }
+
+    fn evaluate_aux_transition<F, E>(
+        &self,
+        _main_frame: &EvaluationFrame<F>,
+        _aux_frame: &EvaluationFrame<E>,
+        _periodic_values: &[F],
+        _aux_rand_elements: &AuxTraceRandElements<E>,
+        _result: &mut [E],
+    ) where
+        F: FieldElement<BaseField = Self::BaseField>,
+        E: FieldElement<BaseField = Self::BaseField> + ExtensionOf<F>,
+    {
+        // do nothing
+    }
+
+    fn get_aux_assertions<E: FieldElement<BaseField = Self::BaseField>>(
+        &self,
+        _aux_rand_elements: &AuxTraceRandElements<E>,
+    ) -> Vec<Assertion<E>> {
+        vec![Assertion::single(1, 0, E::ZERO)]
+    }
 }
 
 // LagrangeProver