Implement Lagrange kernel constraints

air/src/air/boundary/constraint.rs

@@ -9,6 +9,7 @@ use math::{fft, polynom};
 
 // BOUNDARY CONSTRAINT
 // ================================================================================================
+
 /// The numerator portion of a boundary constraint.
 ///
 /// A boundary constraint is described by a rational function $\frac{f(x) - b(x)}{z(x)}$, where:
@@ -45,8 +46,9 @@ where
     F: FieldElement,
     E: FieldElement<BaseField = F::BaseField> + ExtensionOf<F>,
 {
-    // CONSTRUCTOR
+    // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
+
     /// Creates a new boundary constraint from the specified assertion.
     pub(super) fn new(
         assertion: Assertion<F>,

air/src/air/coefficients.rs

@@ -66,8 +66,16 @@ impl<E: FieldElement> Default for AuxTraceRandElements<E> {
 pub struct ConstraintCompositionCoefficients<E: FieldElement> {
     pub transition: Vec<E>,
     pub boundary: Vec<E>,
+    pub lagrange: Option<LagrangeConstraintsCompositionCoefficients<E>>,
 }
 
+/// Stores the constraint composition coefficients for the Lagrange kernel transition and boundary
+/// constraints.
+#[derive(Debug, Clone)]
+pub struct LagrangeConstraintsCompositionCoefficients<E: FieldElement> {
+    pub transition: Vec<E>,
+    pub boundary: E,
+}
 // DEEP COMPOSITION COEFFICIENTS
 // ================================================================================================
 /// Coefficients used in construction of DEEP composition polynomial.

air/src/air/context.rs

@@ -19,6 +19,7 @@ pub struct AirContext<B: StarkField> {
     pub(super) aux_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
     pub(super) num_main_assertions: usize,
     pub(super) num_aux_assertions: usize,
+    pub(super) lagrange_kernel_aux_column_idx: Option<usize>,
     pub(super) ce_blowup_factor: usize,
     pub(super) trace_domain_generator: B,
     pub(super) lde_domain_generator: B,
@@ -59,6 +60,7 @@ impl<B: StarkField> AirContext<B> {
             Vec::new(),
             num_assertions,
             0,
+            None,
             options,
         )
     }
@@ -91,6 +93,7 @@ impl<B: StarkField> AirContext<B> {
         aux_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
         num_main_assertions: usize,
         num_aux_assertions: usize,
+        lagrange_kernel_aux_column_idx: Option<usize>,
         options: ProofOptions,
     ) -> Self {
         assert!(
@@ -103,7 +106,7 @@ impl<B: StarkField> AirContext<B> {
             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"
@@ -119,6 +122,11 @@ impl<B: StarkField> AirContext<B> {
             );
         }
 
+        // validate Lagrange kernel aux column, if any
+        if let Some(lagrange_kernel_aux_column_idx) = lagrange_kernel_aux_column_idx {
+            assert!(lagrange_kernel_aux_column_idx < trace_info.get_aux_segment_width(0), "Lagrange kernel column index out of bounds: index={}, but only {} columns in segment", lagrange_kernel_aux_column_idx, trace_info.get_aux_segment_width(0));
+        }
+
         // determine minimum blowup factor needed to evaluate transition constraints by taking
         // the blowup factor of the highest degree constraint
         let mut ce_blowup_factor = 0;
@@ -151,6 +159,7 @@ impl<B: StarkField> AirContext<B> {
             aux_transition_constraint_degrees,
             num_main_assertions,
             num_aux_assertions,
+            lagrange_kernel_aux_column_idx,
             ce_blowup_factor,
             trace_domain_generator: B::get_root_of_unity(trace_length.ilog2()),
             lde_domain_generator: B::get_root_of_unity(lde_domain_size.ilog2()),
@@ -161,9 +170,14 @@ impl<B: StarkField> AirContext<B> {
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
+    /// Returns the trace info for an instance of a computation.
+    pub fn trace_info(&self) -> &TraceInfo {
+        &self.trace_info
+    }
+
     /// Returns length of the execution trace for an instance of a computation.
     ///
-    // This is guaranteed to be a power of two greater than or equal to 8.
+    /// This is guaranteed to be a power of two greater than or equal to 8.
     pub fn trace_len(&self) -> usize {
         self.trace_info.length()
     }
@@ -189,12 +203,13 @@ impl<B: StarkField> AirContext<B> {
         self.trace_info.length() * self.options.blowup_factor()
     }
 
-    /// Returns the number of transition constraints for a computation.
+    /// Returns the number of transition constraints for a computation, excluding the Lagrange
+    /// kernel transition constraints, which are managed separately.
     ///
-    /// The number of transition constraints is defined by the total number of transition
-    /// constraint degree descriptors (for both the main and the auxiliary trace constraints).
-    /// This number is used to determine how many transition constraint coefficients need to be
-    /// generated for merging transition constraints into a composition polynomial.
+    /// The number of transition constraints is defined by the total number of transition constraint
+    /// degree descriptors (for both the main and the auxiliary trace constraints). This number is
+    /// used to determine how many transition constraint coefficients need to be generated for
+    /// merging transition constraints into a constraint composition polynomial.
     pub fn num_transition_constraints(&self) -> usize {
         self.main_transition_constraint_degrees.len() + self.aux_transition_constraint_degrees.len()
     }
@@ -209,7 +224,18 @@ impl<B: StarkField> AirContext<B> {
         self.aux_transition_constraint_degrees.len()
     }
 
-    /// Returns the total number of assertions defined for a computation.
+    /// Returns the index of the auxiliary column which implements the Lagrange kernel, if any
+    pub fn lagrange_kernel_aux_column_idx(&self) -> Option<usize> {
+        self.lagrange_kernel_aux_column_idx
+    }
+
+    /// Returns true if the auxiliary trace segment contains a Lagrange kernel column
+    pub fn has_lagrange_kernel_aux_column(&self) -> bool {
+        self.lagrange_kernel_aux_column_idx().is_some()
+    }
+
+    /// Returns the total number of assertions defined for a computation, excluding the Lagrange
+    /// kernel assertion, which is managed separately.
     ///
     /// The number of assertions consists of the assertions placed against the main segment of an
     /// execution trace as well as assertions placed against all auxiliary trace segments.
@@ -230,24 +256,26 @@ impl<B: StarkField> AirContext<B> {
     /// Returns the number of columns needed to store the constraint composition polynomial.
     ///
     /// This is the maximum of:
-    /// 1. The maximum evaluation degree over all transition constraints minus the degree
-    /// of the transition constraint divisor divided by trace length.
+    /// 1. The maximum evaluation degree over all transition constraints minus the degree of the
+    /// transition constraint divisor divided by trace length.
     /// 2. `1`, because the constraint composition polynomial requires at least one column.
     ///
-    /// Since the degree of a constraint `C(x)` can be well approximated by
-    /// `[constraint.base + constraint.cycles.len()] * [trace_length - 1]` the degree of the
-    /// constraint composition polynomial can be approximated by:
-    /// `([constraint.base + constraint.cycles.len()] * [trace_length - 1] - [trace_length - n])`
-    /// where `constraint` is the constraint attaining the maximum and `n` is the number of
-    /// exemption points.
-    /// In the case `n = 1`, the expression simplifies to:
-    /// `[constraint.base + constraint.cycles.len() - 1] * [trace_length - 1]`
-    /// Thus, if each column is of length `trace_length`, we would need
-    /// `[constraint.base + constraint.cycles.len() - 1]` columns to store the coefficients of
-    /// the constraint composition polynomial.
-    /// This means that if the highest constraint degree is equal to `5`, the constraint
-    /// composition polynomial will require four columns and if the highest constraint degree is
-    /// equal to `7`, it will require six columns to store.
+    /// Since the degree of a constraint `C(x)` can be computed as `[constraint.base +
+    /// constraint.cycles.len()] * [trace_length - 1]` the degree of the constraint composition
+    /// polynomial can be computed as: `([constraint.base + constraint.cycles.len()] * [trace_length
+    /// - 1] - [trace_length - n])` where `constraint` is the constraint attaining the maximum and
+    /// `n` is the number of exemption points. In the case `n = 1`, the expression simplifies to:
+    /// `[constraint.base + constraint.cycles.len() - 1] * [trace_length - 1]` Thus, if each column
+    /// is of length `trace_length`, we would need `[constraint.base + constraint.cycles.len() - 1]`
+    /// columns to store the coefficients of the constraint composition polynomial. This means that
+    /// if the highest constraint degree is equal to `5`, the constraint composition polynomial will
+    /// require four columns and if the highest constraint degree is equal to `7`, it will require
+    /// six columns to store.
+    ///
+    /// Note that the Lagrange kernel constraints require only 1 column, since the degree of the
+    /// numerator is `trace_len - 1` for all transition constraints (i.e. the base degree is 1).
+    /// Hence, no matter what the degree of the divisor is for each, the degree of the fraction will
+    /// be at most `trace_len - 1`.
     pub fn num_constraint_composition_columns(&self) -> usize {
         let mut highest_constraint_degree = 0_usize;
         for degree in self

air/src/air/divisor.rs

@@ -44,24 +44,24 @@ impl<B: StarkField> ConstraintDivisor<B> {
     ///
     /// For transition constraints, the divisor polynomial $z(x)$ is always the same:
     ///
-    /// $$
-    /// z(x) = \frac{x^n - 1}{ \prod_{i=1}^k (x - g^{n-i})}
-    /// $$
+    /// $$ z(x) = \frac{x^n - 1}{ \prod_{i=1}^k (x - g^{n-i})} $$
     ///
-    /// where, $n$ is the length of the execution trace, $g$ is the generator of the trace
-    /// domain, and $k$ is the number of exemption points. The default value for $k$ is $1$.
+    /// where, $n$ is the length of the execution trace, $g$ is the generator of the trace domain,
+    /// and $k$ is the number of exemption points. The default value for $k$ is $1$.
     ///
     /// The above divisor specifies that transition constraints must hold on all steps of the
-    /// execution trace except for the last $k$ steps.
-    pub fn from_transition(trace_length: usize, num_exemptions: usize) -> Self {
-        assert!(
-            num_exemptions > 0,
-            "invalid number of transition exemptions: must be greater than zero"
-        );
-        let exemptions = (trace_length - num_exemptions..trace_length)
-            .map(|step| get_trace_domain_value_at::<B>(trace_length, step))
+    /// constraint enforcement domain except for the last $k$ steps. The constraint enforcement
+    /// domain is the entire trace in the case of transition constraints, but only a subset of the
+    /// trace for Lagrange kernel transition constraints.
+    pub fn from_transition(
+        constraint_enforcement_domain_size: usize,
+        num_exemptions: usize,
+    ) -> Self {
+        let exemptions = (constraint_enforcement_domain_size - num_exemptions
+            ..constraint_enforcement_domain_size)
+            .map(|step| get_trace_domain_value_at::<B>(constraint_enforcement_domain_size, step))
             .collect();
-        Self::new(vec![(trace_length, B::ONE)], exemptions)
+        Self::new(vec![(constraint_enforcement_domain_size, B::ONE)], exemptions)
     }
 
     /// Builds a divisor for a boundary constraint described by the assertion.

air/src/air/lagrange/boundary.rs

@@ -0,0 +1,63 @@
+use math::FieldElement;
+
+use crate::LagrangeKernelEvaluationFrame;
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct LagrangeKernelBoundaryConstraint<E>
+where
+    E: FieldElement,
+{
+    assertion_value: E,
+    composition_coefficient: E,
+}
+
+impl<E> LagrangeKernelBoundaryConstraint<E>
+where
+    E: FieldElement,
+{
+    /// Creates a new Lagrange kernel boundary constraint.
+    pub fn new(composition_coefficient: E, lagrange_kernel_rand_elements: &[E]) -> Self {
+        Self {
+            assertion_value: Self::assertion_value(lagrange_kernel_rand_elements),
+            composition_coefficient,
+        }
+    }
+
+    /// Returns the evaluation of the boundary constraint at `x`, multiplied by the composition
+    /// coefficient.
+    ///
+    /// `frame` is the evaluation frame of the Lagrange kernel column `c`, starting at `c(x)`
+    pub fn evaluate_at(&self, x: E, frame: &LagrangeKernelEvaluationFrame<E>) -> E {
+        let numerator = self.evaluate_numerator_at(frame);
+        let denominator = self.evaluate_denominator_at(x);
+
+        numerator / denominator
+    }
+
+    /// Returns the evaluation of the boundary constraint numerator, multiplied by the composition
+    /// coefficient.
+    ///
+    /// `frame` is the evaluation frame of the Lagrange kernel column `c`, starting at `c(x)` for
+    /// some `x`
+    pub fn evaluate_numerator_at(&self, frame: &LagrangeKernelEvaluationFrame<E>) -> E {
+        let trace_value = frame.inner()[0];
+        let constraint_evaluation = trace_value - self.assertion_value;
+
+        constraint_evaluation * self.composition_coefficient
+    }
+
+    /// Returns the evaluation of the boundary constraint denominator at point `x`.
+    pub fn evaluate_denominator_at(&self, x: E) -> E {
+        x - E::ONE
+    }
+
+    /// Computes the assertion value given the provided random elements.
+    pub fn assertion_value(lagrange_kernel_rand_elements: &[E]) -> E {
+        let mut assertion_value = E::ONE;
+        for &rand_ele in lagrange_kernel_rand_elements {
+            assertion_value *= E::ONE - rand_ele;
+        }
+
+        assertion_value
+    }
+}

air/src/air/lagrange/frame.rs

@@ -0,0 +1,80 @@
+use alloc::vec::Vec;
+use math::{polynom, FieldElement, StarkField};
+
+/// The evaluation frame for the Lagrange kernel.
+///
+/// The Lagrange kernel's evaluation frame is different from [`crate::EvaluationFrame`].
+/// Specifically,
+/// - it only contains evaluations from the Lagrange kernel column compared to all columns in the
+///   case of [`crate::EvaluationFrame`])
+/// - The column is evaluated at points `x`, `gx`, `g^2 x`, ..., `g^(2^(v-1)) x`, where `x` is an
+///   arbitrary point, and `g` is the trace domain generator
+#[derive(Debug, Clone)]
+pub struct LagrangeKernelEvaluationFrame<E: FieldElement> {
+    frame: Vec<E>,
+}
+
+impl<E: FieldElement> LagrangeKernelEvaluationFrame<E> {
+    // CONSTRUCTORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Constructs a Lagrange kernel evaluation frame from the raw column polynomial evaluations.
+    pub fn new(frame: Vec<E>) -> Self {
+        Self { frame }
+    }
+
+    /// Constructs an empty Lagrange kernel evaluation frame from the raw column polynomial
+    /// evaluations. The frame can subsequently be filled using [`Self::frame_mut`].
+    pub fn new_empty() -> Self {
+        Self { frame: Vec::new() }
+    }
+
+    /// Constructs the frame from the Lagrange kernel column trace polynomial coefficients for an
+    /// evaluation point.
+    pub fn from_lagrange_kernel_column_poly(lagrange_kernel_col_poly: &[E], z: E) -> Self {
+        let log_trace_len = lagrange_kernel_col_poly.len().ilog2();
+        let g = E::from(E::BaseField::get_root_of_unity(log_trace_len));
+
+        let mut frame = Vec::with_capacity(log_trace_len as usize + 1);
+
+        // push c(x)
+        frame.push(polynom::eval(lagrange_kernel_col_poly, z));
+
+        // push c(z * g), c(z * g^2), c(z * g^4), ..., c(z * g^(2^(v-1)))
+        let mut g_exp = g;
+        for _ in 0..log_trace_len {
+            let x = g_exp * z;
+            let lagrange_poly_at_x = polynom::eval(lagrange_kernel_col_poly, x);
+
+            frame.push(lagrange_poly_at_x);
+
+            // takes on the values `g`, `g^2`, `g^4`, `g^8`, ...
+            g_exp *= g_exp;
+        }
+
+        Self { frame }
+    }
+
+    // MUTATORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns a mutable reference to the inner frame.
+    pub fn frame_mut(&mut self) -> &mut Vec<E> {
+        &mut self.frame
+    }
+
+    // ACCESSORS
+    // --------------------------------------------------------------------------------------------
+
+    /// Returns a reference to the inner frame.
+    pub fn inner(&self) -> &[E] {
+        &self.frame
+    }
+
+    /// Returns the number of rows in the frame.
+    ///
+    /// This is equal to `log(trace_length) + 1`.
+    pub fn num_rows(&self) -> usize {
+        self.frame.len()
+    }
+}