Add support for s-column as part of the support for LogUp-GKR (#297)

air/src/air/aux.rs

@@ -5,9 +5,9 @@
 
 use alloc::vec::Vec;
 
-use math::FieldElement;
+use math::{ExtensionOf, FieldElement};
 
-use super::{lagrange::LagrangeKernelRandElements, LogUpGkrOracle};
+use super::{LagrangeKernelRandElements, LogUpGkrOracle};
 
 /// Holds the randomly generated elements used in defining the auxiliary segment of the trace.
 ///
@@ -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

@@ -250,13 +250,9 @@ impl<B: StarkField, P> AirContext<B, P> {
         self.aux_transition_constraint_degrees.len()
     }
 
-    /// 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)
-        } else {
-            None
-        }
+    /// Returns the index of the auxiliary column which implements the Lagrange kernel, if any.
+    pub fn lagrange_kernel_column_idx(&self) -> Option<usize> {
+        self.trace_info.lagrange_kernel_column_idx()
     }
 
     /// Returns true if LogUp-GKR is enabled.

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

@@ -5,7 +5,7 @@
 
 use math::FieldElement;
 
-use crate::{LagrangeKernelEvaluationFrame, LagrangeKernelRandElements};
+use super::{LagrangeKernelEvaluationFrame, LagrangeKernelRandElements};
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct LagrangeKernelBoundaryConstraint<E>

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

@@ -3,17 +3,18 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 
-mod boundary;
 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;
-use math::FieldElement;
 pub use transition::LagrangeKernelTransitionConstraints;
 
 use crate::LagrangeConstraintsCompositionCoefficients;
@@ -22,15 +23,13 @@ use crate::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> {
     /// Constructs a new [`LagrangeKernelConstraints`].
     pub fn new(
         lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients<E>,
         lagrange_kernel_rand_elements: &LagrangeKernelRandElements<E>,
-        lagrange_kernel_col_idx: usize,
     ) -> Self {
         Self {
             transition: LagrangeKernelTransitionConstraints::new(
@@ -40,7 +39,6 @@ impl<E: FieldElement> LagrangeKernelConstraints<E> {
                 lagrange_composition_coefficients.boundary,
                 lagrange_kernel_rand_elements,
             ),
-            lagrange_kernel_col_idx,
         }
     }
 }

air/src/air/logup_gkr.rs → air/src/air/logup_gkr/mod.rs

@@ -9,7 +9,15 @@ use core::marker::PhantomData;
 use crypto::{ElementHasher, RandomCoin};
 use math::{ExtensionOf, FieldElement, StarkField, ToElements};
 
-use super::{EvaluationFrame, GkrData, LagrangeKernelRandElements};
+use super::{EvaluationFrame, GkrData, LagrangeConstraintsCompositionCoefficients};
+mod s_column;
+use s_column::SColumnConstraint;
+
+mod lagrange;
+pub use lagrange::{
+    LagrangeKernelBoundaryConstraint, LagrangeKernelConstraints, LagrangeKernelEvaluationFrame,
+    LagrangeKernelRandElements, LagrangeKernelTransitionConstraints,
+};
 
 /// A trait containing the necessary information in order to run the LogUp-GKR protocol of [1].
 ///
@@ -116,6 +124,27 @@ pub trait LogUpGkrEvaluator: Clone + Sync {
             self.get_oracles().to_vec(),
         )
     }
+
+    /// Returns a new [`LagrangeKernelConstraints`].
+    fn get_lagrange_kernel_constraints<E: FieldElement<BaseField = Self::BaseField>>(
+        &self,
+        lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients<E>,
+        lagrange_kernel_rand_elements: &LagrangeKernelRandElements<E>,
+    ) -> LagrangeKernelConstraints<E> {
+        LagrangeKernelConstraints::new(
+            lagrange_composition_coefficients,
+            lagrange_kernel_rand_elements,
+        )
+    }
+
+    /// Returns a new [`SColumnConstraints`].
+    fn get_s_column_constraints<E: FieldElement<BaseField = Self::BaseField>>(
+        &self,
+        gkr_data: GkrData<E>,
+        composition_coefficient: E,
+    ) -> SColumnConstraint<E> {
+        SColumnConstraint::new(gkr_data, composition_coefficient)
+    }
 }
 
 #[derive(Clone, Default)]

air/src/air/logup_gkr/s_column.rs

@@ -0,0 +1,56 @@
+// 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::{super::Air, EvaluationFrame, GkrData};
+use crate::LogUpGkrEvaluator;
+
+/// Represents the transition constraint for the s-column, as well as the random coefficient used
+/// to linearly combine the constraint into the constraint composition polynomial.
+///
+/// The s-column implements the cohomological sum-check argument of [1] and the constraint in
+/// [`SColumnConstraint`] is exactly Eq (4) in Lemma 1 in [1].
+///
+///
+/// [1]: https://eprint.iacr.org/2021/930
+pub struct SColumnConstraint<E: FieldElement> {
+    gkr_data: GkrData<E>,
+    composition_coefficient: E,
+}
+
+impl<E: FieldElement> SColumnConstraint<E> {
+    pub fn new(gkr_data: GkrData<E>, composition_coefficient: E) -> Self {
+        Self { gkr_data, composition_coefficient }
+    }
+
+    /// Evaluates the transition constraint over the specificed main trace segment, s-column,
+    /// and Lagrange kernel evaluation frames.
+    pub fn evaluate<A>(
+        &self,
+        air: &A,
+        main_trace_frame: &EvaluationFrame<E>,
+        s_cur: E,
+        s_nxt: E,
+        l_cur: E,
+        x: E,
+    ) -> E
+    where
+        A: Air<BaseField = E::BaseField>,
+    {
+        let batched_claim = self.gkr_data.compute_batched_claim();
+        let mean = batched_claim
+            .mul_base(E::BaseField::ONE / E::BaseField::from(air.trace_length() as u32));
+
+        let mut query = vec![E::ZERO; air.get_logup_gkr_evaluator().get_oracles().len()];
+        air.get_logup_gkr_evaluator().build_query(main_trace_frame, &[], &mut query);
+        let batched_claim_at_query = self.gkr_data.compute_batched_query::<E>(&query);
+        let rhs = s_cur - mean + batched_claim_at_query * l_cur;
+        let lhs = s_nxt;
+
+        let divisor = x.exp((air.trace_length() as u32).into()) - E::ONE;
+        self.composition_coefficient * (rhs - lhs) / divisor
+    }
+}

air/src/air/mod.rs

@@ -6,7 +6,6 @@
 use alloc::{collections::BTreeMap, vec::Vec};
 
 use crypto::{RandomCoin, RandomCoinError};
-use logup_gkr::PhantomLogUpGkrEval;
 use math::{fft, ExtensibleField, ExtensionOf, FieldElement, StarkField, ToElements};
 
 use crate::ProofOptions;
@@ -29,15 +28,14 @@ pub use boundary::{BoundaryConstraint, BoundaryConstraintGroup, BoundaryConstrai
 mod transition;
 pub use transition::{EvaluationFrame, TransitionConstraintDegree, TransitionConstraints};
 
-mod lagrange;
-pub use lagrange::{
+mod logup_gkr;
+use logup_gkr::PhantomLogUpGkrEval;
+pub use logup_gkr::{
     LagrangeKernelBoundaryConstraint, LagrangeKernelConstraints, LagrangeKernelEvaluationFrame,
-    LagrangeKernelRandElements, LagrangeKernelTransitionConstraints,
+    LagrangeKernelRandElements, LagrangeKernelTransitionConstraints, LogUpGkrEvaluator,
+    LogUpGkrOracle,
 };
 
-mod logup_gkr;
-pub use logup_gkr::{LogUpGkrEvaluator, LogUpGkrOracle};
-
 mod coefficients;
 pub use coefficients::{
     ConstraintCompositionCoefficients, DeepCompositionCoefficients,
@@ -331,25 +329,6 @@ pub trait Air: Send + Sync {
         Ok(rand_elements)
     }
 
-    /// Returns a new [`LagrangeKernelConstraints`] if a Lagrange kernel auxiliary column is present
-    /// in the trace, or `None` otherwise.
-    fn get_lagrange_kernel_constraints<E: FieldElement<BaseField = Self::BaseField>>(
-        &self,
-        lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients<E>,
-        lagrange_kernel_rand_elements: &LagrangeKernelRandElements<E>,
-    ) -> Option<LagrangeKernelConstraints<E>> {
-        if self.context().logup_gkr_enabled() {
-            let col_idx = self.context().trace_info().aux_segment_width() - 1;
-            Some(LagrangeKernelConstraints::new(
-                lagrange_composition_coefficients,
-                lagrange_kernel_rand_elements,
-                col_idx,
-            ))
-        } else {
-            None
-        }
-    }
-
     /// Returns values for all periodic columns used in the computation.
     ///
     /// These values will be used to compute column values at specific states of the computation
@@ -600,7 +579,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 +589,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/air/trace_info.rs

@@ -39,6 +39,10 @@ impl TraceInfo {
     pub const MAX_META_LENGTH: usize = 65535;
     /// Maximum number of random elements in the auxiliary trace segment; currently set to 255.
     pub const MAX_RAND_SEGMENT_ELEMENTS: usize = 255;
+    /// The Lagrange kernel, if present, is the last column of the auxiliary trace.
+    pub const LAGRANGE_KERNEL_OFFSET: usize = 1;
+    /// The s-column, if present, is the second to last column of the auxiliary trace.
+    pub const S_COLUMN_OFFSET: usize = 2;
 
     // CONSTRUCTORS
     // --------------------------------------------------------------------------------------------
@@ -112,7 +116,7 @@ impl TraceInfo {
 
         // validate trace segment widths
         assert!(main_segment_width > 0, "main trace segment must consist of at least one column");
-        let full_width = main_segment_width + aux_segment_width;
+        let full_width = main_segment_width + aux_segment_width + 2 * logup_gkr as usize;
         assert!(
             full_width <= TraceInfo::MAX_TRACE_WIDTH,
             "total number of columns in the trace cannot be greater than {}, but was {}",
@@ -170,9 +174,9 @@ impl TraceInfo {
         &self.trace_meta
     }
 
-    /// Returns true if an execution trace contains the auxiliary trace segment.
+    /// Returns true if an execution trace contains an auxiliary trace segment.
     pub fn is_multi_segment(&self) -> bool {
-        self.aux_segment_width > 0
+        self.aux_segment_width > 0 || self.logup_gkr
     }
 
     /// Returns the number of columns in the main segment of an execution trace.
@@ -210,6 +214,24 @@ impl TraceInfo {
         self.logup_gkr
     }
 
+    /// Returns the index of the auxiliary column which implements the Lagrange kernel, if any.
+    pub fn lagrange_kernel_column_idx(&self) -> Option<usize> {
+        if self.logup_gkr_enabled() {
+            Some(self.aux_segment_width() - TraceInfo::LAGRANGE_KERNEL_OFFSET)
+        } else {
+            None
+        }
+    }
+
+    /// Returns the index of the auxiliary column which implements the s-column, if any.
+    pub fn s_column_idx(&self) -> Option<usize> {
+        if self.logup_gkr_enabled() {
+            Some(self.aux_segment_width() - TraceInfo::S_COLUMN_OFFSET)
+        } else {
+            None
+        }
+    }
+
     /// Returns the number of random elements needed to build all auxiliary columns, except for the
     /// Lagrange kernel column.
     pub fn get_num_aux_segment_rand_elements(&self) -> usize {