chore: Add coset_fft API (#294)

* add coset fft API

* modify code to match coset fft

* remove hardcoded coset_gens

* small clean-up

* add initial comment

* add a note on the computed interpolation value

* Add more comments and move random linear commitment to the proofs in the same place

* fix: doc numbering

cryptography/erasure_codes/src/reed_solomon.rs

@@ -1,7 +1,11 @@
-use bls12_381::{batch_inversion::batch_inverse, ff::Field, Scalar};
+use bls12_381::{
+    batch_inversion::batch_inverse,
+    ff::{Field, PrimeField},
+    Scalar,
+};
 
 use crate::errors::RSError;
-use polynomial::{domain::Domain, poly_coeff::vanishing_poly};
+use polynomial::{domain::Domain, poly_coeff::vanishing_poly, CosetFFT};
 
 /// ErasurePattern is an abstraction created to capture the idea
 /// that erasures do not appear in completely random locations.
@@ -62,6 +66,8 @@ pub struct ReedSolomon {
     /// The domain that we will use to efficiently compute the vanishing polynomial with, when the erasure pattern
     /// being used is `BlockSynchronizedErasures`.
     block_size_domain: Domain,
+
+    fft_coset_gen: CosetFFT,
 }
 
 impl ReedSolomon {
@@ -77,13 +83,16 @@ impl ReedSolomon {
 
         let block_size_domain = Domain::new(block_size);
 
+        let fft_coset_gen = CosetFFT::new(Scalar::MULTIPLICATIVE_GENERATOR);
+
         Self {
             poly_len,
             evaluation_domain,
             expansion_factor,
             block_size,
             block_size_domain,
             num_blocks,
+            fft_coset_gen,
         }
     }
 
@@ -294,8 +303,12 @@ impl ReedSolomon {
 
         let dz_poly = self.evaluation_domain.ifft_scalars(ez_eval);
 
-        let coset_dz_eval = self.evaluation_domain.coset_fft_scalars(dz_poly);
-        let mut inv_coset_z_x_eval = self.evaluation_domain.coset_fft_scalars(z_x);
+        let coset_dz_eval = self
+            .evaluation_domain
+            .coset_fft_scalars(dz_poly, &self.fft_coset_gen);
+        let mut inv_coset_z_x_eval = self
+            .evaluation_domain
+            .coset_fft_scalars(z_x, &self.fft_coset_gen);
         // We know that none of the values will be zero since we are evaluating z_x
         // over a coset, that we know it has no roots in.
         batch_inverse(&mut inv_coset_z_x_eval);
@@ -307,7 +320,7 @@ impl ReedSolomon {
 
         let coefficients = self
             .evaluation_domain
-            .coset_ifft_scalars(coset_quotient_eval);
+            .coset_ifft_scalars(coset_quotient_eval, &self.fft_coset_gen);
 
         // Check that the polynomial being returned has the correct degree
         //

cryptography/kzg_multi_open/src/fk20/verifier.rs

@@ -3,10 +3,10 @@ use crate::{
     verification_key::VerificationKey,
 };
 use bls12_381::{
-    batch_inversion::batch_inverse, ff::Field, g1_batch_normalize, lincomb::g1_lincomb,
-    multi_pairings, reduce_bytes_to_scalar_bias, G1Point, G2Point, G2Prepared, Scalar,
+    ff::Field, g1_batch_normalize, lincomb::g1_lincomb, multi_pairings,
+    reduce_bytes_to_scalar_bias, G1Point, G2Point, G2Prepared, Scalar,
 };
-use polynomial::{domain::Domain, poly_coeff::poly_add};
+use polynomial::{domain::Domain, poly_coeff::poly_add, CosetFFT};
 use sha2::{Digest, Sha256};
 use std::mem::size_of;
 
@@ -45,10 +45,12 @@ pub struct FK20Verifier {
     tau_pow_n: G2Prepared,
     // [-1]_2
     neg_g2_gen: G2Prepared,
-    //
-    pub coset_gens_pow_n: Vec<Scalar>,
-    //
-    inv_coset_gens_pow_n: Vec<Vec<Scalar>>,
+    // Bit reversed vector of the coset generators raised
+    // to the power of `n`, needed to verify a multi opening proof.
+    pub bit_reversed_coset_gens_pow_n: Vec<Scalar>,
+    // Bit reversed vector of the coset generators that
+    // we will use to compute an inverse coset IFFT.
+    bit_reversed_coset_fft_gens: Vec<CosetFFT>,
 }
 
 impl FK20Verifier {
@@ -78,23 +80,16 @@ impl FK20Verifier {
             .map(|&coset_gen| coset_gen.pow_vartime([n as u64]))
             .collect();
 
-        let inv_coset_gens_pow_n: Vec<_> = coset_gens
-            .iter()
-            .map(|&coset_gen| {
-                let mut inv_coset_gen_powers = compute_powers(coset_gen, n);
-                batch_inverse(&mut inv_coset_gen_powers); // The coset generators are all roots of unity, so none of them will be zero
-                inv_coset_gen_powers
-            })
-            .collect();
+        let coset_fft_gens: Vec<_> = coset_gens.iter().map(|gen| CosetFFT::new(*gen)).collect();
 
         Self {
             verification_key,
             coset_gens_bit_reversed: coset_gens,
             coset_domain,
             tau_pow_n,
             neg_g2_gen,
-            coset_gens_pow_n,
-            inv_coset_gens_pow_n,
+            bit_reversed_coset_gens_pow_n: coset_gens_pow_n,
+            bit_reversed_coset_fft_gens: coset_fft_gens,
         }
     }
 
@@ -145,7 +140,9 @@ impl FK20Verifier {
         // The batch size corresponds to how many openings, we ultimately want to be verifying.
         let batch_size = bit_reversed_coset_indices.len();
 
-        // Compute random challenges for batching the opening together.
+        // 1. Compute random challenges for batching the opening together.
+        //
+        // From hereon out, `random` will refer to using these random challenges.
         //
         // We compute one challenge `r` using fiat-shamir and the rest are powers of `r`
         // This is safe because of the Schwartz-Zippel Lemma.
@@ -160,14 +157,28 @@ impl FK20Verifier {
         let r_powers = compute_powers(r, batch_size);
         let num_unique_commitments = deduplicated_commitments.len();
 
-        // First compute a random linear combination of the proofs
+        // 2. Compute a random linear combination of the proofs
         //
         // Safety: This unwrap can never trigger because `r_powers.len()` is `batch_size`
         // and `bit_reversed_proofs.len()` will equal `batch_size` since we must have a proof for each item in the batch.
         let comm_random_sum_proofs = g1_lincomb(bit_reversed_proofs, &r_powers)
             .expect("number of proofs and number of r_powers should be the same");
 
-        // Now compute a random linear combination of the commitments
+        // 3. Compute a weighted random linear combination of the proofs
+        //
+        // Where the `weight` refers to the coset_generators to the power of `n`
+        let mut weighted_r_powers = Vec::with_capacity(batch_size);
+        for (bit_reversed_coset_index, r_power) in bit_reversed_coset_indices.iter().zip(&r_powers)
+        {
+            let coset_gen_pow_n =
+                self.bit_reversed_coset_gens_pow_n[*bit_reversed_coset_index as usize];
+            weighted_r_powers.push(r_power * coset_gen_pow_n);
+        }
+        // Safety: This should never panic since `bit_reversed_proofs.len()` is equal to the batch_size.
+        let random_weighted_sum_proofs = g1_lincomb(bit_reversed_proofs, &weighted_r_powers)
+            .expect("number of proofs and number of weighted_r_powers should be the same");
+
+        // 4. Compute a random linear combination of the commitments
         //
         // For each commitment_index/commitment, we add its contribution of `r` to
         // the associated weight for that commitment.
@@ -181,62 +192,36 @@ impl FK20Verifier {
         //
         // The extra field additions are being calculated in the for loop below.
         let mut weights = vec![Scalar::ZERO; num_unique_commitments];
-        for (commitment_index, r_power) in commitment_indices.iter().zip(r_powers.iter()) {
+        for (commitment_index, r_power) in commitment_indices.iter().zip(&r_powers) {
             weights[*commitment_index as usize] += r_power;
         }
-
         // Safety: This unwrap will never trigger because the length of `weights` has been initialized
         // to be `deduplicated_commitments.len()`.
         //
         // This only panics, if `deduplicated_commitments.len()` != `weights.len()`
         let random_sum_commitments = g1_lincomb(deduplicated_commitments, &weights)
             .expect("number of row_commitments and number of weights should be the same");
 
-        // Linearly combine the interpolation polynomials using the same randomness `r`
-        let mut random_sum_interpolation_poly = Vec::new();
-        let coset_evals = bit_reversed_coset_evals.to_vec();
-        for (k, mut coset_eval) in coset_evals.into_iter().enumerate() {
-            // Reverse the order, so it matches the fft domain
-            reverse_bit_order(&mut coset_eval);
-
-            // Compute the interpolation polynomial
-            let ifft_scalars = self.coset_domain.ifft_scalars(coset_eval);
-            let inv_coset_gen_pow_n =
-                &self.inv_coset_gens_pow_n[bit_reversed_coset_indices[k] as usize];
-            let ifft_scalars: Vec<_> = ifft_scalars
-                .into_iter()
-                .zip(inv_coset_gen_pow_n)
-                .map(|(scalar, inv_h_k_pow)| scalar * inv_h_k_pow)
-                .collect();
-
-            // Scale the interpolation polynomial by the challenge
-            let scale_factor = r_powers[k];
-            let scaled_interpolation_poly = ifft_scalars
-                .into_iter()
-                .map(|coeff| coeff * scale_factor)
-                .collect::<Vec<_>>();
-
-            random_sum_interpolation_poly =
-                poly_add(random_sum_interpolation_poly, scaled_interpolation_poly);
-        }
+        // 5. Compute random linear combination of the interpolation polynomials
+        let random_sum_interpolation_poly = compute_sum_interpolation_poly(
+            &self.coset_domain,
+            &self.bit_reversed_coset_fft_gens,
+            &bit_reversed_coset_evals,
+            &bit_reversed_coset_indices,
+            &r_powers,
+        );
         let comm_random_sum_interpolation_poly = self
             .verification_key
             .commit_g1(&random_sum_interpolation_poly);
 
-        let mut weighted_r_powers = Vec::with_capacity(batch_size);
-        for (coset_index, r_power) in bit_reversed_coset_indices.iter().zip(r_powers) {
-            let coset_gen_pow_n = self.coset_gens_pow_n[*coset_index as usize];
-            weighted_r_powers.push(r_power * coset_gen_pow_n);
-        }
-
-        // Safety: This should never panic since `bit_reversed_proofs.len()` is equal to the batch_size.
-        let random_weighted_sum_proofs = g1_lincomb(bit_reversed_proofs, &weighted_r_powers)
-            .expect("number of proofs and number of weighted_r_powers should be the same");
-
-        // This is `rl` in the specs.
+        // 6. Compute pairing check
+        //
+        // Note: This variable is `rl` in the specs.
         let pairing_input_g1 = (random_sum_commitments - comm_random_sum_interpolation_poly)
             + random_weighted_sum_proofs;
 
+        // The pairings function requires elements in affine representation, so we must batch normalize the
+        // pairing inputs.
         let normalized_vectors = g1_batch_normalize(&[comm_random_sum_proofs, pairing_input_g1]);
         let random_sum_proofs = normalized_vectors[0];
         let pairing_input_g1 = normalized_vectors[1];
@@ -337,6 +322,46 @@ fn compute_powers(value: Scalar, num_elements: usize) -> Vec<Scalar> {
     powers
 }
 
+/// Computes `k` Interpolation polynomials and then combines
+/// them linearly using `k` values from `r_powers`.
+/// The computed value is I(X) = I_0(x) + r * I_1(x) + ... + r^{n-1} * I_{n-1}(x)
+fn compute_sum_interpolation_poly(
+    coset_domain: &Domain,
+    bit_reversed_coset_fft_gens: &[CosetFFT],
+    bit_reversed_coset_evals: &[Vec<Scalar>],
+    bit_reversed_coset_indices: &[CosetIndex],
+    r_powers: &[Scalar],
+) -> Vec<Scalar> {
+    let mut random_sum_interpolation_poly = Vec::new();
+
+    for ((mut bit_reversed_coset_eval, bit_reversed_coset_index), scale_factor) in
+        bit_reversed_coset_evals
+            .to_vec()
+            .into_iter()
+            .zip(bit_reversed_coset_indices)
+            .zip(r_powers)
+    {
+        // Reverse the order, so it matches the fft domain
+        reverse_bit_order(&mut bit_reversed_coset_eval);
+        let coset_eval = bit_reversed_coset_eval; // variable rename since we un-bit reversed the vector
+
+        // Compute the interpolation polynomial using a coset fft
+        let coset_gen = &bit_reversed_coset_fft_gens[*bit_reversed_coset_index as usize];
+        let ifft_scalars = coset_domain.coset_ifft_scalars(coset_eval, coset_gen);
+
+        // Scale the interpolation polynomial by the challenge
+        let scaled_interpolation_poly = ifft_scalars
+            .into_iter()
+            .map(|coeff| coeff * scale_factor)
+            .collect::<Vec<_>>();
+
+        random_sum_interpolation_poly =
+            poly_add(random_sum_interpolation_poly, scaled_interpolation_poly);
+    }
+
+    random_sum_interpolation_poly
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

cryptography/polynomial/src/coset_fft.rs

@@ -0,0 +1,19 @@
+use bls12_381::{ff::Field, Scalar};
+
+/// CosetFFt contains a generator(coset) element that can be used
+/// to compute a coset FFT and its inverse which consequently can be used to
+/// compute a coset IFFT
+#[derive(Debug, Clone)]
+pub struct CosetFFT {
+    pub generator: Scalar,
+    pub generator_inv: Scalar,
+}
+
+impl CosetFFT {
+    pub fn new(gen: Scalar) -> Self {
+        Self {
+            generator: gen,
+            generator_inv: gen.invert().expect("cosets should be non-zero"),
+        }
+    }
+}

cryptography/polynomial/src/domain.rs

@@ -1,3 +1,4 @@
+use crate::coset_fft::CosetFFT;
 use crate::fft::{fft_g1_inplace, fft_scalar_inplace, precompute_twiddle_factors};
 use crate::poly_coeff::PolyCoeff;
 use bls12_381::ff::{Field, PrimeField};
@@ -23,11 +24,6 @@ pub struct Domain {
     /// Inverse of the generator for the domain
     /// This is cached for IFFT
     pub generator_inv: Scalar,
-    /// Element used to generate a coset
-    /// of the domain
-    coset_generator: Scalar,
-    /// Inverse of the coset generator
-    coset_generator_inv: Scalar,
     /// Precomputed values for the generator to speed up
     /// the forward FFT
     twiddle_factors: Vec<Scalar>,
@@ -61,11 +57,6 @@ impl Domain {
             roots.push(prev_root * generator)
         }
 
-        let coset_generator = Scalar::MULTIPLICATIVE_GENERATOR;
-        let coset_generator_inv = coset_generator
-            .invert()
-            .expect("coset generator should not be zero");
-
         let twiddle_factors = precompute_twiddle_factors(&generator, size);
         let twiddle_factors_inv = precompute_twiddle_factors(&generator_inv, size);
 
@@ -75,8 +66,6 @@ impl Domain {
             domain_size_inv: size_as_scalar_inv,
             generator,
             generator_inv,
-            coset_generator,
-            coset_generator_inv,
             twiddle_factors,
             twiddle_factors_inv,
         }
@@ -127,15 +116,15 @@ impl Domain {
 
     /// Evaluates a polynomial at the points in the domain multiplied by a coset
     /// generator `g`.
-    pub fn coset_fft_scalars(&self, mut points: PolyCoeff) -> Vec<Scalar> {
+    pub fn coset_fft_scalars(&self, mut points: PolyCoeff, coset: &CosetFFT) -> Vec<Scalar> {
         // Pad the polynomial with zeroes, so that it is the same size as the
         // domain.
         points.resize(self.size(), Scalar::ZERO);
 
         let mut coset_scale = Scalar::ONE;
         for point in points.iter_mut() {
             *point *= coset_scale;
-            coset_scale *= self.coset_generator;
+            coset_scale *= coset.generator;
         }
         fft_scalar_inplace(&self.twiddle_factors, &mut points);
 
@@ -212,13 +201,13 @@ impl Domain {
     }
 
     /// Interpolates a polynomial over the coset of a domain
-    pub fn coset_ifft_scalars(&self, points: Vec<Scalar>) -> Vec<Scalar> {
+    pub fn coset_ifft_scalars(&self, points: Vec<Scalar>, coset: &CosetFFT) -> Vec<Scalar> {
         let mut coset_coeffs = self.ifft_scalars(points);
 
         let mut coset_scale = Scalar::ONE;
         for element in coset_coeffs.iter_mut() {
             *element *= coset_scale;
-            coset_scale *= self.coset_generator_inv;
+            coset_scale *= coset.generator_inv;
         }
         coset_coeffs
     }
@@ -268,9 +257,9 @@ mod tests {
         let polynomial: Vec<_> = (0..32).map(|i| -Scalar::from(i)).collect();
 
         let domain = Domain::new(32);
-
-        let coset_evals = domain.coset_fft_scalars(polynomial.clone());
-        let got_poly = domain.coset_ifft_scalars(coset_evals);
+        let coset_fft = CosetFFT::new(Scalar::MULTIPLICATIVE_GENERATOR);
+        let coset_evals = domain.coset_fft_scalars(polynomial.clone(), &coset_fft);
+        let got_poly = domain.coset_ifft_scalars(coset_evals, &coset_fft);
 
         assert_eq!(got_poly, polynomial);
     }

cryptography/polynomial/src/lib.rs

@@ -1,3 +1,6 @@
+mod coset_fft;
 pub mod domain;
 mod fft;
 pub mod poly_coeff;
+
+pub use coset_fft::CosetFFT;