feat: add support for s-col and merge it with the that of Lagrange ke…

…rnel under LogUp-GKR
facebook · Aug 15, 2024 · 15ad52f · 15ad52f
1 parent f28a64d
commit 15ad52f
Show file tree

Hide file tree

Showing 8 changed files with 21 additions and 12 deletions.
diff --git a/air/src/air/context.rs b/air/src/air/context.rs
@@ -10,6 +10,8 @@ use math::StarkField;
 
 use crate::{air::TransitionConstraintDegree, ProofOptions, TraceInfo};
 
+use super::LAGRANGE_KERNEL_OFFSET;
+
 // AIR CONTEXT
 // ================================================================================================
 /// STARK parameters and trace properties for a specific execution of a computation.
@@ -242,7 +244,7 @@ impl<B: StarkField> AirContext<B> {
     /// 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.is_with_logup_gkr() {
-            Some(self.trace_info().aux_segment_width() - 1)
+            Some(self.trace_info().aux_segment_width() - LAGRANGE_KERNEL_OFFSET)
         } else {
             None
         }

diff --git a/air/src/air/logup_gkr.rs b/air/src/air/logup_gkr.rs
@@ -9,6 +9,14 @@ use math::{fields::f64::BaseElement, ExtensionOf, FieldElement, StarkField, ToEl
 
 use super::EvaluationFrame;
 
+// CONSTANTS
+// ===============================================================================================
+pub const LAGRANGE_KERNEL_OFFSET: usize = 1;
+pub const S_COLUMN_OFFSET: usize = 2;
+
+// LOGUP-GKR EVALUATOR TRAIT
+// ===============================================================================================
+
 pub trait LogUpGkrEvaluator: Clone + Sync {
     /// Defines the base field of the evaluator.
     type BaseField: StarkField;

diff --git a/air/src/air/mod.rs b/air/src/air/mod.rs
@@ -35,7 +35,7 @@ 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::{

diff --git a/air/src/lib.rs b/air/src/lib.rs
@@ -48,5 +48,5 @@ pub use air::{
     LagrangeConstraintsCompositionCoefficients, LagrangeKernelBoundaryConstraint,
     LagrangeKernelConstraints, LagrangeKernelEvaluationFrame, LagrangeKernelRandElements,
     LagrangeKernelTransitionConstraints, LogUpGkrEvaluator, LogUpGkrOracle, TraceInfo,
-    TransitionConstraintDegree, TransitionConstraints,
+    TransitionConstraintDegree, TransitionConstraints, LAGRANGE_KERNEL_OFFSET, S_COLUMN_OFFSET,
 };
diff --git a/prover/src/constraints/evaluator/default.rs b/prover/src/constraints/evaluator/default.rs
@@ -13,7 +13,7 @@ use utils::iter_mut;
 use utils::{iterators::*, rayon};
 
 use super::{
-    super::EvaluationTableFragment, lagrange::LogUpGkrConstraintsEvaluator, BoundaryConstraints,
+    super::EvaluationTableFragment, logup_gkr::LogUpGkrConstraintsEvaluator, BoundaryConstraints,
     CompositionPolyTrace, ConstraintEvaluationTable, ConstraintEvaluator, PeriodicValueTable,
     StarkDomain, TraceLde,
 };

diff --git a/prover/src/constraints/evaluator/lagrange.rs → ...er/src/constraints/evaluator/logup_gkr.rs b/prover/src/constraints/evaluator/lagrange.rs → ...er/src/constraints/evaluator/logup_gkr.rs
@@ -6,8 +6,7 @@
 use alloc::vec::Vec;
 
 use air::{
-    Air, EvaluationFrame, GkrData, LagrangeConstraintsCompositionCoefficients,
-    LagrangeKernelConstraints, LagrangeKernelEvaluationFrame, LogUpGkrEvaluator,
+    Air, EvaluationFrame, GkrData, LagrangeConstraintsCompositionCoefficients, LagrangeKernelConstraints, LagrangeKernelEvaluationFrame, LogUpGkrEvaluator, LAGRANGE_KERNEL_OFFSET, S_COLUMN_OFFSET
 };
 use math::{batch_inversion, FieldElement};
 
@@ -49,7 +48,7 @@ where
 
     /// Evaluates the transition and boundary constraints. Specifically, the constraint evaluations
     /// are divided by their corresponding divisors, and the resulting terms are linearly combined
-    /// using the composition coefficients.
+    /// using the constraint composition coefficients.
     ///
     /// Writes the evaluations in `combined_evaluations_acc` at the corresponding (constraint
     /// evaluation) domain index.
@@ -73,8 +72,8 @@ where
         let evaluator = self.air.get_logup_gkr_evaluator::<E>();
         let s_col_constraint_divisor =
             compute_s_col_divisor::<E>(domain.ce_domain_size(), domain, self.air.trace_length());
-        let s_col_idx = trace.trace_info().aux_segment_width() - 2;
-        let l_col_idx = trace.trace_info().aux_segment_width() - 1;
+        let s_col_idx = trace.trace_info().aux_segment_width() - S_COLUMN_OFFSET;
+        let l_col_idx = trace.trace_info().aux_segment_width() - LAGRANGE_KERNEL_OFFSET;
         let mut main_frame = EvaluationFrame::new(trace.trace_info().main_trace_width());
         let mut aux_frame = EvaluationFrame::new(trace.trace_info().aux_segment_width());
 

diff --git a/prover/src/constraints/evaluator/mod.rs b/prover/src/constraints/evaluator/mod.rs
@@ -14,7 +14,7 @@ pub use default::DefaultConstraintEvaluator;
 mod boundary;
 use boundary::BoundaryConstraints;
 
-mod lagrange;
+mod logup_gkr;
 
 mod periodic_table;
 use periodic_table::PeriodicValueTable;

diff --git a/verifier/src/evaluator.rs b/verifier/src/evaluator.rs
@@ -7,7 +7,7 @@ use alloc::vec::Vec;
 
 use air::{
     Air, AuxRandElements, ConstraintCompositionCoefficients, EvaluationFrame,
-    LagrangeKernelEvaluationFrame, LogUpGkrEvaluator,
+    LagrangeKernelEvaluationFrame, LogUpGkrEvaluator, S_COLUMN_OFFSET,
 };
 use math::{polynom, FieldElement};
 
@@ -116,7 +116,7 @@ pub fn evaluate_constraints<A: Air, E: FieldElement<BaseField = A::BaseField>>(
 
         // s-column constraints
 
-        let s_col_idx = air.trace_info().aux_segment_width() - 2;
+        let s_col_idx = air.trace_info().aux_segment_width() - S_COLUMN_OFFSET;
         let s_cur = aux_trace_frame
             .as_ref()
             .expect("expected aux rand elements to be present")