chore: re-organise logup_gkr module

air/src/air/logup_gkr/lagrange/boundary.rs

@@ -0,0 +1,71 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+//
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+use math::FieldElement;
+
+use super::{LagrangeKernelEvaluationFrame, LagrangeKernelRandElements};
+
+#[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: &LagrangeKernelRandElements<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: &LagrangeKernelRandElements<E>) -> E {
+        let mut assertion_value = E::ONE;
+        for &rand_ele in lagrange_kernel_rand_elements.as_ref() {
+            assertion_value *= E::ONE - rand_ele;
+        }
+
+        assertion_value
+    }
+}

air/src/air/logup_gkr/lagrange/frame.rs

@@ -0,0 +1,86 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+//
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+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()
+    }
+}

air/src/air/logup_gkr/lagrange/mod.rs

@@ -0,0 +1,81 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+//
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+use alloc::vec::Vec;
+use core::ops::Deref;
+
+use math::FieldElement;
+
+mod boundary;
+pub use boundary::LagrangeKernelBoundaryConstraint;
+
+mod frame;
+pub use frame::LagrangeKernelEvaluationFrame;
+
+mod transition;
+pub use transition::LagrangeKernelTransitionConstraints;
+
+use crate::LagrangeConstraintsCompositionCoefficients;
+
+/// Represents the Lagrange kernel transition and boundary constraints.
+pub struct LagrangeKernelConstraints<E: FieldElement> {
+    pub transition: LagrangeKernelTransitionConstraints<E>,
+    pub boundary: LagrangeKernelBoundaryConstraint<E>,
+}
+
+impl<E: FieldElement> LagrangeKernelConstraints<E> {
+    /// Constructs a new [`LagrangeKernelConstraints`].
+    pub fn new(
+        lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients<E>,
+        lagrange_kernel_rand_elements: &LagrangeKernelRandElements<E>,
+    ) -> Self {
+        Self {
+            transition: LagrangeKernelTransitionConstraints::new(
+                lagrange_composition_coefficients.transition,
+            ),
+            boundary: LagrangeKernelBoundaryConstraint::new(
+                lagrange_composition_coefficients.boundary,
+                lagrange_kernel_rand_elements,
+            ),
+        }
+    }
+}
+
+/// Holds the randomly generated elements needed to build the Lagrange kernel auxiliary column.
+///
+/// The Lagrange kernel consists of evaluating the function $eq(x, r)$, where $x$ is the binary
+/// decomposition of the row index, and $r$ is some random point. The "Lagrange kernel random
+/// elements" refer to this (multidimensional) point $r$.
+#[derive(Debug, Clone, Default)]
+pub struct LagrangeKernelRandElements<E> {
+    elements: Vec<E>,
+}
+
+impl<E> LagrangeKernelRandElements<E> {
+    /// Creates a new [`LagrangeKernelRandElements`].
+    pub fn new(elements: Vec<E>) -> Self {
+        Self { elements }
+    }
+}
+
+impl<E> Deref for LagrangeKernelRandElements<E> {
+    type Target = Vec<E>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.elements
+    }
+}
+
+impl<E> AsRef<[E]> for LagrangeKernelRandElements<E> {
+    fn as_ref(&self) -> &[E] {
+        &self.elements
+    }
+}
+
+impl<E> From<LagrangeKernelRandElements<E>> for Vec<E> {
+    fn from(lagrange_rand_elements: LagrangeKernelRandElements<E>) -> Self {
+        lagrange_rand_elements.elements
+    }
+}

air/src/air/logup_gkr/lagrange/transition.rs

@@ -0,0 +1,141 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+//
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+use alloc::vec::Vec;
+
+use math::{ExtensionOf, FieldElement};
+
+use crate::{ConstraintDivisor, LagrangeKernelEvaluationFrame};
+
+/// Represents the transition constraints for the Lagrange kernel column, as well as the random
+/// coefficients used to linearly combine all the constraints.
+///
+/// There are `log(trace_len)` constraints, each with its own divisor, as described in
+/// [this issue](https://github.com/facebook/winterfell/issues/240).
+pub struct LagrangeKernelTransitionConstraints<E: FieldElement> {
+    lagrange_constraint_coefficients: Vec<E>,
+    divisors: Vec<ConstraintDivisor<E::BaseField>>,
+}
+
+impl<E: FieldElement> LagrangeKernelTransitionConstraints<E> {
+    /// Creates a new [`LagrangeKernelTransitionConstraints`], which represents the Lagrange kernel
+    /// transition constraints as well as the random coefficients necessary to combine the
+    /// constraints together.
+    pub fn new(lagrange_constraint_coefficients: Vec<E>) -> Self {
+        let num_lagrange_kernel_transition_constraints = lagrange_constraint_coefficients.len();
+
+        let divisors = {
+            let mut divisors = Vec::with_capacity(num_lagrange_kernel_transition_constraints);
+            for i in 0..num_lagrange_kernel_transition_constraints {
+                let constraint_domain_size = 2_usize.pow(i as u32);
+                let divisor = ConstraintDivisor::from_transition(constraint_domain_size, 0);
+
+                divisors.push(divisor);
+            }
+            divisors
+        };
+
+        Self {
+            lagrange_constraint_coefficients,
+            divisors,
+        }
+    }
+
+    /// Evaluates the numerator of the `constraint_idx`th transition constraint.
+    pub fn evaluate_ith_numerator<F>(
+        &self,
+        lagrange_kernel_column_frame: &LagrangeKernelEvaluationFrame<E>,
+        lagrange_kernel_rand_elements: &[E],
+        constraint_idx: usize,
+    ) -> E
+    where
+        F: FieldElement<BaseField = E::BaseField>,
+        E: ExtensionOf<F>,
+    {
+        let c = lagrange_kernel_column_frame.inner();
+        let v = c.len() - 1;
+        let r = lagrange_kernel_rand_elements;
+        let k = constraint_idx + 1;
+
+        let eval = (r[v - k] * c[0]) - ((E::ONE - r[v - k]) * c[v - k + 1]);
+
+        self.lagrange_constraint_coefficients[constraint_idx].mul_base(eval)
+    }
+
+    /// Evaluates the divisor of the `constraint_idx`th transition constraint.
+    pub fn evaluate_ith_divisor<F>(&self, constraint_idx: usize, x: F) -> E
+    where
+        F: FieldElement<BaseField = E::BaseField>,
+        E: ExtensionOf<F>,
+    {
+        self.divisors[constraint_idx].evaluate_at(x.into())
+    }
+
+    /// Evaluates the transition constraints over the specificed Lagrange kernel evaluation frame,
+    /// and combines them.
+    ///
+    /// By "combining transition constraints evaluations", we mean computing a linear combination of
+    /// each transition constraint evaluation, where each transition evaluation is divided by its
+    /// corresponding divisor.
+    pub fn evaluate_and_combine<F>(
+        &self,
+        lagrange_kernel_column_frame: &LagrangeKernelEvaluationFrame<E>,
+        lagrange_kernel_rand_elements: &[E],
+        x: F,
+    ) -> E
+    where
+        F: FieldElement<BaseField = E::BaseField>,
+        E: ExtensionOf<F>,
+    {
+        let numerators = self
+            .evaluate_numerators::<F>(lagrange_kernel_column_frame, lagrange_kernel_rand_elements);
+
+        numerators
+            .iter()
+            .zip(self.divisors.iter())
+            .fold(E::ZERO, |acc, (&numerator, divisor)| {
+                let z = divisor.evaluate_at(x);
+
+                acc + (numerator / z.into())
+            })
+    }
+
+    /// Returns the number of constraints.
+    pub fn num_constraints(&self) -> usize {
+        self.lagrange_constraint_coefficients.len()
+    }
+
+    // HELPERS
+    // ---------------------------------------------------------------------------------------------
+
+    /// Evaluates the transition constraints' numerators over the specified Lagrange kernel
+    /// evaluation frame.
+    fn evaluate_numerators<F>(
+        &self,
+        lagrange_kernel_column_frame: &LagrangeKernelEvaluationFrame<E>,
+        lagrange_kernel_rand_elements: &[E],
+    ) -> Vec<E>
+    where
+        F: FieldElement<BaseField = E::BaseField>,
+        E: ExtensionOf<F>,
+    {
+        let log2_trace_len = lagrange_kernel_column_frame.num_rows() - 1;
+        let mut transition_evals = vec![E::ZERO; log2_trace_len];
+
+        let c = lagrange_kernel_column_frame.inner();
+        let v = c.len() - 1;
+        let r = lagrange_kernel_rand_elements;
+
+        for k in 1..v + 1 {
+            transition_evals[k - 1] = (r[v - k] * c[0]) - ((E::ONE - r[v - k]) * c[v - k + 1]);
+        }
+
+        transition_evals
+            .into_iter()
+            .zip(self.lagrange_constraint_coefficients.iter())
+            .map(|(transition_eval, &coeff)| coeff.mul_base(transition_eval))
+            .collect()
+    }
+}