do more in-place

Author: zlangley

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

@@ -125,20 +125,14 @@ impl<F: Field> MultivariatePolyOracle<F> for GkrMultivariatePolyOracle<'_, F> {
         correct_sum_as_poly_in_first_variable(eval_at_0, eval_at_2, claim, y, n_variables)
     }
 
-    fn partial_evaluation(self, alpha: F) -> Self {
+    fn fix_first_in_place(&mut self, alpha: F) {
         if self.has_zero_arity() {
-            return self;
+            return;
         }
 
         let z0 = self.eq_evals.y[self.eq_evals.y.len() - self.arity()];
-        let eq_fixed_var_correction = self.eq_fixed_var_correction * (alpha * z0 + (F::ONE - alpha) * (F::ONE - z0));
-
-        Self {
-            eq_evals: self.eq_evals,
-            eq_fixed_var_correction,
-            input_layer: self.input_layer.fix_first_variable(alpha),
-            lambda: self.lambda,
-        }
+        self.eq_fixed_var_correction *= alpha * z0 + (F::ONE - alpha) * (F::ONE - z0);
+        self.input_layer.fix_first_variable_in_place(alpha);
     }
 }
 

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

@@ -166,79 +166,3 @@ fn test_logup_with_generic_trace() -> Result<(), GkrError<BabyBear>> {
     );
     Ok(())
 }
-
-#[test]
-fn test_logup_with_singles_trace() -> Result<(), GkrError<BabyBear>> {
-    const N: usize = 1 << 5;
-    type F = BabyBear;
-
-    let mut rng = create_seeded_rng();
-    let denominator_values = (0..N).map(|_| rng.gen()).collect_vec();
-    let sum: Fraction<F> = denominator_values
-        .iter()
-        .map(|&d| Fraction::new(F::ONE, d))
-        .sum();
-    let denominators = Mle::from_vec(denominator_values);
-    let top_layer = Layer::LogUpSingles {
-        denominators: denominators.clone(),
-    };
-
-    let engine = default_engine();
-    let (proof, _) = gkr::prove_batch(&mut engine.new_challenger(), vec![top_layer]);
-
-    let GkrArtifact {
-        ood_point,
-        claims_to_verify_by_instance,
-        n_variables_by_instance: _,
-    } = gkr::partially_verify_batch(vec![Gate::LogUp], &proof, &mut engine.new_challenger())?;
-
-    assert_eq!(claims_to_verify_by_instance.len(), 1);
-    assert_eq!(proof.output_claims_by_instance.len(), 1);
-    assert_eq!(
-        claims_to_verify_by_instance[0],
-        [F::ONE, denominators.eval(&ood_point)]
-    );
-    assert_eq!(
-        proof.output_claims_by_instance[0],
-        [sum.numerator, sum.denominator]
-    );
-    Ok(())
-}
-
-#[test]
-fn test_logup_with_multiplicities_trace() -> Result<(), GkrError<BabyBear>> {
-    const N: usize = 1 << 5;
-    let mut rng = create_seeded_rng();
-    let numerator_values = (0..N).map(|_| rng.gen::<BabyBear>()).collect_vec();
-    let denominator_values = (0..N).map(|_| rng.gen::<BabyBear>()).collect_vec();
-    let sum: Fraction<BabyBear> = zip(&numerator_values, &denominator_values)
-        .map(|(&n, &d)| Fraction::new(n, d))
-        .sum();
-    let numerators = Mle::from_vec(numerator_values);
-    let denominators = Mle::from_vec(denominator_values);
-    let top_layer = Layer::LogUpMultiplicities {
-        numerators: numerators.clone(),
-        denominators: denominators.clone(),
-    };
-
-    let engine = default_engine();
-    let (proof, _) = gkr::prove_batch(&mut engine.new_challenger(), vec![top_layer]);
-
-    let GkrArtifact {
-        ood_point,
-        claims_to_verify_by_instance,
-        n_variables_by_instance: _,
-    } = gkr::partially_verify_batch(vec![Gate::LogUp], &proof, &mut engine.new_challenger())?;
-
-    assert_eq!(claims_to_verify_by_instance.len(), 1);
-    assert_eq!(proof.output_claims_by_instance.len(), 1);
-    assert_eq!(
-        claims_to_verify_by_instance[0],
-        [numerators.eval(&ood_point), denominators.eval(&ood_point)]
-    );
-    assert_eq!(
-        proof.output_claims_by_instance[0],
-        [sum.numerator, sum.denominator]
-    );
-    Ok(())
-}

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

@@ -197,19 +197,19 @@ impl<F: Field> Layer<F> {
     }
 
     /// Returns a transformed layer with the first variable of each column fixed to `assignment`.
-    pub fn fix_first_variable(self, x0: F) -> Self {
+    pub fn fix_first_variable_in_place(&mut self, x0: F) {
         if self.n_variables() == 0 {
-            return self;
+            return;
         }
 
         match self {
-            Self::GrandProduct(mle) => Self::GrandProduct(mle.partial_evaluation(x0)),
+            Self::GrandProduct(mle) => mle.fix_first_in_place(x0),
             Self::LogUpGeneric {
                 numerators,
                 denominators,
-            } => Self::LogUpGeneric {
-                numerators: numerators.partial_evaluation(x0),
-                denominators: denominators.partial_evaluation(x0),
+            } =>  {
+                numerators.fix_first_in_place(x0);
+                denominators.fix_first_in_place(x0);
             },
         }
     }

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

@@ -19,7 +19,7 @@ pub trait MultivariatePolyOracle<F>: Send + Sync {
     fn marginalize_first(&self, claim: F) -> UnivariatePolynomial<F>;
 
     /// Returns the multivariate polynomial `h(x_2, ..., x_n) = g(alpha, x_2, ..., x_n)`.
-    fn partial_evaluation(self, alpha: F) -> Self;
+    fn fix_first_in_place(&mut self, alpha: F);
 }
 
 /// Multilinear extension of the function defined on the boolean hypercube.
@@ -90,7 +90,7 @@ impl<F: Field> MultivariatePolyOracle<F> for Mle<F> {
     }
 
 
-    fn partial_evaluation(mut self, alpha: F) -> Self {
+    fn fix_first_in_place(&mut self, alpha: F) {
         let midpoint = self.len() / 2;
         let (lhs_evals, rhs_evals) = self.split_at_mut(midpoint);
         lhs_evals
@@ -100,7 +100,6 @@ impl<F: Field> MultivariatePolyOracle<F> for Mle<F> {
                 *lhs += alpha * (rhs - *lhs);
             });
         self.evals.truncate(midpoint);
-        self
     }
 }
 
@@ -254,17 +253,17 @@ mod test {
             BabyBear::from_canonical_u32(4),
         ];
         // (1 - x_1)(1 - x_2) + 2 (1 - x_1) x_2 + 3 x_1 (1 - x_2) + 4 x_1 x_2
-        let mle = Mle::from_vec(evals);
+        let mut mle = Mle::from_vec(evals);
         let alpha = BabyBear::from_canonical_u32(2);
         // -(1 - x_2) - 2 x_2 + 6 (1 - x_2) + 8 x_2 = x_2 + 5
-        let partial_eval = mle.partial_evaluation(alpha);
+        mle.fix_first_in_place(alpha);
 
         assert_eq!(
-            partial_eval.eval(&[BabyBear::ZERO]),
+            mle.eval(&[BabyBear::ZERO]),
             BabyBear::from_canonical_u32(5)
         );
         assert_eq!(
-            partial_eval.eval(&[BabyBear::ONE]),
+            mle.eval(&[BabyBear::ONE]),
             BabyBear::from_canonical_u32(6)
         );
     }

crates/stark-backend/src/sumcheck.rs

@@ -121,16 +121,13 @@ where
             .map(|round_poly| round_poly.evaluate(challenge))
             .collect();
 
-        polys = polys
-            .into_par_iter()
-            .map(|multivariate_poly| {
-                if n_remaining_rounds != multivariate_poly.arity() {
-                    multivariate_poly
-                } else {
-                    multivariate_poly.partial_evaluation(challenge)
+        polys
+            .par_iter_mut()
+            .for_each(|multivariate_poly| {
+                if n_remaining_rounds == multivariate_poly.arity() {
+                    multivariate_poly.fix_first_in_place(challenge)
                 }
-            })
-            .collect();
+            });
 
         round_polys.push(round_poly);
         evaluation_point.push(challenge);