chore: address feeback 3

air/src/air/context.rs

@@ -251,7 +251,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> {
+    pub fn lagrange_kernel_column_idx(&self) -> Option<usize> {
         self.trace_info.lagrange_kernel_column_idx()
     }
 

air/src/air/trace_info.rs

@@ -116,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 {}",
@@ -174,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.

prover/src/constraints/evaluator/logup_gkr.rs

@@ -71,8 +71,7 @@ where
         let mut lagrange_frame = LagrangeKernelEvaluationFrame::new_empty();
 
         let evaluator = self.air.get_logup_gkr_evaluator();
-        let s_col_constraint_divisor =
-            compute_s_col_divisor::<E>(domain.ce_domain_size(), domain, self.air.trace_length());
+        let s_col_constraint_divisor = compute_s_col_divisor::<E>(domain, self.air.trace_length());
         let s_col_idx = trace.trace_info().s_column_idx().expect("S-column should be present");
         let l_col_idx = trace
             .trace_info()
@@ -145,10 +144,10 @@ where
                 let rhs = s_cur - mean + batched_query * l_cur;
                 let lhs = s_nxt;
 
-                (rhs - lhs)
-                    * self
-                        .s_col_composition_coefficient
-                        .mul_base(s_col_constraint_divisor[step % (domain.trace_to_ce_blowup())])
+                let divisor_at_step =
+                    s_col_constraint_divisor[step % (domain.trace_to_ce_blowup())];
+
+                (rhs - lhs) * self.s_col_composition_coefficient.mul_base(divisor_at_step)
             };
 
             combined_evaluations_acc[step] +=
@@ -348,12 +347,11 @@ impl<E: FieldElement> TransitionDivisorEvaluator<E> {
 /// The divisor for the s-column is $X^n - 1$ where $n$ is the trace length. This means that
 /// we need only compute `ce_blowup` many values and thus only that many exponentiations.
 fn compute_s_col_divisor<E: FieldElement>(
-    ce_domain_size: usize,
     domain: &StarkDomain<E::BaseField>,
     trace_length: usize,
 ) -> Vec<E::BaseField> {
     let degree = trace_length as u32;
-    let mut result = Vec::with_capacity(ce_domain_size);
+    let mut result = Vec::with_capacity(domain.trace_to_ce_blowup());
 
     for row in 0..domain.trace_to_ce_blowup() {
         let x = domain.get_ce_x_at(row).exp(degree.into()) - E::BaseField::ONE;

prover/src/lib.rs

@@ -346,7 +346,7 @@ pub trait Prover {
             };
 
             trace_polys
-                .add_aux_segment(aux_segment_polys, air.context().lagrange_kernel_aux_column_idx());
+                .add_aux_segment(aux_segment_polys, air.context().lagrange_kernel_column_idx());
 
             Some(AuxTraceWithMetadata { aux_trace, aux_rand_elements, gkr_proof })
         } else {

prover/src/trace/mod.rs

@@ -142,7 +142,7 @@ pub trait Trace: Sized {
             }
 
             // then, check the Lagrange kernel assertion, if any
-            if let Some(lagrange_kernel_col_idx) = air.context().lagrange_kernel_aux_column_idx() {
+            if let Some(lagrange_kernel_col_idx) = air.context().lagrange_kernel_column_idx() {
                 let boundary_constraint_assertion_value =
                     LagrangeKernelBoundaryConstraint::assertion_value(
                         aux_rand_elements

verifier/src/composer.rs

@@ -43,7 +43,7 @@ impl<E: FieldElement> DeepComposer<E> {
             x_coordinates,
             z: [z, z * E::from(g_trace)],
             g_trace,
-            lagrange_kernel_column_idx: air.context().lagrange_kernel_aux_column_idx(),
+            lagrange_kernel_column_idx: air.context().lagrange_kernel_column_idx(),
         }
     }
 

verifier/src/evaluator.rs

@@ -89,6 +89,8 @@ pub fn evaluate_constraints<A: Air, E: FieldElement<BaseField = A::BaseField>>(
     // 3 ----- evaluate Lagrange kernel constraints ------------------------------------
 
     if let Some(lagrange_kernel_column_frame) = lagrange_kernel_frame {
+        let logup_gkr_evaluator = air.get_logup_gkr_evaluator();
+
         let lagrange_coefficients = composition_coefficients
             .lagrange
             .expect("expected Lagrange kernel composition coefficients to be present");
@@ -100,7 +102,7 @@ pub fn evaluate_constraints<A: Air, E: FieldElement<BaseField = A::BaseField>>(
 
         // Lagrange kernel constraints
 
-        let lagrange_constraints = air.get_logup_gkr_evaluator().get_lagrange_kernel_constraints(
+        let lagrange_constraints = logup_gkr_evaluator.get_lagrange_kernel_constraints(
             lagrange_coefficients,
             &gkr_data.lagrange_kernel_eval_point,
         );
@@ -115,22 +117,20 @@ pub fn evaluate_constraints<A: Air, E: FieldElement<BaseField = A::BaseField>>(
         // s-column constraints
 
         let s_col_idx = air.trace_info().s_column_idx().expect("s-column should be present");
-        let s_cur = aux_trace_frame
-            .as_ref()
-            .expect("expected aux rand elements to be present")
-            .current()[s_col_idx];
-        let s_nxt = aux_trace_frame
-            .as_ref()
-            .expect("expected aux rand elements to be present")
-            .next()[s_col_idx];
+
+        let aux_trace_frame =
+            aux_trace_frame.as_ref().expect("expected aux rand elements to be present");
+
+        let s_cur = aux_trace_frame.current()[s_col_idx];
+        let s_nxt = aux_trace_frame.next()[s_col_idx];
         let l_cur = lagrange_kernel_column_frame.inner()[0];
 
         let s_column_cc = composition_coefficients
             .s_col
             .expect("expected constraint composition coefficient for s-column to be present");
 
         let s_column_constraint =
-            air.get_logup_gkr_evaluator().get_s_column_constraints(gkr_data, s_column_cc);
+            logup_gkr_evaluator.get_s_column_constraints(gkr_data, s_column_cc);
 
         result += s_column_constraint.evaluate(air, main_trace_frame, s_cur, s_nxt, l_cur, x);
     }

winterfell/src/tests.rs

@@ -26,14 +26,14 @@ fn test_logup_gkr() {
     let trace = LogUpGkrSimple::new(2_usize.pow(7), aux_trace_width);
     let prover = LogUpGkrSimpleProver::new(aux_trace_width);
 
-    let _proof = prover.prove(trace).unwrap();
+    let proof = prover.prove(trace).unwrap();
 
     verify::<
         LogUpGkrSimpleAir,
         Blake3_256<BaseElement>,
         DefaultRandomCoin<Blake3_256<BaseElement>>,
         MerkleTree<Blake3_256<BaseElement>>,
-    >(_proof, (), &AcceptableOptions::MinConjecturedSecurity(0))
+    >(proof, (), &AcceptableOptions::MinConjecturedSecurity(0))
     .unwrap()
 }