cleanup

Author: zlangley

crates/stark-backend/src/gkr/prover.rs

@@ -54,11 +54,13 @@ impl<'a, F: Field> FixedFirstHypercubeEqEvals<'a, F> {
 
         for &y_i in y.iter().rev() {
             let (left, right) = evals.split_at_mut(curr_len);
-            left.par_iter_mut().zip(right.par_iter_mut()).for_each(|(l, r)| {
-                let tmp = *l * y_i;
-                *r = tmp;
-                *l -= tmp;
-            });
+            left.par_iter_mut()
+                .zip(right.par_iter_mut())
+                .for_each(|(l, r)| {
+                    let tmp = *l * y_i;
+                    *r = tmp;
+                    *l -= tmp;
+                });
             curr_len *= 2;
         }
         evals
@@ -91,7 +93,7 @@ impl<F> Deref for FixedFirstHypercubeEqEvals<'_, F> {
 ///
 /// P(x) = eq(x, y) * (numer(x) + lambda * denom(x))
 /// ```
-struct GkrMultivariatePolyOracle<'a, F: Clone> {
+struct GkrMultivariatePolyOracle<'a, F> {
     pub eq_evals: &'a FixedFirstHypercubeEqEvals<'a, F>,
     pub input_layer: Layer<F>,
     pub eq_fixed_var_correction: F,
@@ -114,12 +116,24 @@ impl<F: Field> MultivariatePolyOracle<F> for GkrMultivariatePolyOracle<'_, F> {
             Layer::LogUpGeneric {
                 numerators,
                 denominators,
-            } => eval_logup_sum(self.eq_evals, numerators, denominators, n_terms, self.lambda),
+            } => eval_logup_sum(
+                self.eq_evals,
+                numerators,
+                denominators,
+                n_terms,
+                self.lambda,
+            ),
         };
 
         eval_at_0 *= self.eq_fixed_var_correction;
         eval_at_2 *= self.eq_fixed_var_correction;
-        correct_sum_as_poly_in_first_variable(eval_at_0, eval_at_2, claim, self.eq_evals.y, n_variables)
+        correct_sum_as_poly_in_first_variable(
+            eval_at_0,
+            eval_at_2,
+            claim,
+            self.eq_evals.y,
+            n_variables,
+        )
     }
 
     fn fix_first_in_place(&mut self, alpha: F) {
@@ -197,14 +211,8 @@ fn eval_logup_sum<F: Field>(
             let (n2_0, d2_0) = (n1_0.double() - n0_0, d1_0.double() - d0_0);
             let (n2_1, d2_1) = (n1_1.double() - n0_1, d1_1.double() - d0_1);
 
-            let (num_t0, den_t0) = (
-                n0_0 * d0_1 + n0_1 * d0_0,
-                d0_0 * d0_1,
-            );
-            let (num_t2, den_t2) = (
-                n2_0 * d2_1 + n2_1 * d2_0,
-                d2_0 * d2_1,
-            );
+            let (num_t0, den_t0) = (n0_0 * d0_1 + n0_1 * d0_0, d0_0 * d0_1);
+            let (num_t2, den_t2) = (n2_0 * d2_1 + n2_1 * d2_0, d2_0 * d2_1);
 
             let eq = eq_evals[i];
             let eval_t0 = eq * (num_t0 + lambda * den_t0);
@@ -281,23 +289,25 @@ pub fn prove_batch<F: Field, EF: ExtensionField<F>>(
         .collect();
 
     let mut output_claims_by_instance = vec![None; n_instances];
-    let mut layer_masks_by_instance = (0..n_instances).map(|_| Vec::new()).collect_vec();
+    let mut layer_masks_by_instance = vec![vec![]; n_instances];
     let mut sumcheck_proofs = Vec::new();
 
     let mut ood_point = Vec::new();
     let mut claims_to_verify_by_instance = vec![None; n_instances];
 
-    for layer in 0..n_layers {
-        let n_remaining_layers = n_layers - layer;
+    let mut output_instances_by_layer: Vec<Vec<usize>> = vec![Vec::new(); n_layers];
+    for (instance, &instance_n_layers) in n_layers_by_instance.iter().enumerate() {
+        let output_layer = n_layers - instance_n_layers;
+        output_instances_by_layer[output_layer].push(instance);
+    }
 
+    for layer in 0..n_layers {
         // Check all the instances for output layers.
-        for (instance, layers) in layers_by_instance.iter_mut().enumerate() {
-            if n_layers_by_instance[instance] == n_remaining_layers {
-                let output_layer = layers.next().unwrap();
-                let output_layer_values = output_layer.try_into_output_layer_values().unwrap();
-                claims_to_verify_by_instance[instance] = Some(output_layer_values.clone());
-                output_claims_by_instance[instance] = Some(output_layer_values);
-            }
+        for &instance in &output_instances_by_layer[layer] {
+            let output_layer = layers_by_instance[instance].next().unwrap();
+            let output_layer_values = output_layer.try_into_output_layer_values().unwrap();
+            claims_to_verify_by_instance[instance] = Some(output_layer_values.clone());
+            output_claims_by_instance[instance] = Some(output_layer_values);
         }
 
         // Seed the channel with layer claims.

crates/stark-backend/src/gkr/types.rs

@@ -1,15 +1,10 @@
-use std::ops::Index;
-
-use p3_maybe_rayon::prelude::*;
 use p3_field::Field;
+use p3_maybe_rayon::prelude::*;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 use crate::{
-    poly::{
-        multi::{fold_mle_evals, Mle, MultivariatePolyOracle},
-        uni::Fraction,
-    },
+    poly::multi::{fold_mle_evals, Mle, MultivariatePolyOracle},
     sumcheck::{SumcheckError, SumcheckProof},
 };
 
@@ -207,10 +202,10 @@ impl<F: Field> Layer<F> {
             Self::LogUpGeneric {
                 numerators,
                 denominators,
-            } =>  {
+            } => {
                 numerators.fix_first_in_place(x0);
                 denominators.fix_first_in_place(x0);
-            },
+            }
         }
     }
 }

crates/stark-backend/src/interaction/gkr_log_up.rs

@@ -17,6 +17,7 @@ use rayon::iter::IntoParallelRefIterator;
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
+use crate::interaction::SymbolicInteraction;
 use crate::{
     air_builders::symbolic::SymbolicConstraints,
     gkr,
@@ -27,7 +28,6 @@ use crate::{
     },
     rap::PermutationAirBuilderWithExposedValues,
 };
-use crate::interaction::SymbolicInteraction;
 
 pub struct GkrLogUpPhase<F, EF, Challenger> {
     // FIXME: USE THIS IN POW
@@ -112,11 +112,7 @@ where
         _params_per_air: &[&Self::PartialProvingKey],
         trace_view_per_air: &[PairTraceView<'_, F>],
     ) -> Option<(Self::PartialProof, RapPhaseProverData<EF>)> {
-        let all_interactions = interactions_per_air
-            .iter()
-            .cloned()
-            .flatten()
-            .collect_vec();
+        let all_interactions = interactions_per_air.iter().cloned().flatten().collect_vec();
         if all_interactions.is_empty() {
             return None;
         }

crates/stark-backend/src/interaction/gkr_log_up/prove.rs

@@ -1,6 +1,6 @@
 use std::iter;
 
-use itertools::{izip, Itertools};
+use itertools::Itertools;
 use p3_challenger::FieldChallenger;
 use p3_field::{ExtensionField, Field};
 use p3_matrix::{
@@ -11,10 +11,7 @@ use p3_maybe_rayon::prelude::*;
 use p3_util::log2_strict_usize;
 
 use crate::{
-    air_builders::symbolic::{
-        symbolic_expression::{SymbolicEvaluator, SymbolicExpression},
-        SymbolicConstraints,
-    },
+    air_builders::symbolic::symbolic_expression::{SymbolicEvaluator, SymbolicExpression},
     gkr::{GkrArtifact, Layer, Layer::LogUpGeneric},
     interaction::{
         gkr_log_up::{num_interaction_dimensions, GkrAuxData, GkrLogUpPhase},
@@ -91,7 +88,9 @@ impl<EF: Field> GkrLogUpInstance<EF> {
                         .iter()
                         .chain(iter::once(&SymbolicExpression::Constant(b)))
                         .zip(beta_pows)
-                        .fold(EF::ZERO, |acc, (expr, &beta)| acc + beta * evaluator.eval_expr(expr));
+                        .fold(EF::ZERO, |acc, (expr, &beta)| {
+                            acc + beta * evaluator.eval_expr(expr)
+                        });
 
                     sigma_row[col] = sigma;
                 }
@@ -167,26 +166,23 @@ where
     ) -> GkrAuxData<EF> {
         let ood_point = &gkr_artifact.ood_point;
 
-        let max_interactions = interactions_per_air
-            .iter()
-            .map(|v| v.len())
-            .max()
-            .unwrap();
+        let max_interactions = interactions_per_air.iter().map(|v| v.len()).max().unwrap();
         let gamma_pows = gamma.powers().take(2 * max_interactions).collect_vec();
 
-        let trace_view_per_air_filtered =
-            interactions_per_air
-                .iter()
-                .zip(trace_view_per_air.iter())
-                .filter_map(|(interactions, view)| {
-                    if interactions.is_empty() {
-                        None
-                    } else {
-                        Some(view)
-                    }
-                }).collect_vec();
-
-        let results: Vec<_> = trace_view_per_air_filtered.par_iter()
+        let trace_view_per_air_filtered = interactions_per_air
+            .iter()
+            .zip(trace_view_per_air.iter())
+            .filter_map(|(interactions, view)| {
+                if interactions.is_empty() {
+                    None
+                } else {
+                    Some(view)
+                }
+            })
+            .collect_vec();
+
+        let results: Vec<_> = trace_view_per_air_filtered
+            .par_iter()
             .zip(gkr_instances.par_iter())
             .zip(gkr_artifact.n_variables_by_instance.par_iter())
             .map(|((trace_view, gkr_instance), &n_vars)| {
@@ -222,8 +218,7 @@ where
                 after_challenge_trace_per_air.push(None);
                 exposed_values_per_air.push(None);
             } else {
-                let
-                    (after_trace, exposed, count_mle, sigma_mle) = results_iter.next().unwrap();
+                let (after_trace, exposed, count_mle, sigma_mle) = results_iter.next().unwrap();
 
                 for (count_mle_claim, sigma_mle_claim) in count_mle.iter().zip(sigma_mle.iter()) {
                     challenger.observe_ext_element(*count_mle_claim);

crates/stark-backend/src/poly/multi.rs

@@ -89,7 +89,6 @@ impl<F: Field> MultivariatePolyOracle<F> for Mle<F> {
         UnivariatePolynomial::from_coeffs(vec![y0, slope])
     }
 
-
     fn fix_first_in_place(&mut self, alpha: F) {
         let midpoint = self.len() / 2;
         let (lhs_evals, rhs_evals) = self.split_at_mut(midpoint);
@@ -258,14 +257,8 @@ mod test {
         // -(1 - x_2) - 2 x_2 + 6 (1 - x_2) + 8 x_2 = x_2 + 5
         mle.fix_first_in_place(alpha);
 
-        assert_eq!(
-            mle.eval(&[BabyBear::ZERO]),
-            BabyBear::from_canonical_u32(5)
-        );
-        assert_eq!(
-            mle.eval(&[BabyBear::ONE]),
-            BabyBear::from_canonical_u32(6)
-        );
+        assert_eq!(mle.eval(&[BabyBear::ZERO]), BabyBear::from_canonical_u32(5));
+        assert_eq!(mle.eval(&[BabyBear::ONE]), BabyBear::from_canonical_u32(6));
     }
 
     #[test]

crates/stark-backend/src/poly/uni.rs

@@ -1,7 +1,6 @@
 //! Copied from starkware-libs/stwo under Apache-2.0 license.
 use p3_field::Field;
 use serde::{Deserialize, Serialize};
-use std::ops::MulAssign;
 use std::{
     iter::Sum,
     ops::{Add, Deref, Mul, Neg, Sub},
@@ -85,7 +84,10 @@ impl<F: Field> UnivariatePolynomial<F> {
 
     fn remove_trailing_zeroes(&mut self) {
         self.coeffs.truncate(
-            self.coeffs.iter().rposition(|c| !c.is_zero()).map_or(0, |i| i + 1)
+            self.coeffs
+                .iter()
+                .rposition(|c| !c.is_zero())
+                .map_or(0, |i| i + 1),
         );
     }
 

crates/stark-backend/src/sumcheck.rs

@@ -88,13 +88,13 @@ where
                 let eval_at_0 = round_poly.evaluate_at_zero();
                 let eval_at_1 = round_poly.evaluate_at_one();
 
-                assert_eq!(
+                debug_assert_eq!(
                     eval_at_0 + eval_at_1,
                     claim,
                     "Round {round}, poly {i}: eval(0) + eval(1) != claim ({} != {claim})",
                     eval_at_0 + eval_at_1,
                 );
-                assert!(
+                debug_assert!(
                     round_poly.degree() <= MAX_DEGREE,
                     "Round {round}, poly {i}: degree {} > max {MAX_DEGREE}",
                     round_poly.degree(),
@@ -115,17 +115,16 @@ where
 
         let challenge = challenger.sample_ext_element();
 
-        claims.par_iter_mut().zip(this_round_polys
-            .par_iter())
+        claims
+            .par_iter_mut()
+            .zip(this_round_polys.par_iter())
             .for_each(|(claim, round_poly)| *claim = round_poly.evaluate(challenge));
 
-        polys
-            .par_iter_mut()
-            .for_each(|multivariate_poly| {
-                if n_remaining_rounds == multivariate_poly.arity() {
-                    multivariate_poly.fix_first_in_place(challenge)
-                }
-            });
+        polys.par_iter_mut().for_each(|multivariate_poly| {
+            if n_remaining_rounds == multivariate_poly.arity() {
+                multivariate_poly.fix_first_in_place(challenge)
+            }
+        });
 
         round_polys.push(round_poly);
         evaluation_point.push(challenge);