Feat/higher arity fri

crypto/stark/benches/profile_prover.rs

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

crypto/stark/benches/prover_benchmark.rs

@@ -61,6 +61,8 @@ fn benchmark_proof_options() -> ProofOptions {
         fri_number_of_queries: 30,
         coset_offset: 3,
         grinding_factor: 0,
+        fri_log_arity: 1,
+        fri_log_final_poly_len: 0,
     }
 }
 

crypto/stark/src/config.rs

@@ -1,5 +1,5 @@
 use crypto::merkle_tree::{
-    backends::types::{BatchKeccak256Backend, Keccak256Backend, PairKeccak256Backend},
+    backends::types::{BatchKeccak256Backend, Keccak256Backend},
     merkle::MerkleTree,
 };
 
@@ -19,6 +19,6 @@ pub type Commitment = [u8; COMMITMENT_SIZE];
 pub type BatchedMerkleTreeBackend<F> = BatchKeccak256Backend<F>;
 pub type BatchedMerkleTree<F> = MerkleTree<BatchedMerkleTreeBackend<F>>;
 
-// FRI layer uses fixed-size pairs for efficiency (avoids Vec allocation per pair)
-pub type FriLayerMerkleTreeBackend<F> = PairKeccak256Backend<F>;
+// FRI layer uses batch backend to support variable-arity FRI (generalizes over leaf width)
+pub type FriLayerMerkleTreeBackend<F> = BatchKeccak256Backend<F>;
 pub type FriLayerMerkleTree<F> = MerkleTree<FriLayerMerkleTreeBackend<F>>;

crypto/stark/src/fri/fri_decommit.rs

@@ -8,5 +8,8 @@ use crate::config::Commitment;
 #[serde(bound = "")]
 pub struct FriDecommitment<F: IsField> {
     pub layers_auth_paths: Vec<Proof<Commitment>>,
-    pub layers_evaluations_sym: Vec<FieldElement<F>>,
+    /// For each FRI layer, the sibling evaluations at the queried position.
+    /// Vec<Vec<FE>> supports higher-arity FRI where arity-k folds k values per layer.
+    /// For arity-2 (current default) each inner Vec has exactly one element.
+    pub layers_evaluations_sym: Vec<Vec<FieldElement<F>>>,
 }

crypto/stark/src/fri/mod.rs

@@ -17,83 +17,145 @@ 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).
+///
+/// Supports configurable folding arity (`log_arity` binary folds per Merkle commit)
+/// and early termination (`log_final_poly_len` remaining evaluations sent as coefficients).
+///
+/// Protocol structure (matches the verifier's expectations):
+/// - Round 0: receive challenge ζ₀, fold **1×** (always arity-2 initial fold), commit p₁.
+///   This corresponds to the verifier's "initial fold: always arity-2 from DEEP composition".
+/// - Rounds 1+: receive challenge ζₖ, fold **log_arity×**, commit pₖ.
+///   The verifier folds each committed layer's group by log_arity sub-folds using ζₖ.
+/// - Last round: fold log_arity×, **no commit** — send final polynomial coefficients.
+///
+/// `total_folds` is adjusted so that `(total_folds - 1)` is divisible by `log_arity`,
+/// ensuring every non-initial round does exactly `log_arity` folds.
+/// The final polynomial may have > 1 element when `log_arity > 1`.
+///
+/// For `log_arity=1` this reproduces the standard arity-2 FRI exactly.
+#[allow(clippy::type_complexity)]
 pub fn commit_phase_from_evaluations<F: IsFFTField + IsSubFieldOf<E>, E: IsField>(
     number_layers: usize,
     mut evals: Vec<FieldElement<E>>,
     transcript: &mut impl IsStarkTranscript<E, F>,
     coset_offset: &FieldElement<F>,
     domain_size: usize,
+    log_arity: usize,
+    log_final_poly_len: usize,
 ) -> (
-    FieldElement<E>,
+    Vec<FieldElement<E>>,
     Vec<FriLayer<E, FriLayerMerkleTreeBackend<E>>>,
 )
 where
     FieldElement<F>: AsBytes + Sync + Send,
     FieldElement<E>: AsBytes + Sync + Send,
 {
-    // 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(number_layers);
+    // Maximum binary folds available given the domain and desired final poly length.
+    let max_folds = number_layers.max(1).saturating_sub(log_final_poly_len);
+
+    // Adjust total_folds so that (total_folds - 1) is divisible by log_arity.
+    // This ensures every non-initial round does exactly log_arity folds, matching
+    // the verifier which always folds log_arity times per committed layer.
+    // For log_arity=1 this is a no-op (any value satisfies (T-1) % 1 == 0).
+    let total_folds = if max_folds == 0 {
+        0
+    } else {
+        // Largest T ≤ max_folds with T = 1 + n*log_arity for integer n ≥ 0.
+        let n = (max_folds - 1) / log_arity;
+        1 + n * log_arity
+    };
+
+    // Number of committed FRI layers: all rounds except the last (final poly round).
+    let num_committed_rounds = if total_folds > 0 {
+        // Round 0 always commits (1 fold). Rounds 1..k commit every log_arity folds.
+        // Total committed = 1 + (total_folds - 1) / log_arity - 1 = (total_folds - 1) / log_arity.
+        (total_folds - 1) / log_arity
+    } else {
+        0
+    };
+
+    // Leaf grouping for Merkle trees: verifier always folds log_arity sub-folds per layer.
+    let group_size = 1usize << log_arity;
+
+    let mut fri_layer_list = Vec::with_capacity(num_committed_rounds);
     let mut current_coset_offset = coset_offset.clone();
     let mut current_domain_size = domain_size;
+    let mut folds_done = 0;
+    let mut is_first_round = true;
 
-    for _ in 1..number_layers {
-        // <<<< Receive challenge 𝜁ₖ₋₁
-        let zeta = transcript.sample_field_element();
-        current_coset_offset = current_coset_offset.square();
-        current_domain_size /= 2;
-
-        // Fold evaluations in-place (no FFT needed)
-        fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles);
-
-        // Build Merkle tree from consecutive pairs
-        let leaves: Vec<[FieldElement<E>; 2]> = evals
-            .chunks_exact(2)
-            .map(|chunk| [chunk[0].clone(), chunk[1].clone()])
-            .collect();
-        let merkle_tree = FriLayerMerkleTree::build(&leaves)
-            .expect("FRI commit: Merkle tree construction must succeed");
-        let root = merkle_tree.root;
-        fri_layer_list.push(FriLayer::new(
-            &evals,
-            merkle_tree,
-            current_coset_offset.clone().to_extension(),
-            current_domain_size,
-        ));
-
-        // >>>> Send commitment: [pₖ]
-        transcript.append_bytes(&root);
-
-        // Update twiddles for next level
-        update_twiddles_in_place(&mut inv_twiddles);
-    }
+    while folds_done < total_folds {
+        // Round 0 always folds exactly 1 time (the "initial arity-2 fold").
+        // All subsequent rounds fold exactly log_arity times.
+        let folds_this_round = if is_first_round { 1 } else { log_arity };
+        is_first_round = false;
+        let is_last_round = folds_done + folds_this_round >= total_folds;
 
-    // <<<< Receive challenge: 𝜁ₙ₋₁
-    let zeta = transcript.sample_field_element();
-
-    // Final fold
-    fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles);
+        // <<<< Receive challenge: one per committed round
+        let zeta = transcript.sample_field_element();
 
-    let last_value = evals.first().unwrap_or(&FieldElement::zero()).clone();
+        // Apply folds_this_round sequential binary folds.
+        // Challenges: zeta, zeta^2, zeta^4, ..., zeta^(2^(k-1))
+        let mut challenge = zeta;
+        for _ in 0..folds_this_round {
+            fold_evaluations_in_place(&mut evals, &challenge, &inv_twiddles);
+            current_coset_offset = current_coset_offset.square();
+            current_domain_size /= 2;
+            update_twiddles_in_place(&mut inv_twiddles);
+            challenge = challenge.square();
+        }
+        folds_done += folds_this_round;
+
+        // Commit post-fold evaluations (skip for last round — final poly sent separately).
+        // Leaves always group `group_size = 2^log_arity` elements so the verifier
+        // can open and fold an entire arity-k coset at each query.
+        if !is_last_round {
+            let leaves: Vec<Vec<FieldElement<E>>> = evals
+                .chunks_exact(group_size)
+                .map(|chunk| chunk.to_vec())
+                .collect();
+            let merkle_tree = FriLayerMerkleTree::build(&leaves)
+                .expect("FRI commit: Merkle tree construction must succeed");
+            let root = merkle_tree.root;
+            fri_layer_list.push(FriLayer::new(
+                &evals,
+                merkle_tree,
+                current_coset_offset.clone().to_extension(),
+                current_domain_size,
+            ));
+
+            // >>>> Send commitment: [pₖ]
+            transcript.append_bytes(&root);
+        }
+    }
 
-    // >>>> Send value: pₙ
-    transcript.append_field_element(&last_value);
+    // >>>> Send final polynomial coefficients
+    let fri_final_poly = evals;
+    for coeff in &fri_final_poly {
+        transcript.append_field_element(coeff);
+    }
 
-    (last_value, fri_layer_list)
+    (fri_final_poly, fri_layer_list)
 }
 
+/// FRI query phase: for each query index, collect siblings and auth paths from each layer.
+///
+/// `log_arity` controls how many elements are in each leaf group.
+/// For arity-2 (log_arity=1): each layer has 1 sibling.
+/// For arity-4 (log_arity=2): each layer has 3 siblings.
 pub fn query_phase<F: IsField>(
-    fri_layers: &Vec<FriLayer<F, FriLayerMerkleTreeBackend<F>>>,
+    fri_layers: &[FriLayer<F, FriLayerMerkleTreeBackend<F>>],
     iotas: &[usize],
+    log_arity: usize,
 ) -> Vec<FriDecommitment<F>>
 where
     FieldElement<F>: AsBytes + Sync + Send,
 {
     if !fri_layers.is_empty() {
         let num_layers = fri_layers.len();
+        let group_size = 1usize << log_arity;
+
         iotas
             .iter()
             .map(|iota_s| {
@@ -102,13 +164,19 @@ where
 
                 let mut index = *iota_s;
                 for layer in fri_layers {
-                    // symmetric element
-                    let evaluation_sym = layer.evaluation[index ^ 1].clone();
-                    let auth_path_sym = layer.merkle_tree.get_proof_by_pos(index >> 1).unwrap();
-                    layers_evaluations_sym.push(evaluation_sym);
-                    layers_auth_paths_sym.push(auth_path_sym);
-
-                    index >>= 1;
+                    // Collect all siblings in the group (K-1 elements, excluding the queried one)
+                    let group_start = (index / group_size) * group_size;
+                    let siblings: Vec<FieldElement<F>> = (0..group_size)
+                        .filter(|&k| group_start + k != index)
+                        .map(|k| layer.evaluation[group_start + k].clone())
+                        .collect();
+
+                    let leaf_index = index / group_size;
+                    let auth_path = layer.merkle_tree.get_proof_by_pos(leaf_index).unwrap();
+                    layers_evaluations_sym.push(siblings);
+                    layers_auth_paths_sym.push(auth_path);
+
+                    index >>= log_arity;
                 }
 
                 FriDecommitment {
@@ -118,9 +186,6 @@ where
             })
             .collect()
     } else {
-        // For 0 FRI layers (small traces), return empty decommitments for each query.
-        // The verifier still needs one decommitment entry per query, even if the
-        // FRI layer data is empty.
         iotas
             .iter()
             .map(|_| FriDecommitment {

crypto/stark/src/proof/options.rs

@@ -45,6 +45,13 @@ pub struct ProofOptions {
     pub fri_number_of_queries: usize,
     pub coset_offset: u64,
     pub grinding_factor: u8,
+    /// Log2 of FRI folding arity per round. 1 = arity-2 (current default),
+    /// 2 = arity-4, 3 = arity-8. Each round applies this many sequential
+    /// binary folds before committing a Merkle tree.
+    pub fri_log_arity: u8,
+    /// Log2 of final polynomial length. 0 = fold to constant (current default).
+    /// FRI stops when the polynomial degree drops below 2^fri_log_final_poly_len.
+    pub fri_log_final_poly_len: u8,
 }
 
 impl ProofOptions {
@@ -56,6 +63,8 @@ impl ProofOptions {
             fri_number_of_queries: 3,
             coset_offset: 3,
             grinding_factor: 1,
+            fri_log_arity: 1,
+            fri_log_final_poly_len: 0,
         }
     }
 }
@@ -112,6 +121,8 @@ impl GoldilocksCubicProofOptions {
             fri_number_of_queries,
             coset_offset: 3,
             grinding_factor,
+            fri_log_arity: 2, // arity-4: halves FRI Merkle tree count
+            fri_log_final_poly_len: 0,
         })
     }
 }

crypto/stark/src/proof/stark.rs

@@ -52,8 +52,8 @@ pub struct StarkProof<F: IsSubFieldOf<E>, E: IsField, PI> {
     pub composition_poly_parts_ood_evaluation: Vec<FieldElement<E>>,
     // [pₖ]
     pub fri_layers_merkle_roots: Vec<Commitment>,
-    // pₙ
-    pub fri_last_value: FieldElement<E>,
+    // pₙ (vector to support higher-arity FRI where the final poly has degree > 0)
+    pub fri_final_poly: Vec<FieldElement<E>>,
     // Open(pₖ(Dₖ), −𝜐ₛ^(2ᵏ))
     pub query_list: Vec<FriDecommitment<E>>,
     // Open(H₁(D_LDE, 𝜐ᵢ), Open(H₂(D_LDE, 𝜐ᵢ), Open(tⱼ(D_LDE), 𝜐ᵢ)

crypto/stark/src/prover.rs

@@ -311,8 +311,9 @@ pub struct Round3<F: IsField> {
 
 /// A container for the results of the fourth round of the STARK Prove protocol.
 pub struct Round4<F: IsSubFieldOf<E>, E: IsField> {
-    /// The final value resulting from folding the Deep composition polynomial all the way down to a constant value.
-    fri_last_value: FieldElement<E>,
+    /// The final polynomial resulting from folding the Deep composition polynomial all the way down.
+    /// For arity-2 FRI this is a single constant; for higher arity it may have degree > 0.
+    fri_final_poly: Vec<FieldElement<E>>,
     /// The commitments to the fold polynomials of the inner layers of FRI.
     fri_layers_merkle_roots: Vec<Commitment>,
     /// The values and proofs of validity of the evaluations of the trace polynomials and the composition polynomials
@@ -1090,13 +1091,17 @@ pub trait IsStarkProver<
         // FRI commit phase from pre-computed evaluations (no initial FFT)
         #[cfg(feature = "instruments")]
         let t_sub = Instant::now();
-        let (fri_last_value, fri_layers) =
+        let log_arity = air.context().proof_options.fri_log_arity as usize;
+        let log_final_poly_len = air.context().proof_options.fri_log_final_poly_len as usize;
+        let (fri_final_poly, fri_layers) =
             fri::commit_phase_from_evaluations::<Field, FieldExtension>(
                 domain.root_order as usize,
                 lde_evals,
                 transcript,
                 &coset_offset,
                 domain_size,
+                log_arity,
+                log_final_poly_len,
             );
         #[cfg(feature = "instruments")]
         let r4_merkle_dur = t_sub.elapsed();
@@ -1116,7 +1121,7 @@ pub trait IsStarkProver<
         let number_of_queries = air.options().fri_number_of_queries;
         let iotas = Self::sample_query_indexes(number_of_queries, domain, transcript);
 
-        let query_list = fri::query_phase(&fri_layers, &iotas);
+        let query_list = fri::query_phase(&fri_layers, &iotas, log_arity);
 
         let fri_layers_merkle_roots: Vec<_> = fri_layers
             .iter()
@@ -1133,7 +1138,7 @@ pub trait IsStarkProver<
         }
 
         Round4 {
-            fri_last_value,
+            fri_final_poly,
             fri_layers_merkle_roots,
             deep_poly_openings,
             query_list,
@@ -2040,7 +2045,7 @@ pub trait IsStarkProver<
             // [pₖ]
             fri_layers_merkle_roots: round_4_result.fri_layers_merkle_roots,
             // pₙ
-            fri_last_value: round_4_result.fri_last_value,
+            fri_final_poly: round_4_result.fri_final_poly,
             // Open(p₀(D₀), 𝜐ₛ), Open(pₖ(Dₖ), −𝜐ₛ^(2ᵏ))
             query_list: round_4_result.query_list,
             // Open(H₁(D_LDE, 𝜐₀), Open(H₂(D_LDE, 𝜐₀), Open(tⱼ(D_LDE), 𝜐₀)