Refactor lde_onto_coset

0xPolygonZero · Aug 1, 2023 · 2b9133a · 2b9133a
1 parent 53d8fd2
commit 2b9133a
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 46 deletions.
diff --git a/evm/src/prover.rs b/evm/src/prover.rs
@@ -4,7 +4,6 @@ use anyhow::{ensure, Result};
 use itertools::Itertools;
 use once_cell::sync::Lazy;
 use plonky2::field::extension::Extendable;
-use plonky2::field::fft::ifft;
 use plonky2::field::packable::Packable;
 use plonky2::field::packed::PackedField;
 use plonky2::field::polynomial::{PolynomialCoeffs, PolynomialValues};
@@ -499,29 +498,16 @@ where
     // When opening the `Z`s polys at the "next" point, need to look at the point `next_step` steps away.
     let next_step = 1 << quotient_degree_bits;
 
-    let powers_vec = SmallPowers::<F>::new(F::coset_shift().to_noncanonical_u64() as u32)
+    let powers_vec: Vec<F> = SmallPowers::new(F::coset_shift().to_noncanonical_u64() as u32)
         .take(degree << quotient_degree_bits)
         .collect_vec();
-    // Evaluation of the first Lagrange polynomial on the LDE domain, with custom lde_onto_coset.
-    let lagrange_first = PolynomialValues::<F>::selector(degree, 0);
-    let coeffs = ifft(lagrange_first).lde(quotient_degree_bits);
-    let modified_poly: PolynomialCoeffs<F> = powers_vec
-        .iter()
-        .zip(&coeffs.coeffs)
-        .map(|(&r, &c)| r * c)
-        .collect::<Vec<_>>()
-        .into();
-    let lagrange_first = modified_poly.fft_with_options(None, None);
-    // Evaluation of the last Lagrange polynomial on the LDE domain, with custom lde_onto_coset.
-    let lagrange_last = PolynomialValues::<F>::selector(degree, degree - 1);
-    let coeffs = ifft(lagrange_last).lde(quotient_degree_bits);
-    let modified_poly: PolynomialCoeffs<F> = powers_vec
-        .into_iter()
-        .zip(&coeffs.coeffs)
-        .map(|(r, &c)| r * c)
-        .collect::<Vec<_>>()
-        .into();
-    let lagrange_last = modified_poly.fft_with_options(None, None);
+
+    // Evaluation of the first Lagrange polynomial on the LDE domain.
+    let lagrange_first = PolynomialValues::selector(degree, 0)
+        .lde_onto_coset(quotient_degree_bits, powers_vec.iter());
+    // Evaluation of the last Lagrange polynomial on the LDE domain.
+    let lagrange_last = PolynomialValues::selector(degree, degree - 1)
+        .lde_onto_coset(quotient_degree_bits, powers_vec.iter());
 
     let z_h_on_coset = ZeroPolyOnCoset::<F>::new(degree_bits, quotient_degree_bits);
 

diff --git a/field/src/polynomial/mod.rs b/field/src/polynomial/mod.rs
@@ -82,9 +82,17 @@ impl<F: Field> PolynomialValues<F> {
     }
 
     /// Low-degree extend `Self` (seen as evaluations over the subgroup) onto a coset.
-    pub fn lde_onto_coset(self, rate_bits: usize) -> Self {
+    pub fn lde_onto_coset<I>(self, rate_bits: usize, shift_powers: I) -> Self
+    where
+        I: Iterator<Item = F>,
+    {
         let coeffs = ifft(self).lde(rate_bits);
-        coeffs.coset_fft_with_options(F::coset_shift(), Some(rate_bits), None)
+        let modified_poly: PolynomialCoeffs<F> = shift_powers
+            .zip(&coeffs.coeffs)
+            .map(|(r, &c)| r * c)
+            .collect::<Vec<_>>()
+            .into();
+        modified_poly.fft_with_options(Some(rate_bits), None)
     }
 
     pub fn degree(&self) -> usize {

diff --git a/starky/src/prover.rs b/starky/src/prover.rs
@@ -4,7 +4,6 @@ use core::iter::once;
 use anyhow::{ensure, Result};
 use itertools::Itertools;
 use plonky2::field::extension::Extendable;
-use plonky2::field::fft::ifft;
 use plonky2::field::packable::Packable;
 use plonky2::field::packed::PackedField;
 use plonky2::field::polynomial::{PolynomialCoeffs, PolynomialValues};
@@ -228,29 +227,16 @@ where
     // When opening the `Z`s polys at the "next" point, need to look at the point `next_step` steps away.
     let next_step = 1 << quotient_degree_bits;
 
-    let powers_vec = SmallPowers::<F>::new(F::coset_shift().to_noncanonical_u64() as u32)
+    let powers_vec: Vec<F> = SmallPowers::new(F::coset_shift().to_noncanonical_u64() as u32)
         .take(degree << quotient_degree_bits)
         .collect_vec();
-    // Evaluation of the first Lagrange polynomial on the LDE domain, with custom lde_onto_coset.
-    let lagrange_first = PolynomialValues::<F>::selector(degree, 0);
-    let coeffs = ifft(lagrange_first).lde(quotient_degree_bits);
-    let modified_poly: PolynomialCoeffs<F> = powers_vec
-        .iter()
-        .zip(&coeffs.coeffs)
-        .map(|(&r, &c)| r * c)
-        .collect::<Vec<_>>()
-        .into();
-    let lagrange_first = modified_poly.fft_with_options(None, None);
-    // Evaluation of the last Lagrange polynomial on the LDE domain, with custom lde_onto_coset.
-    let lagrange_last = PolynomialValues::<F>::selector(degree, degree - 1);
-    let coeffs = ifft(lagrange_last).lde(quotient_degree_bits);
-    let modified_poly: PolynomialCoeffs<F> = powers_vec
-        .into_iter()
-        .zip(&coeffs.coeffs)
-        .map(|(r, &c)| r * c)
-        .collect::<Vec<_>>()
-        .into();
-    let lagrange_last = modified_poly.fft_with_options(None, None);
+
+    // Evaluation of the first Lagrange polynomial on the LDE domain.
+    let lagrange_first = PolynomialValues::selector(degree, 0)
+        .lde_onto_coset(quotient_degree_bits, powers_vec.iter());
+    // Evaluation of the last Lagrange polynomial on the LDE domain.
+    let lagrange_last = PolynomialValues::selector(degree, degree - 1)
+        .lde_onto_coset(quotient_degree_bits, powers_vec.iter());
     let z_h_on_coset = ZeroPolyOnCoset::<F>::new(degree_bits, quotient_degree_bits);
 
     // Retrieve the LDE values at index `i`.