Add support for s-column as part of the support for LogUp-GKR

air/src/air/aux.rs

@@ -5,7 +5,7 @@
 
 use alloc::vec::Vec;
 
-use math::FieldElement;
+use math::{ExtensionOf, FieldElement};
 
 use super::{lagrange::LagrangeKernelRandElements, LogUpGkrOracle};
 
@@ -130,7 +130,11 @@ impl<E: FieldElement> GkrData<E> {
                 .fold(E::ZERO, |acc, (a, b)| acc + *a * *b)
     }
 
-    pub fn compute_batched_query(&self, query: &[E::BaseField]) -> E {
+    pub fn compute_batched_query<F>(&self, query: &[F]) -> E
+    where
+        F: FieldElement<BaseField = E::BaseField>,
+        E: ExtensionOf<F>,
+    {
         E::from(query[0])
             + query
                 .iter()

air/src/air/context.rs

@@ -8,6 +8,7 @@ use core::cmp;
 
 use math::StarkField;
 
+use super::LAGRANGE_KERNEL_OFFSET;
 use crate::{air::TransitionConstraintDegree, ProofOptions, TraceInfo};
 
 // AIR CONTEXT
@@ -253,7 +254,7 @@ impl<B: StarkField, P> AirContext<B, P> {
     /// Returns the index of the auxiliary column which implements the Lagrange kernel, if any
     pub fn lagrange_kernel_aux_column_idx(&self) -> Option<usize> {
         if self.logup_gkr_enabled() {
-            Some(self.trace_info().aux_segment_width() - 1)
+            Some(self.trace_info().aux_segment_width() - LAGRANGE_KERNEL_OFFSET)
         } else {
             None
         }

air/src/air/logup_gkr.rs

@@ -11,6 +11,11 @@ use math::{ExtensionOf, FieldElement, StarkField, ToElements};
 
 use super::{EvaluationFrame, GkrData, LagrangeKernelRandElements};
 
+// CONSTANTS
+// ===============================================================================================
+pub const LAGRANGE_KERNEL_OFFSET: usize = 1;
+pub const S_COLUMN_OFFSET: usize = 2;
+
 /// A trait containing the necessary information in order to run the LogUp-GKR protocol of [1].
 ///
 /// The trait contains useful information for running the GKR protocol as well as for implementing
@@ -200,3 +205,62 @@ pub enum LogUpGkrOracle<B: StarkField> {
     /// must be a power of 2.
     PeriodicValue(Vec<B>),
 }
+
+#[derive(Clone, Default)]
+pub struct DummyLogUpGkrEval<B: StarkField, P: Clone + Send + Sync + ToElements<B>> {
+    _field: PhantomData<B>,
+    _public_inputs: PhantomData<P>,
+}
+
+impl<B, P> LogUpGkrEvaluator for DummyLogUpGkrEval<B, P>
+where
+    B: StarkField,
+    P: Clone + Send + Sync + ToElements<B>,
+{
+    type BaseField = B;
+
+    type PublicInputs = P;
+
+    fn get_oracles(&self) -> &[LogUpGkrOracle<Self::BaseField>] {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn get_num_rand_values(&self) -> usize {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn get_num_fractions(&self) -> usize {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn max_degree(&self) -> usize {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn build_query<E>(&self, _frame: &EvaluationFrame<E>, _periodic_values: &[E], _query: &mut [E])
+    where
+        E: FieldElement<BaseField = Self::BaseField>,
+    {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn evaluate_query<F, E>(
+        &self,
+        _query: &[F],
+        _rand_values: &[E],
+        _numerator: &mut [E],
+        _denominator: &mut [E],
+    ) where
+        F: FieldElement<BaseField = Self::BaseField>,
+        E: FieldElement<BaseField = Self::BaseField> + ExtensionOf<F>,
+    {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn compute_claim<E>(&self, _inputs: &Self::PublicInputs, _rand_values: &[E]) -> E
+    where
+        E: FieldElement<BaseField = Self::BaseField>,
+    {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+}

air/src/air/mod.rs

@@ -36,7 +36,9 @@ pub use lagrange::{
 };
 
 mod logup_gkr;
-pub use logup_gkr::{LogUpGkrEvaluator, LogUpGkrOracle};
+pub use logup_gkr::{
+    LogUpGkrEvaluator, LogUpGkrOracle, LAGRANGE_KERNEL_OFFSET, S_COLUMN_OFFSET,
+};
 
 mod coefficients;
 pub use coefficients::{
@@ -600,7 +602,7 @@ pub trait Air: Send + Sync {
             None
         };
 
-        let s_col = if self.context().logup_gkr_enabled() {
+        let s_col_cc = if self.context().logup_gkr_enabled() {
             Some(public_coin.draw()?)
         } else {
             None
@@ -610,7 +612,7 @@ pub trait Air: Send + Sync {
             trace: t_coefficients,
             constraints: c_coefficients,
             lagrange: lagrange_cc,
-            s_col,
+            s_col: s_col_cc,
         })
     }
 }

air/src/air/tests.rs

@@ -225,7 +225,6 @@ impl MockAir {
 impl Air for MockAir {
     type BaseField = BaseElement;
     type PublicInputs = ();
-    //type LogUpGkrEvaluator = DummyLogUpGkrEval<Self::BaseField, ()>;
 
     fn new(trace_info: TraceInfo, _pub_inputs: (), _options: ProofOptions) -> Self {
         let num_assertions = trace_info.meta()[0] as usize;

air/src/lib.rs

@@ -47,6 +47,7 @@ pub use air::{
     DeepCompositionCoefficients, EvaluationFrame, GkrData,
     LagrangeConstraintsCompositionCoefficients, LagrangeKernelBoundaryConstraint,
     LagrangeKernelConstraints, LagrangeKernelEvaluationFrame, LagrangeKernelRandElements,
-    LagrangeKernelTransitionConstraints, LogUpGkrEvaluator, LogUpGkrOracle, TraceInfo,
-    TransitionConstraintDegree, TransitionConstraints,
+    LagrangeKernelTransitionConstraints, LogUpGkrEvaluator, LogUpGkrOracle, 
+    TraceInfo, TransitionConstraintDegree, TransitionConstraints, LAGRANGE_KERNEL_OFFSET,
+    S_COLUMN_OFFSET,
 };

prover/src/constraints/evaluator/default.rs

@@ -13,9 +13,9 @@ use utils::iter_mut;
 use utils::{iterators::*, rayon};
 
 use super::{
-    super::EvaluationTableFragment, lagrange::LagrangeKernelConstraintsBatchEvaluator,
-    BoundaryConstraints, CompositionPolyTrace, ConstraintEvaluationTable, ConstraintEvaluator,
-    PeriodicValueTable, StarkDomain, TraceLde,
+    super::EvaluationTableFragment, logup_gkr::LogUpGkrConstraintsEvaluator, BoundaryConstraints,
+    CompositionPolyTrace, ConstraintEvaluationTable, ConstraintEvaluator, PeriodicValueTable,
+    StarkDomain, TraceLde,
 };
 
 // CONSTANTS
@@ -40,7 +40,7 @@ pub struct DefaultConstraintEvaluator<'a, A: Air, E: FieldElement<BaseField = A:
     air: &'a A,
     boundary_constraints: BoundaryConstraints<E>,
     transition_constraints: TransitionConstraints<E>,
-    lagrange_constraints_evaluator: Option<LagrangeKernelConstraintsBatchEvaluator<E>>,
+    logup_gkr_constraints_evaluator: Option<LogUpGkrConstraintsEvaluator<'a, E, A>>,
     aux_rand_elements: Option<AuxRandElements<E>>,
     periodic_values: PeriodicValueTable<E::BaseField>,
 }
@@ -117,10 +117,10 @@ where
         evaluation_table.validate_transition_degrees();
 
         // combine all constraint evaluations into a single column, including the evaluations of the
-        // Lagrange kernel constraints (if present)
+        // LogUp-GKR constraints (if present)
         let combined_evaluations = {
             let mut constraints_evaluations = evaluation_table.combine();
-            self.evaluate_lagrange_kernel_constraints(trace, domain, &mut constraints_evaluations);
+            self.evaluate_logup_gkr_constraints(trace, domain, &mut constraints_evaluations);
 
             constraints_evaluations
         };
@@ -158,18 +158,21 @@ where
             &composition_coefficients.boundary,
         );
 
-        let lagrange_constraints_evaluator = if air.context().logup_gkr_enabled() {
+        let logup_gkr_constraints_evaluator = if air.context().logup_gkr_enabled() {
             let aux_rand_elements =
                 aux_rand_elements.as_ref().expect("expected aux rand elements to be present");
-            let lagrange_rand_elements = aux_rand_elements
-                .lagrange()
-                .expect("expected lagrange rand elements to be present");
-            Some(LagrangeKernelConstraintsBatchEvaluator::new(
+
+            Some(LogUpGkrConstraintsEvaluator::new(
                 air,
-                lagrange_rand_elements.clone(),
+                aux_rand_elements
+                    .gkr_data()
+                    .expect("expected LogUp-GKR randomness to be present"),
                 composition_coefficients
                     .lagrange
                     .expect("expected Lagrange kernel composition coefficients to be present"),
+                composition_coefficients
+                    .s_col
+                    .expect("expected s-column composition coefficient to be present"),
             ))
         } else {
             None
@@ -179,7 +182,7 @@ where
             air,
             boundary_constraints,
             transition_constraints,
-            lagrange_constraints_evaluator,
+            logup_gkr_constraints_evaluator,
             aux_rand_elements,
             periodic_values,
         }
@@ -295,22 +298,22 @@ where
         }
     }
 
-    /// If present, evaluates the Lagrange kernel constraints over the constraint evaluation domain.
+    /// If present, evaluates the LogUp-GKR constraints over the constraint evaluation domain.
     /// The evaluation of each constraint (both boundary and transition) is divided by its divisor,
     /// multiplied by its composition coefficient, the result of which is added to
     /// `combined_evaluations_accumulator`.
     ///
     /// Specifically, `combined_evaluations_accumulator` is a buffer whose length is the size of the
     /// constraint evaluation domain, where each index contains combined evaluations of other
     /// constraints in the system.
-    fn evaluate_lagrange_kernel_constraints<T: TraceLde<E>>(
+    fn evaluate_logup_gkr_constraints<T: TraceLde<E>>(
         &self,
         trace: &T,
         domain: &StarkDomain<A::BaseField>,
         combined_evaluations_accumulator: &mut [E],
     ) {
-        if let Some(ref lagrange_constraints_evaluator) = self.lagrange_constraints_evaluator {
-            lagrange_constraints_evaluator.evaluate_constraints(
+        if let Some(ref logup_gkr_constraints_evaluator) = self.logup_gkr_constraints_evaluator {
+            logup_gkr_constraints_evaluator.evaluate_constraints(
                 trace,
                 domain,
                 combined_evaluations_accumulator,