FRI early-termination review follow-ups (2/3): deduplicate the fold schedule + terminal check

crypto/stark/src/fri/mod.rs

@@ -72,30 +72,17 @@ where
         }
     }
 
-    // Determine how many total folds are needed to reach the terminal codeword.
-    // terminal_len = 2^(blowup_log + k), clamped to initial_len for tiny inputs.
-    let initial_len = evals.len();
-    let k = final_poly_log_degree as usize;
-    // Compute `min(2^(blowup_log + k), initial_len)` without materializing the
-    // shift: an out-of-range `k` (prover self-misconfiguration) would make
-    // `2^(blowup_log + k)` overflow to 0, and the `initial_len / terminal_len`
-    // below would then divide by zero. Clamping to `initial_len` first degrades
-    // gracefully to no early termination, mirroring the verifier's own
-    // `expected_k = min(k, root_order)` clamp.
-    let terminal_shift = blowup_log as usize + k;
-    let terminal_len = if terminal_shift >= initial_len.trailing_zeros() as usize {
-        initial_len
-    } else {
-        1usize << terminal_shift
-    };
-    let total_folds = (initial_len / terminal_len).trailing_zeros() as usize;
-    let num_committed = total_folds.saturating_sub(1);
+    // Fold layout, shared with the GPU prover and the verifier — see `FriFoldLayout`.
+    let layout = crate::fri::terminal::FriFoldLayout::new(
+        evals.len().trailing_zeros(),
+        blowup_log,
+        final_poly_log_degree,
+    );
+    let num_committed = layout.num_committed;
 
     // Inverse twiddle factors for evaluation-form folding.
     let mut inv_twiddles = compute_coset_twiddles_inv(coset_offset, domain_size);
     let mut fri_layer_list = Vec::with_capacity(num_committed);
-    // Track the coset offset as it squares with each fold (needed for iFFT in terminal).
-    let mut terminal_offset = coset_offset.clone();
 
     // Commit `num_committed` folded layers to the transcript.
     for _ in 0..num_committed {
@@ -118,36 +105,38 @@ where
         // >>>> Send commitment: [pₖ]
         transcript.append_bytes(&root);
 
-        // Update twiddles and offset for the next level.
+        // Update twiddles for the next level.
         update_twiddles_in_place(&mut inv_twiddles);
-        terminal_offset = terminal_offset.square();
     }
 
     // One final fold to reach the terminal codeword (size terminal_len), unless
     // already there (total_folds == 0 means initial_len == terminal_len).
-    if total_folds > 0 {
+    if layout.total_folds > 0 {
         // <<<< Receive challenge: 𝜁_final
         let zeta = transcript.sample_field_element();
         fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles);
-        terminal_offset = terminal_offset.square();
     }
-    debug_assert_eq!(evals.len(), terminal_len, "terminal codeword size mismatch");
+    debug_assert_eq!(
+        evals.len(),
+        layout.terminal_len,
+        "terminal codeword size mismatch"
+    );
 
     // Recover the low-degree polynomial coefficients from the terminal codeword
     // and send them to the verifier.
     //
-    // The number of coefficients is determined by the *actual* terminal codeword,
-    // not the requested `final_poly_log_degree`: for tiny inputs `terminal_len`
-    // is clamped to `initial_len`, so the terminal polynomial has degree
-    // < terminal_len / 2^blowup_log = 2^(log2(terminal_len) - blowup_log). Using
-    // this clamped exponent keeps the coefficient count in lockstep with what the
-    // verifier reconstructs (`expected_k = min(k, trace_bits)`); passing the raw
-    // `final_poly_log_degree` would over-pad with zeros and break the round-trip.
-    let effective_log_degree = terminal_len.trailing_zeros() - blowup_log;
+    // The coefficient count follows the *actual* terminal codeword via
+    // `layout.effective_k` (`min(k, trace_bits)`), not the requested
+    // `final_poly_log_degree`: for tiny inputs the codeword is clamped to the
+    // full LDE, so passing the raw `k` would over-pad with zeros and break the
+    // round-trip against the verifier's own `expected_k` reconstruction.
+    // The terminal coset offset is `coset_offset^(2^total_folds)` — the offset
+    // after `total_folds` squarings (matches the GPU prover and the verifier).
+    let terminal_offset = coset_offset.pow(1u64 << layout.total_folds);
     let final_poly_coeffs = crate::fri::terminal::coeffs_from_terminal_codeword::<F, E>(
         &evals,
         &terminal_offset,
-        effective_log_degree,
+        layout.effective_k,
     );
     for c in &final_poly_coeffs {
         transcript.append_field_element(c);

crypto/stark/src/fri/terminal.rs

@@ -1,14 +1,59 @@
-//! Conversion helpers between a FRI terminal codeword and the coefficients of
-//! the low-degree polynomial it encodes.
-//!
-//! These are pure, self-contained helpers — no transcript, no FRI logic.
-//! They are used by the prover (`commit_phase_from_evaluations`) and verifier FRI step.
+//! Shared, pure FRI early-termination helpers used by both the prover
+//! (`commit_phase_from_evaluations`, `try_fri_commit_gpu`) and the verifier
+//! (`step_3_verify_fri`): the fold layout (`FriFoldLayout`) and the conversion
+//! between a terminal codeword and the coefficients of the low-degree
+//! polynomial it encodes. No transcript, no FRI protocol state.
 
 use math::fft::bit_reversing::{in_place_bit_reverse_permute, reverse_index};
 use math::field::element::FieldElement;
 use math::field::traits::{IsFFTField, IsField, IsSubFieldOf};
 use math::polynomial::Polynomial;
 
+/// The FRI early-termination fold layout.
+///
+/// Derived identically by the CPU prover (`commit_phase_from_evaluations`), the
+/// GPU prover (`try_fri_commit_gpu`), and the verifier (`fri_termination_params`).
+/// Keeping the arithmetic in one place is load-bearing: the three callers must
+/// agree exactly or proofs fail to verify, and a CPU/GPU disagreement would
+/// surface only on GPU machines.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(crate) struct FriFoldLayout {
+    /// Folds from the LDE codeword down to the terminal codeword.
+    pub(crate) total_folds: u32,
+    /// Committed (Merkle-rooted) FRI layers = `total_folds - 1`, or 0 when there
+    /// is no fold or only a single final fold.
+    pub(crate) num_committed: usize,
+    /// Terminal codeword length = `2^(blowup_log + effective_k)`.
+    pub(crate) terminal_len: usize,
+    /// Terminal polynomial log-degree bound actually used, `min(k, trace_bits)`.
+    /// This is the verifier's `expected_k` and the prover's `effective_log_degree`.
+    pub(crate) effective_k: u32,
+}
+
+impl FriFoldLayout {
+    /// Derive the layout from the LDE codeword size.
+    ///
+    /// * `lde_log`    — log2 of the LDE (deep-composition) codeword length.
+    /// * `blowup_log` — log2 of the LDE blowup factor.
+    /// * `k`          — requested `fri_final_poly_log_degree`.
+    ///
+    /// Folding stops once the codeword encodes a polynomial of degree `< 2^k`,
+    /// i.e. at codeword length `2^(blowup_log + k)`, clamped to the full LDE
+    /// size for traces too small to fold that far (the `.min(lde_log)`).
+    /// Computing `blowup_log + k` in `u32` (both small) sidesteps the
+    /// `1 << (blowup_log + k)` overflow an out-of-range `k` would otherwise cause.
+    pub(crate) fn new(lde_log: u32, blowup_log: u32, k: u32) -> Self {
+        let terminal_log = (blowup_log + k).min(lde_log);
+        let total_folds = lde_log - terminal_log;
+        Self {
+            total_folds,
+            num_committed: total_folds.saturating_sub(1) as usize,
+            terminal_len: 1usize << terminal_log,
+            effective_k: terminal_log - blowup_log,
+        }
+    }
+}
+
 /// Prover side: given a FRI terminal codeword in **bit-reversed** order,
 /// recover the `2^final_poly_log_degree` coefficients of the underlying
 /// low-degree polynomial.

crypto/stark/src/gpu_lde.rs

@@ -1629,27 +1629,21 @@ where
     // produced had this dispatch never been called.
     let transcript_snapshot = transcript.clone();
 
-    // Early-termination schedule (mirrors commit_phase_from_evaluations):
-    // terminal_len = 2^(blowup_log + final_poly_log_degree), clamped to n0.
-    // Clamp without materializing the shift so an out-of-range `k` cannot
-    // overflow `terminal_len` to 0 (see `commit_phase_from_evaluations`).
-    let k = final_poly_log_degree as usize;
-    let terminal_shift = blowup_log as usize + k;
-    let terminal_len = if terminal_shift >= n0.trailing_zeros() as usize {
-        n0
-    } else {
-        1usize << terminal_shift
-    };
-    let total_folds = (n0 / terminal_len).trailing_zeros() as usize;
+    // Fold layout, shared with the CPU prover and the verifier — see `FriFoldLayout`.
+    let layout = crate::fri::terminal::FriFoldLayout::new(
+        n0.trailing_zeros(),
+        blowup_log,
+        final_poly_log_degree,
+    );
     // The GPU path only runs above gpu_lde_threshold(). Two cases fall back to
     // the CPU path (which handles both correctly): tiny clamped traces
     // (total_folds == 0), and terminal_len == 1 (blowup_log + k == 0), whose
     // final fold would reach n_out == 1 and trip `fold_and_commit_layer`'s
     // `n_out >= 2` assert. The final fold below is therefore always n_out >= 2.
-    if total_folds == 0 || terminal_len < 2 {
+    if layout.total_folds == 0 || layout.terminal_len < 2 {
         return None;
     }
-    let num_committed = total_folds - 1;
+    let num_committed = layout.num_committed;
     let mut fri_layer_list: Vec<FriLayer<E, FriLayerMerkleTreeBackend<E>>> =
         Vec::with_capacity(num_committed);
 
@@ -1696,16 +1690,15 @@ where
             return None;
         }
     };
-    debug_assert_eq!(terminal_evals_u64.len(), terminal_len * 3);
+    debug_assert_eq!(terminal_evals_u64.len(), layout.terminal_len * 3);
     let terminal_codeword = u64_to_ext3_vec::<E>(&terminal_evals_u64);
 
     // CPU-side coefficient extraction, identical to commit_phase_from_evaluations.
-    let terminal_offset = coset_offset.pow(1u64 << total_folds);
-    let effective_log_degree = terminal_len.trailing_zeros() - blowup_log;
+    let terminal_offset = coset_offset.pow(1u64 << layout.total_folds);
     let final_poly_coeffs = crate::fri::terminal::coeffs_from_terminal_codeword::<F, E>(
         &terminal_codeword,
         &terminal_offset,
-        effective_log_degree,
+        layout.effective_k,
     );
 
     // >>>> Send the final polynomial coefficients.