integrated decomposition in fri (#598)

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/starkware-libs/stwo/598)
<!-- Reviewable:end -->

crates/prover/src/core/fri.rs

@@ -21,6 +21,7 @@ use super::poly::twiddles::TwiddleTree;
 use super::poly::BitReversedOrder;
 // TODO(andrew): Create fri/ directory, move queries.rs there and split this file up.
 use super::queries::{Queries, SparseSubCircleDomain};
+use crate::core::air::accumulation::PointEvaluationAccumulator;
 use crate::core::circle::Coset;
 use crate::core::fft::ibutterfly;
 use crate::core::fields::FieldExpOps;
@@ -202,18 +203,27 @@ impl<B: FriOps + MerkleOps<H>, H: MerkleHasher> FriProver<B, H> {
         let circle_poly_alpha = channel.draw_felt();
 
         while layer_evaluation.len() > config.last_layer_domain_size() {
-            // Check for any columns (circle poly evaluations) that should be combined.
+            let mut layer_lambda_acc = PointEvaluationAccumulator::new(circle_poly_alpha.square());
+
+            // The following steps are blowup factor 2 specific.
+            // TODO(Ohad/ShaharS/ShaharP): support larger blowup.
+            assert_eq!(
+                config.log_blowup_factor, 1,
+                "only log blowup factor 2 is supported"
+            );
+
+            // Check for any columns (circle poly evaluations) that should be combined. If there are
+            // multiple, combine them into a single column.
             while let Some(column) = columns.next_if(|c| folded_len(c) == layer_evaluation.len()) {
-                B::fold_circle_into_line(
-                    &mut layer_evaluation,
-                    column,
-                    circle_poly_alpha,
-                    twiddles,
-                );
+                let (g, lambda) = B::decompose(column);
+                layer_lambda_acc.accumulate(lambda);
+                B::fold_circle_into_line(&mut layer_evaluation, &g, circle_poly_alpha, twiddles);
             }
 
-            let layer = FriLayerProver::new(layer_evaluation);
+            let layer_lambda = layer_lambda_acc.finalize();
+            let layer = FriLayerProver::new(layer_evaluation, layer_lambda);
             channel.mix_digest(layer.merkle_tree.root());
+            channel.mix_felts(&[layer_lambda]);
             let folding_alpha = channel.draw_felt();
             let folded_layer_evaluation = B::fold_line(&layer.evaluation, folding_alpha, twiddles);
 
@@ -354,6 +364,10 @@ impl<H: MerkleHasher> FriVerifier<H> {
         for (layer_index, proof) in proof.inner_layers.into_iter().enumerate() {
             channel.mix_digest(proof.commitment);
 
+            // The merkle verification, combined with the decomposition being unique, asserts the
+            // decomposition correctness.
+            channel.mix_felts(&[proof.decomposition_coeff]);
+
             let folding_alpha = channel.draw_felt();
 
             inner_layers.push(FriLayerVerifier {
@@ -445,17 +459,34 @@ impl<H: MerkleHasher> FriVerifier<H> {
         let mut layer_query_evals = vec![SecureField::zero(); layer_queries.len()];
 
         for layer in self.inner_layers.iter() {
-            // Check for column evals that need to folded into this layer.
-            while column_bounds
+            if column_bounds
                 .next_if(|b| b.fold_to_line() == layer.degree_bound)
                 .is_some()
             {
-                let sparse_evaluation = decommitted_values.next().unwrap();
-                let folded_evals = sparse_evaluation.fold(circle_poly_alpha);
-                assert_eq!(folded_evals.len(), layer_query_evals.len());
+                let mut n_columns_in_layer = 1;
+                let mut combined_sparse_evals = decommitted_values.next().unwrap();
+
+                // Check for more columns that were folded into the current layer.
+                while column_bounds
+                    .next_if(|b| b.fold_to_line() == layer.degree_bound)
+                    .is_some()
+                {
+                    combined_sparse_evals
+                        .accumulate(decommitted_values.next().unwrap(), circle_poly_alpha_sq);
+                    n_columns_in_layer += 1;
+                }
+
+                Self::project_to_fft_space(
+                    &layer_queries,
+                    &mut combined_sparse_evals,
+                    layer.proof.decomposition_coeff,
+                );
 
+                let folded_evals = combined_sparse_evals.fold(circle_poly_alpha);
+                let prev_layer_combination_factor = circle_poly_alpha_sq.pow(n_columns_in_layer);
+                assert_eq!(folded_evals.len(), layer_query_evals.len());
                 for (layer_eval, folded_eval) in zip(&mut layer_query_evals, folded_evals) {
-                    *layer_eval = *layer_eval * circle_poly_alpha_sq + folded_eval;
+                    *layer_eval = *layer_eval * prev_layer_combination_factor + folded_eval;
                 }
             }
 
@@ -470,6 +501,27 @@ impl<H: MerkleHasher> FriVerifier<H> {
         Ok((layer_queries, layer_query_evals))
     }
 
+    /// Projects a polynomial to the fft space using the coefficient from the proof.
+    /// See [`CPUBackend::decompose`] for explanation about the decomposition.
+    /// NOTE: specific to when Blowup factor is 2.
+    // TODO(Ohad): generalize this.
+    fn project_to_fft_space(
+        layer_queries: &Queries,
+        evals: &mut SparseCircleEvaluation,
+        lambda: SecureField,
+    ) {
+        let domain_size = 1 << layer_queries.log_domain_size;
+        layer_queries
+            .iter()
+            .zip(evals.subcircle_evals.iter_mut())
+            .for_each(|(&q, e)| {
+                let lambda = if q < domain_size / 2 { lambda } else { -lambda };
+                for v in e.values.iter_mut() {
+                    *v -= lambda;
+                }
+            });
+    }
+
     /// Verifies the last layer.
     fn decommit_last_layer(
         self,
@@ -639,6 +691,7 @@ pub struct FriLayerProof<H: MerkleHasher> {
     /// to fold and verify the merkle decommitment.
     pub evals_subset: Vec<SecureField>,
     pub decommitment: MerkleDecommitment<H>,
+    pub decomposition_coeff: SecureField,
     pub commitment: H::Hash,
 }
 
@@ -789,17 +842,19 @@ impl<H: MerkleHasher> FriLayerVerifier<H> {
 // TODO(andrew): Support different step sizes.
 struct FriLayerProver<B: FriOps + MerkleOps<H>, H: MerkleHasher> {
     evaluation: LineEvaluation<B>,
+    decomposition_coeff: SecureField,
     merkle_tree: MerkleProver<B, H>,
 }
 
 impl<B: FriOps + MerkleOps<H>, H: MerkleHasher> FriLayerProver<B, H> {
-    fn new(evaluation: LineEvaluation<B>) -> Self {
+    fn new(evaluation: LineEvaluation<B>, decomposition_coeff: SecureField) -> Self {
         // TODO(spapini): Commit on slice.
         // TODO(spapini): Merkle tree in backend.
         let merkle_tree = MerkleProver::commit(evaluation.values.columns.iter().collect_vec());
         #[allow(unreachable_code)]
         FriLayerProver {
             evaluation,
+            decomposition_coeff,
             merkle_tree,
         }
     }
@@ -839,10 +894,12 @@ impl<B: FriOps + MerkleOps<H>, H: MerkleHasher> FriLayerProver<B, H> {
                 .collect(),
             self.evaluation.values.columns.iter().collect_vec(),
         );
+        let decomposition_coeff = self.decomposition_coeff;
 
         FriLayerProof {
             evals_subset,
             decommitment,
+            decomposition_coeff,
             commitment,
         }
     }
@@ -884,6 +941,17 @@ impl SparseCircleEvaluation {
             })
             .collect()
     }
+
+    /// Computes `self = self * alpha + rhs`.
+    fn accumulate(&mut self, rhs: Self, alpha: SecureField) {
+        assert_eq!(self.subcircle_evals.len(), rhs.subcircle_evals.len());
+        for (lhs, rhs) in self.subcircle_evals.iter_mut().zip(rhs.subcircle_evals) {
+            assert_eq!(lhs.len(), rhs.len());
+            for (lhs, rhs) in lhs.values.iter_mut().zip(rhs.values) {
+                *lhs = *lhs * alpha + rhs;
+            }
+        }
+    }
 }
 
 impl<'a> IntoIterator for &'a mut SparseCircleEvaluation {
@@ -1001,16 +1069,17 @@ mod tests {
         fold_circle_into_line, fold_line, CirclePolyDegreeBound, FriConfig, FriVerifier,
         CIRCLE_TO_LINE_FOLD_STEP,
     };
-    use crate::core::poly::circle::{CircleDomain, PolyOps, SecureEvaluation};
+    use crate::core::poly::circle::{CanonicCoset, CircleDomain, PolyOps, SecureEvaluation};
     use crate::core::poly::line::{LineDomain, LineEvaluation, LinePoly};
     use crate::core::poly::NaturalOrder;
     use crate::core::queries::{Queries, SparseSubCircleDomain};
     use crate::core::test_utils::test_channel;
     use crate::core::utils::bit_reverse_index;
     use crate::core::vcs::blake2_merkle::Blake2sMerkleHasher;
+    use crate::m31;
 
     /// Default blowup factor used for tests.
-    const LOG_BLOWUP_FACTOR: u32 = 2;
+    const LOG_BLOWUP_FACTOR: u32 = 1;
 
     type FriProver = super::FriProver<CPUBackend, Blake2sMerkleHasher>;
 
@@ -1379,6 +1448,56 @@ mod tests {
         let _ = verifier.decommit_on_queries(&invalid_queries, vec![decommitment_value]);
     }
 
+    #[test]
+    fn low_degree_polynomials_outside_fft_space_pass_verification(
+    ) -> Result<(), FriVerificationError> {
+        const LOG_DEGREE: u32 = 9;
+        const LOG_DOMAIN_SIZE: u32 = LOG_DEGREE + 1;
+        let evaluation = riemann_roch_polynomial_evaluation(LOG_DOMAIN_SIZE);
+        let queries = Queries::from_positions(vec![2, 10, 11], LOG_DOMAIN_SIZE);
+        let config = FriConfig::new(1, 1, queries.len());
+        let decommitment_value = query_polynomial(&evaluation, &queries);
+        let prover = FriProver::commit(
+            &mut test_channel(),
+            config,
+            &[evaluation.clone(), evaluation.clone()],
+            &CPUBackend::precompute_twiddles(evaluation.domain.half_coset),
+        );
+        let proof = prover.decommit_on_queries(&queries);
+        let bound = vec![
+            CirclePolyDegreeBound::new(LOG_DEGREE),
+            CirclePolyDegreeBound::new(LOG_DEGREE),
+        ];
+        let verifier = FriVerifier::commit(&mut test_channel(), config, proof, bound).unwrap();
+
+        verifier.decommit_on_queries(
+            &queries,
+            vec![decommitment_value.clone(), decommitment_value],
+        )
+    }
+
+    #[test]
+    fn low_mixed_degree_polynomials_outside_fft_space_pass_verification(
+    ) -> Result<(), FriVerificationError> {
+        const LOG_DEGREES: [u32; 6] = [8, 5, 5, 5, 5, 4];
+        let evaluations = LOG_DEGREES.map(|log_d| riemann_roch_polynomial_evaluation(log_d + 1));
+        let log_domain_size = evaluations[0].domain.log_size();
+        let queries = Queries::from_positions(vec![3, 250, 500], log_domain_size);
+        let config = FriConfig::new(2, 1, queries.len());
+        let prover = FriProver::commit(
+            &mut test_channel(),
+            config,
+            &evaluations,
+            &CPUBackend::precompute_twiddles(evaluations[0].domain.half_coset),
+        );
+        let decommitment_values = evaluations.map(|p| query_polynomial(&p, &queries)).to_vec();
+        let proof = prover.decommit_on_queries(&queries);
+        let bounds = LOG_DEGREES.map(CirclePolyDegreeBound::new).to_vec();
+        let verifier = FriVerifier::commit(&mut test_channel(), config, proof, bounds).unwrap();
+
+        verifier.decommit_on_queries(&queries, decommitment_values)
+    }
+
     /// Returns an evaluation of a random polynomial with degree `2^log_degree`.
     ///
     /// The evaluation domain size is `2^(log_degree + log_blowup_factor)`.
@@ -1446,4 +1565,29 @@ mod tests {
 
         SparseCircleEvaluation::new(coset_evals)
     }
+
+    /// Returns an evaluation of a polynomial that is out of FFT-space but inside
+    /// riemann-roch-space.
+    fn riemann_roch_polynomial_evaluation(log_domain_size: u32) -> SecureEvaluation<CPUBackend> {
+        let mut coeffs = vec![BaseField::zero(); 1 << log_domain_size];
+        let domain_log_half_size = log_domain_size - 1;
+        // Polynomial is out of FFT space.
+        coeffs[1 << domain_log_half_size] = m31!(1);
+        coeffs[(1 << domain_log_half_size) - 1] = m31!(2);
+        let poly = CPUCirclePoly::new(coeffs);
+
+        let domain = CanonicCoset::new(log_domain_size).circle_domain();
+        let values = poly.evaluate(domain);
+        SecureEvaluation {
+            domain,
+            values: SecureColumn {
+                columns: [
+                    values.values,
+                    Col::<CPUBackend, BaseField>::zeros(1 << log_domain_size),
+                    Col::<CPUBackend, BaseField>::zeros(1 << log_domain_size),
+                    Col::<CPUBackend, BaseField>::zeros(1 << log_domain_size),
+                ],
+            },
+        }
+    }
 }

crates/prover/src/examples/fibonacci/mod.rs

@@ -227,7 +227,7 @@ mod tests {
         let fib = Fibonacci::new(FIB_LOG_SIZE, m31!(443693538));
 
         let mut invalid_proof = fib.prove().unwrap();
-        invalid_proof.commitment_scheme_proof.queried_values.0[0][0][4] += BaseField::one();
+        invalid_proof.commitment_scheme_proof.queried_values.0[0][0][3] += BaseField::one();
 
         let error = fib.verify(invalid_proof).unwrap_err();
         assert_matches!(error, VerificationError::Merkle(_));