Feat/fri early termination

crypto/stark/benches/profile_prover.rs

@@ -21,6 +21,7 @@ fn main() {
         fri_number_of_queries: 100,
         coset_offset: 3,
         grinding_factor: 0,
+        fri_final_poly_log_degree: 7,
     };
 
     let num_columns = 16;

crypto/stark/benches/prover_benchmark.rs

@@ -61,6 +61,7 @@ fn benchmark_proof_options() -> ProofOptions {
         fri_number_of_queries: 30,
         coset_offset: 3,
         grinding_factor: 0,
+        fri_final_poly_log_degree: 7,
     }
 }
 

crypto/stark/src/domain.rs

@@ -49,7 +49,6 @@ use super::traits::AIR;
 /// Full domain with pre-computed roots of unity. Used by the prover which needs
 /// all elements for FFT operations.
 pub struct Domain<F: IsFFTField> {
-    pub(crate) root_order: u32,
     pub(crate) lde_roots_of_unity_coset: Vec<FieldElement<F>>,
     pub(crate) trace_primitive_root: FieldElement<F>,
     pub(crate) trace_roots_of_unity: Vec<FieldElement<F>>,
@@ -88,7 +87,6 @@ impl<F: IsFFTField> Domain<F> {
         .unwrap();
 
         Self {
-            root_order,
             lde_roots_of_unity_coset,
             trace_primitive_root,
             trace_roots_of_unity,

crypto/stark/src/fri/mod.rs

@@ -1,6 +1,7 @@
 pub mod fri_commitment;
 pub mod fri_decommit;
 pub(crate) mod fri_functions;
+pub(crate) mod terminal;
 
 use crypto::fiat_shamir::is_transcript::IsStarkTranscript;
 use math::field::element::FieldElement;
@@ -16,25 +17,28 @@ use self::fri_functions::{
 };
 
 /// FRI commit phase from pre-computed bit-reversed evaluations, skipping the
-/// initial FFT. Use this when the caller already has the evaluation vector
-/// (e.g. from a fused LDE pipeline).
+/// initial FFT. Stops folding when the remaining codeword encodes a polynomial
+/// of degree < 2^`final_poly_log_degree` with blowup 2^`blowup_log`, and
+/// returns the coefficient vector of that terminal polynomial.
 ///
 /// The `T: Clone` and `F/E: 'static` bounds are required by the cuda GPU
 /// fast path (`try_fri_commit_gpu` snapshots the transcript and TypeId-
 /// checks the field types). They are present unconditionally (including
 /// in builds without the `cuda` feature) to keep one stable signature.
+#[allow(clippy::type_complexity)]
 pub fn commit_phase_from_evaluations<
     F: IsFFTField + IsSubFieldOf<E> + 'static,
-    E: IsField + 'static,
+    E: IsField + 'static + Send + Sync,
     T: IsStarkTranscript<E, F> + Clone,
 >(
-    number_layers: usize,
     mut evals: Vec<FieldElement<E>>,
     transcript: &mut T,
     coset_offset: &FieldElement<F>,
     domain_size: usize,
+    blowup_log: u32,
+    final_poly_log_degree: u32,
 ) -> (
-    FieldElement<E>,
+    Vec<FieldElement<E>>,
     Vec<FriLayer<E, FriLayerMerkleTreeBackend<E>>>,
 )
 where
@@ -50,27 +54,52 @@ where
     // had never been tried.
     #[cfg(feature = "cuda")]
     {
+        // Try the GPU early-termination FRI commit first. `try_fri_commit_gpu`
+        // drives the same commit phase on-device (Goldilocks + Ext3, above the
+        // LDE size threshold, and only when folding actually happens) and returns
+        // `Some` with the final-polynomial coefficients. It returns `None` on any
+        // precondition miss or cudarc error — restoring the transcript first — so
+        // the CPU path below then runs as if the GPU had never been tried.
         if let Some(result) = crate::gpu_lde::try_fri_commit_gpu::<F, E, T>(
-            number_layers,
             &evals,
             transcript,
             coset_offset,
             domain_size,
+            blowup_log,
+            final_poly_log_degree,
         ) {
             return result;
         }
     }
 
+    // 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);
+
     // 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();
 
-    // The loop commits `number_layers - 1` folded layers; the final fold below
-    // produces the (uncommitted) last value.
-    let num_committed_layers = number_layers.saturating_sub(1);
-    let mut fri_layer_list = Vec::with_capacity(num_committed_layers);
-
-    for _ in 0..num_committed_layers {
-        // <<<< Receive challenge 𝜁ₖ₋₁
+    // Commit `num_committed` folded layers to the transcript.
+    for _ in 0..num_committed {
+        // <<<< Receive challenge 𝜁ₖ
         let zeta = transcript.sample_field_element();
 
         // Fold evaluations in-place (no FFT needed).
@@ -89,25 +118,42 @@ where
         // >>>> Send commitment: [pₖ]
         transcript.append_bytes(&root);
 
-        // Update twiddles for the next level.
+        // Update twiddles and offset for the next level.
         update_twiddles_in_place(&mut inv_twiddles);
+        terminal_offset = terminal_offset.square();
     }
 
-    // <<<< Receive challenge: 𝜁ₙ₋₁
-    let zeta = transcript.sample_field_element();
-
-    // Final fold.
-    fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles);
-
-    let last_value = evals
-        .first()
-        .expect("FRI evals are non-empty after folding")
-        .clone();
-
-    // >>>> Send value: pₙ
-    transcript.append_field_element(&last_value);
+    // 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 {
+        // <<<< 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");
+
+    // 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;
+    let final_poly_coeffs = crate::fri::terminal::coeffs_from_terminal_codeword::<F, E>(
+        &evals,
+        &terminal_offset,
+        effective_log_degree,
+    );
+    for c in &final_poly_coeffs {
+        transcript.append_field_element(c);
+    }
 
-    (last_value, fri_layer_list)
+    (final_poly_coeffs, fri_layer_list)
 }
 
 pub fn query_phase<F: IsField>(

crypto/stark/src/fri/terminal.rs

@@ -0,0 +1,109 @@
+//! 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.
+
+use math::fft::bit_reversing::in_place_bit_reverse_permute;
+use math::field::element::FieldElement;
+use math::field::traits::{IsFFTField, IsField, IsSubFieldOf};
+use math::polynomial::Polynomial;
+
+/// 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.
+///
+/// The codeword is a coset evaluation of a polynomial of degree less than
+/// `2^final_poly_log_degree` on the coset `terminal_offset·⟨ω⟩` of size
+/// `blowup·2^k`.
+///
+/// Algorithm:
+/// 1. Bit-reverse permute to convert from FRI order to natural (DFT) order.
+/// 2. Decimate: extract the size-`2^k` sub-coset
+///    `terminal_offset·⟨ω^blowup⟩` = every `blowup`-th natural-order point.
+/// 3. Coset iFFT on the small (`2^k`-point) sub-domain — a `blowup×`-smaller
+///    transform that recovers the `2^k` coefficients directly (no oversized
+///    transform and no wasteful truncation).
+pub(crate) fn coeffs_from_terminal_codeword<F, E>(
+    codeword_bitrev: &[FieldElement<E>],
+    terminal_offset: &FieldElement<F>,
+    final_poly_log_degree: u32,
+) -> Vec<FieldElement<E>>
+where
+    F: IsFFTField + IsSubFieldOf<E>,
+    E: IsField + Send + Sync,
+{
+    // Bit-reversed -> natural order.
+    let mut natural = codeword_bitrev.to_vec();
+    in_place_bit_reverse_permute(&mut natural);
+
+    // A degree-<2^k poly is determined by 2^k points: take the size-2^k sub-coset
+    // terminal_offset*<w^blowup> = every `blowup`-th natural-order evaluation.
+    let keep = 1usize << final_poly_log_degree;
+    let blowup = natural.len() / keep;
+    let sub_coset: Vec<FieldElement<E>> = natural.into_iter().step_by(blowup).collect();
+    debug_assert_eq!(sub_coset.len(), keep);
+
+    // Coset iFFT on the small domain -> the 2^k coefficients directly (no oversized trim).
+    let poly = Polynomial::interpolate_offset_fft::<F>(&sub_coset, terminal_offset)
+        .expect("terminal sub-coset must have power-of-two length and non-zero offset");
+
+    // Pad with zeros only if interpolation dropped trailing-zero coeffs, so the
+    // proof always carries exactly 2^k coefficients (the verifier length-checks).
+    let mut coeffs = poly.coefficients().to_vec();
+    coeffs.resize(keep, FieldElement::<E>::zero());
+    coeffs
+}
+
+/// Verifier side: given `2^k` coefficients of the low-degree polynomial,
+/// reconstruct the full FRI terminal codeword in **bit-reversed** order.
+///
+/// Algorithm:
+/// 1. FFT (coset): evaluate the polynomial on the full coset of size
+///    `codeword_len` with shift `terminal_offset` to get natural order.
+/// 2. Bit-reverse permute to convert natural order to FRI order.
+///
+/// # Panics
+///
+/// Panics if any of the following preconditions are violated:
+/// - `coeffs` is non-empty,
+/// - `coeffs.len()` is a power of two,
+/// - `codeword_len` is a power of two,
+/// - `coeffs.len() <= codeword_len`, and
+/// - `codeword_len` is divisible by `coeffs.len()`.
+///
+/// In the normal verifier flow these conditions are guaranteed by the
+/// final-polynomial length check that the verifier performs before calling
+/// this helper, so the assert should never fire in production.
+pub(crate) fn terminal_codeword_from_coeffs<F, E>(
+    coeffs: &[FieldElement<E>],
+    terminal_offset: &FieldElement<F>,
+    codeword_len: usize,
+) -> Vec<FieldElement<E>>
+where
+    F: IsFFTField + IsSubFieldOf<E>,
+    E: IsField + Send + Sync,
+{
+    assert!(
+        !coeffs.is_empty()
+            && coeffs.len().is_power_of_two()
+            && codeword_len.is_power_of_two()
+            && coeffs.len() <= codeword_len
+            && codeword_len.is_multiple_of(coeffs.len()),
+        "terminal_codeword_from_coeffs: coeffs.len() ({}) must be a non-zero power of two dividing codeword_len ({}); the verifier must length-check coeffs before calling",
+        coeffs.len(),
+        codeword_len,
+    );
+
+    let poly = Polynomial::new(coeffs);
+    let blowup = codeword_len / coeffs.len();
+
+    // Step 1: coset FFT to get natural-order evaluations.
+    let mut natural =
+        Polynomial::evaluate_offset_fft::<F>(&poly, blowup, Some(coeffs.len()), terminal_offset)
+            .expect("terminal coset size must be a power of two within the field's two-adicity");
+
+    // Step 2: convert natural order to bit-reversed (FRI) order.
+    in_place_bit_reverse_permute(&mut natural);
+    natural
+}