Fix reduction strategy

plonky2/src/fri/mod.rs

@@ -67,7 +67,7 @@ pub struct FriParams {
 }
 
 impl FriParams {
-    pub(crate) fn total_arities(&self) -> usize {
+    pub fn total_arities(&self) -> usize {
         self.reduction_arity_bits.iter().sum()
     }
 

plonky2/src/fri/reduction_strategies.rs

@@ -8,11 +8,12 @@ pub enum FriReductionStrategy {
     /// Specifies the exact sequence of arities (expressed in bits) to use.
     Fixed(Vec<usize>),
 
-    /// `ConstantArityBits(arity_bits, final_poly_bits)` applies reductions of arity `2^arity_bits`
-    /// until the polynomial degree is `2^final_poly_bits` or less. This tends to work well in the
-    /// recursive setting, as it avoids needing multiple configurations of gates used in FRI
-    /// verification, such as `InterpolationGate`.
-    ConstantArityBits(usize, usize),
+    /// `ConstantArityBits(arity_bits, final_poly_bits, cap_height)` applies reductions of arity `2^arity_bits`
+    /// until the polynomial degree is less than or equal to `2^final_poly_bits` or until any further
+    /// `arity_bits`-reduction makes the polynomial degree smaller than `2^cap_height` (which would make FRI fail).
+    /// This tends to work well in the recursive setting, as it avoids needing multiple configurations
+    /// of gates used in FRI verification, such as `InterpolationGate`.
+    ConstantArityBits(usize, usize, usize),
 
     /// `MinSize(opt_max_arity_bits)` searches for an optimal sequence of reduction arities, with an
     /// optional max `arity_bits`. If this proof will have recursive proofs on top of it, a max
@@ -31,12 +32,12 @@ impl FriReductionStrategy {
         match self {
             FriReductionStrategy::Fixed(reduction_arity_bits) => reduction_arity_bits.to_vec(),
 
-            FriReductionStrategy::ConstantArityBits(arity_bits, final_poly_bits) => {
+            &FriReductionStrategy::ConstantArityBits(arity_bits, final_poly_bits, cap_height) => {
                 let mut result = Vec::new();
-                while degree_bits > *final_poly_bits {
-                    result.push(*arity_bits);
-                    assert!(degree_bits >= *arity_bits);
-                    degree_bits -= *arity_bits;
+                while degree_bits > final_poly_bits && degree_bits - arity_bits >= cap_height {
+                    result.push(arity_bits);
+                    assert!(degree_bits >= arity_bits);
+                    degree_bits -= arity_bits;
                 }
                 result.shrink_to_fit();
                 result

plonky2/src/plonk/circuit_builder.rs

@@ -664,7 +664,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let degree_bits = log2_strict(degree);
         let fri_params = self.fri_params(degree_bits);
         assert!(
-            fri_params.total_arities() <= degree_bits,
+            fri_params.total_arities() <= degree_bits - self.config.fri_config.cap_height,
             "FRI total reduction arity is too large.",
         );
 

plonky2/src/plonk/circuit_data.rs

@@ -73,7 +73,7 @@ impl CircuitConfig {
                 rate_bits: 3,
                 cap_height: 4,
                 proof_of_work_bits: 16,
-                reduction_strategy: FriReductionStrategy::ConstantArityBits(4, 5),
+                reduction_strategy: FriReductionStrategy::ConstantArityBits(4, 5, 4),
                 num_query_rounds: 28,
             },
         }

starky/src/config.rs

@@ -22,7 +22,7 @@ impl StarkConfig {
                 rate_bits: 1,
                 cap_height: 4,
                 proof_of_work_bits: 10,
-                reduction_strategy: FriReductionStrategy::ConstantArityBits(4, 5),
+                reduction_strategy: FriReductionStrategy::ConstantArityBits(4, 5, 4),
                 num_query_rounds: 90,
             },
         }

starky/src/prover.rs

@@ -37,6 +37,11 @@ where
 {
     let degree = trace.len();
     let degree_bits = log2_strict(degree);
+    let fri_params = config.fri_params(degree_bits);
+    assert!(
+        fri_params.total_arities() <= degree_bits - config.fri_config.cap_height,
+        "FRI total reduction arity is too large.",
+    );
 
     let trace_vecs = trace.into_iter().map(|row| row.to_vec()).collect_vec();
     let trace_col_major: Vec<Vec<F>> = transpose(&trace_vecs);
@@ -117,7 +122,6 @@ where
 
     // TODO: Add permutation checks
     let initial_merkle_trees = &[&trace_commitment, &quotient_commitment];
-    let fri_params = config.fri_params(degree_bits);
 
     let opening_proof = timed!(
         timing,