Refactor compute_deep_composition_poly function (#200)

* [WIP] Deep FRI

* Finish implementation

* Document code

* Fix vulnerability tests and add more

* Define basic traits and structs for Air

* Write basic skeleton of ConstraintEvaluator

* Save evaluator progress

* Complete evaluator basic methods

* Fix some compilation errors

* Finish method for computing composition poly from Evaluation Table

* More fixes

* More progress

* Comment out tests

* [WIP] Implement verifier

* Debugging stuff

* Remove debugging stuff as the boundary ood evaluation was correct

* Uncomment important line

* Start compute_deep_composition_poly refactor

* Finish compute_deep_composition_poly refactor

* Refactor trace function to take the length as a parameter

* Fix hardcoded boundary poly degree on verifier

* Some fixes

* Fix transition_divisors implementation

* Fix transition_divisor once more

* Fix calculation of trace primitive root on the verifier

* Test finally working

* Remove unusued num_assertions field

* [WIP] fiat-shamir integration

* Comment transcript.append line in fri function

* Fix fiat-shamir integration bugs

* Fix boundaryconstraints zerofier test

* Fix clippy suggestiojns

* Remove outdated comments

* Rename compute_transition_evaluations

* Add Fibonacci AIR test for a 17 finite field

* Save refactor WIP

* Refactor compute_deep_composition_poly to generalize for a trace with more than one column

* Add comments to dummy transcript challenge calls in verifier

* Use a better variable name

* Remove debug comment

* Add comments and documentation

* Fix typo

---------

Co-authored-by: Javier Chatruc <[email protected]>

Author: entropidelic

proving_system/stark/src/air/context.rs

@@ -2,7 +2,7 @@
 pub struct AirContext {
     pub options: ProofOptions,
     pub trace_length: usize,
-    pub trace_info: (usize, usize),
+    pub trace_columns: usize,
     pub transition_degrees: Vec<usize>,
     /// This is a vector with the indices of all the rows that constitute
     /// an evaluation frame. Note that, because of how we write all constraints

proving_system/stark/src/air/frame.rs

@@ -52,6 +52,31 @@ impl<F: IsField> Frame<F> {
         Self::new(data, 1)
     }
 
+    /// Given a slice of trace polynomials, an evaluation point `x`, the frame offsets
+    /// corresponding to the computation of the transitions, and a primitive root,
+    /// outputs the trace evaluations of each trace polynomial over the values used to
+    /// compute a transition.
+    /// Example: For a simple Fibonacci computation, if t(x) is the trace polynomial of
+    /// the computation, this will output evaluations t(x), t(g * x), t(g^2 * z).
+    pub fn get_trace_evaluations(
+        trace_polys: &[Polynomial<FieldElement<F>>],
+        x: &FieldElement<F>,
+        frame_offsets: &[usize],
+        primitive_root: &FieldElement<F>,
+    ) -> Vec<Vec<FieldElement<F>>> {
+        let mut evaluations = Vec::with_capacity(trace_polys.len());
+        let evaluation_points: Vec<FieldElement<F>> = frame_offsets
+            .iter()
+            .map(|offset| x * primitive_root.pow(*offset))
+            .collect();
+
+        trace_polys
+            .iter()
+            .for_each(|p| evaluations.push(p.evaluate_slice(&evaluation_points)));
+
+        evaluations
+    }
+
     /// Returns the Out of Domain Frame for the given trace polynomials, out of domain evaluation point (called `z` in the literature),
     /// frame offsets given by the AIR and primitive root used for interpolating the trace polynomials.
     /// An out of domain frame is nothing more than the evaluation of the trace polynomials in the points required by the
@@ -65,15 +90,7 @@ impl<F: IsField> Frame<F> {
         frame_offsets: &[usize],
         primitive_root: &FieldElement<F>,
     ) -> Self {
-        let mut data = vec![];
-        let evaluation_points: Vec<FieldElement<F>> = frame_offsets
-            .iter()
-            .map(|offset| z * primitive_root.pow(*offset))
-            .collect();
-
-        for poly in trace_polys {
-            data.push(poly.evaluate_slice(&evaluation_points));
-        }
+        let data = Self::get_trace_evaluations(trace_polys, z, frame_offsets, primitive_root);
 
         Self {
             data: data.into_iter().flatten().collect(),

proving_system/stark/src/air/mod.rs

@@ -28,4 +28,7 @@ pub trait AIR: Clone {
     fn blowup_factor(&self) -> u8 {
         self.options().blowup_factor
     }
+    fn num_transition_constraints(&self) -> usize {
+        self.context().num_transition_constraints
+    }
 }

proving_system/stark/src/lib.rs

@@ -85,6 +85,10 @@ mod tests {
     #[test]
     fn test_prove_fib() {
         let trace = fibonacci_trace([FE::from(1), FE::from(1)], 4);
+        let trace_table = TraceTable {
+            table: trace.clone(),
+            num_cols: 1,
+        };
 
         let context = AirContext {
             options: ProofOptions {
@@ -93,18 +97,13 @@ mod tests {
                 coset_offset: 3,
             },
             trace_length: trace.len(),
-            trace_info: (trace.len(), 1),
+            trace_columns: trace_table.num_cols,
             transition_degrees: vec![1],
             transition_exemptions: vec![trace.len() - 2, trace.len() - 1],
             transition_offsets: vec![0, 1, 2],
             num_transition_constraints: 1,
         };
 
-        let trace_table = TraceTable {
-            table: trace.clone(),
-            num_cols: 1,
-        };
-
         let fibonacci_air = FibonacciAIR::new(trace_table, context);
 
         let result = prove(&trace, &fibonacci_air);
@@ -116,25 +115,25 @@ mod tests {
     fn test_prove_fib17() {
         let trace = fibonacci_trace([FE17::new(1), FE17::new(1)], 4);
 
+        let trace_table = TraceTable {
+            table: trace.clone(),
+            num_cols: 1,
+        };
+
         let context = AirContext {
             options: ProofOptions {
                 blowup_factor: 2,
                 fri_number_of_queries: 1,
                 coset_offset: 3,
             },
             trace_length: trace.len(),
-            trace_info: (trace.len(), 1),
+            trace_columns: trace_table.num_cols,
             transition_degrees: vec![1],
             transition_exemptions: vec![trace.len() - 2, trace.len() - 1],
             transition_offsets: vec![0, 1, 2],
             num_transition_constraints: 1,
         };
 
-        let trace_table = TraceTable {
-            table: trace.clone(),
-            num_cols: 1,
-        };
-
         let fibonacci_air = Fibonacci17AIR::new(trace_table, context);
 
         let result = prove(&trace, &fibonacci_air);

proving_system/stark/src/prover.rs

@@ -4,7 +4,7 @@ use lambdaworks_math::{
         element::FieldElement,
         traits::{IsField, IsTwoAdicField},
     },
-    polynomial::{self, Polynomial},
+    polynomial::Polynomial,
     traits::ByteConversion,
 };
 
@@ -97,11 +97,13 @@ where
 
     // Compute DEEP composition polynomial so we can commit to it using FRI.
     let mut deep_composition_poly = compute_deep_composition_poly(
-        &trace_poly,
+        air,
+        &[trace_poly],
         &composition_poly_even,
         &composition_poly_odd,
         &z,
         &trace_primitive_root,
+        transcript,
     );
 
     // * Do FRI on the composition polynomials
@@ -118,7 +120,7 @@ where
 
     for _i in 0..air.context().options.fri_number_of_queries {
         // * Sample q_1, ..., q_m using Fiat-Shamir
-        let q_i: usize = transcript_to_usize(transcript) % 2_usize.pow(lde_root_order);
+        let q_i = transcript_to_usize(transcript) % 2_usize.pow(lde_root_order);
         transcript.append(&q_i.to_be_bytes());
 
         // * For every q_i, do FRI decommitment
@@ -146,52 +148,63 @@ where
 /// Returns the DEEP composition polynomial that the prover then commits to using
 /// FRI. This polynomial is a linear combination of the trace polynomial and the
 /// composition polynomial, with coefficients sampled by the verifier (i.e. using Fiat-Shamir).
-fn compute_deep_composition_poly<F: IsField>(
-    trace_poly: &Polynomial<FieldElement<F>>,
+fn compute_deep_composition_poly<A: AIR, F: IsField>(
+    air: &A,
+    trace_polys: &[Polynomial<FieldElement<F>>],
     even_composition_poly: &Polynomial<FieldElement<F>>,
     odd_composition_poly: &Polynomial<FieldElement<F>>,
     ood_evaluation_point: &FieldElement<F>,
     primitive_root: &FieldElement<F>,
+    transcript: &mut Transcript,
 ) -> Polynomial<FieldElement<F>> {
-    // TODO: Fiat-Shamir
-    let gamma_1 = FieldElement::one();
-    let gamma_2 = FieldElement::one();
-    let gamma_3 = FieldElement::one();
-    let gamma_4 = FieldElement::one();
-
-    let first_term = (trace_poly.clone()
-        - Polynomial::new_monomial(trace_poly.evaluate(ood_evaluation_point), 0))
-        / (Polynomial::new_monomial(FieldElement::one(), 1)
-            - Polynomial::new_monomial(ood_evaluation_point.clone(), 0));
-    let second_term = (trace_poly.clone()
-        - Polynomial::new_monomial(
-            trace_poly.evaluate(&(ood_evaluation_point * primitive_root)),
-            0,
-        ))
-        / (Polynomial::new_monomial(FieldElement::one(), 1)
-            - Polynomial::new_monomial(ood_evaluation_point * primitive_root, 0));
+    let transition_offsets = air.context().transition_offsets;
+
+    // Get the number of trace terms the DEEP composition poly will have.
+    // One coefficient will be sampled for each of them.
+    let n_trace_terms = transition_offsets.len() * trace_polys.len();
+    let mut trace_term_coeffs = Vec::with_capacity(n_trace_terms);
+    for _ in 0..n_trace_terms {
+        trace_term_coeffs.push(transcript_to_field::<F>(transcript));
+    }
 
-    // Evaluate in X^2
-    let even_composition_poly = polynomial::compose(
-        even_composition_poly,
-        &Polynomial::new_monomial(FieldElement::one(), 2),
-    );
-    let odd_composition_poly = polynomial::compose(
-        odd_composition_poly,
-        &Polynomial::new_monomial(FieldElement::one(), 2),
+    // Get coefficients for even and odd terms of the composition polynomial H(x)
+    let gamma_even = transcript_to_field::<F>(transcript);
+    let gamma_odd = transcript_to_field::<F>(transcript);
+
+    // Get trace evaluations needed for the trace terms of the deep composition polynomial
+    let trace_evaluations = Frame::get_trace_evaluations(
+        trace_polys,
+        ood_evaluation_point,
+        &transition_offsets,
+        primitive_root,
     );
 
-    let third_term = (even_composition_poly.clone()
+    // Compute all the trace terms of the deep composition polynomial. There will be one
+    // term for every trace polynomial and every trace evaluation.
+    let mut trace_terms = Polynomial::zero();
+    for (trace_evaluation, trace_poly) in trace_evaluations.iter().zip(trace_polys) {
+        for (eval, coeff) in trace_evaluation.iter().zip(&trace_term_coeffs) {
+            let poly = (trace_poly.clone()
+                - Polynomial::new_monomial(trace_poly.evaluate(eval), 0))
+                / (Polynomial::new_monomial(FieldElement::<F>::one(), 1)
+                    - Polynomial::new_monomial(eval.clone(), 0));
+
+            trace_terms = trace_terms + poly * coeff.clone();
+        }
+    }
+
+    let even_composition_poly_term = (even_composition_poly.clone()
         - Polynomial::new_monomial(
             even_composition_poly.evaluate(&ood_evaluation_point.clone()),
             0,
         ))
         / (Polynomial::new_monomial(FieldElement::one(), 1)
             - Polynomial::new_monomial(ood_evaluation_point * ood_evaluation_point, 0));
-    let fourth_term = (odd_composition_poly.clone()
+
+    let odd_composition_poly_term = (odd_composition_poly.clone()
         - Polynomial::new_monomial(odd_composition_poly.evaluate(ood_evaluation_point), 0))
         / (Polynomial::new_monomial(FieldElement::one(), 1)
             - Polynomial::new_monomial(ood_evaluation_point * ood_evaluation_point, 0));
 
-    first_term * gamma_1 + second_term * gamma_2 + third_term * gamma_3 + fourth_term * gamma_4
+    trace_terms + even_composition_poly_term * gamma_even + odd_composition_poly_term * gamma_odd
 }

proving_system/stark/src/verifier.rs

@@ -116,6 +116,18 @@ where
     let lde_root_order =
         (air.context().trace_length * air.options().blowup_factor as usize).trailing_zeros();
 
+    // We have to make the call to `transcript.challenge()` a number of times since we need
+    // the transcript to be in the same state as the one in the prover at this stage.
+    // The prover samples coefficients when building the deep composition polynomial. These
+    // sampling is not done in the verifier hence we need to make this dummy calls.
+    // There will be one call for each trace term in the deep composition polynomial + 2 from
+    // the even and odd terms of the H(x) polynomial.
+    let deep_poly_challenges =
+        air.context().transition_offsets.len() * air.context().trace_columns + 2;
+    (0..deep_poly_challenges).for_each(|_| {
+        transcript.challenge();
+    });
+
     // construct vector of betas
     let mut beta_list = Vec::new();
     let count_betas = proof.fri_layers_merkle_roots.len() - 1;
@@ -137,7 +149,7 @@ where
         let last_evaluation_bytes = last_evaluation.to_bytes_be();
         transcript.append(&last_evaluation_bytes);
 
-        let q_i: usize = transcript_to_usize(transcript) % (2_usize.pow(lde_root_order));
+        let q_i = transcript_to_usize(transcript) % (2_usize.pow(lde_root_order));
         transcript.append(&q_i.to_be_bytes());
 
         let fri_decommitment = &proof_i.fri_decommitment;