Remove innecesary collect in fft_g1

nomos-da/kzgrs/src/common.rs

@@ -4,10 +4,11 @@ use std::fmt;
 // crates
 use crate::{FieldElement, BYTES_PER_FIELD_ELEMENT};
 use ark_bls12_381::fr::Fr;
-use ark_ff::Zero;
+use ark_ff::{BigInt, FftField, Field, Zero};
 use ark_poly::domain::general::GeneralEvaluationDomain;
 use ark_poly::evaluations::univariate::Evaluations;
 use ark_poly::univariate::DensePolynomial;
+use ark_poly::{EvaluationDomain, Radix2EvaluationDomain};
 use blst::BLST_ERROR;
 use num_bigint::BigUint;
 use thiserror::Error;
@@ -122,6 +123,11 @@ pub fn field_element_from_bytes_le(b: &[u8]) -> FieldElement {
     FieldElement::from(BigUint::from_bytes_le(b))
 }
 
+pub fn compute_roots_of_unity(size: usize) -> Vec<Fr> {
+    let domain = GeneralEvaluationDomain::new(size).unwrap();
+    domain.elements().take(size).collect()
+}
+
 #[cfg(test)]
 mod test {
     use super::{bytes_to_evaluations, bytes_to_polynomial, KzgRsError};

nomos-da/kzgrs/src/fft.rs

@@ -4,6 +4,7 @@ use ark_ec::{AffineRepr, CurveGroup};
 use ark_ff::{BigInt, BigInteger, FftField, Field, PrimeField};
 #[cfg(feature = "parallel")]
 use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
+use std::ops::Neg;
 
 pub fn fft_g1(vals: &[G1Affine], roots_of_unity: &[Fr]) -> Vec<G1Affine> {
     debug_assert_eq!(vals.len(), roots_of_unity.len());
@@ -48,8 +49,8 @@ pub fn fft_g1(vals: &[G1Affine], roots_of_unity: &[Fr]) -> Vec<G1Affine> {
         }
     };
     // Double sized so we can use iterator later on
-    let l: Vec<_> = l.into_iter().cycle().take(original_len).collect();
-    let r: Vec<_> = r.into_iter().cycle().take(original_len).collect();
+    let l = l.into_iter().cycle().take(original_len);
+    let r = r.into_iter().cycle().take(original_len);
 
     let y_times_root = {
         #[cfg(feature = "parallel")]
@@ -62,7 +63,7 @@ pub fn fft_g1(vals: &[G1Affine], roots_of_unity: &[Fr]) -> Vec<G1Affine> {
         }
     }
     .enumerate()
-    .map(|(i, y)| (y * roots_of_unity[i % vals.len()]).into_affine());
+    .map(|(i, y)| (y * roots_of_unity[i % vals.len()]));
 
     {
         #[cfg(feature = "parallel")]
@@ -80,7 +81,7 @@ pub fn fft_g1(vals: &[G1Affine], roots_of_unity: &[Fr]) -> Vec<G1Affine> {
         if i < vals.len() / 2 {
             x + y_times_root
         } else {
-            x - y_times_root
+            x + y_times_root.neg()
         }
         .into_affine()
     })
@@ -106,8 +107,90 @@ pub fn ifft_g1(vals: &[G1Affine], roots_of_unity: &[Fr]) -> Vec<G1Affine> {
     .collect()
 }
 
+pub fn fft_fr(vals: &[Fr], roots_of_unity: &[Fr]) -> Vec<Fr> {
+    debug_assert_eq!(vals.len(), roots_of_unity.len());
+    let original_len = vals.len();
+    if original_len == 1 {
+        return vals.to_vec();
+    }
+    let half_roots: Vec<_> = roots_of_unity.iter().step_by(2).copied().collect();
+
+    let l = || {
+        crate::fft::fft_fr(
+            vals.iter()
+                .step_by(2)
+                .copied()
+                .collect::<Vec<_>>()
+                .as_slice(),
+            half_roots.as_slice(),
+        )
+    };
+
+    let r = || {
+        crate::fft::fft_fr(
+            vals.iter()
+                .skip(1)
+                .step_by(2)
+                .copied()
+                .collect::<Vec<_>>()
+                .as_slice(),
+            half_roots.as_slice(),
+        )
+    };
+
+    let [l, r]: [Vec<Fr>; 2] = {
+        #[cfg(feature = "parallel")]
+        {
+            let (l, r) = rayon::join(l, r);
+            [l, r]
+        }
+        #[cfg(not(feature = "parallel"))]
+        {
+            [l(), r()]
+        }
+    };
+    // Double sized so we can use iterator later on
+    let l: Vec<_> = l.into_iter().cycle().take(original_len).collect();
+    let r: Vec<_> = r.into_iter().cycle().take(original_len).collect();
+
+    let y_times_root = {
+        #[cfg(feature = "parallel")]
+        {
+            r.into_par_iter()
+        }
+        #[cfg(not(feature = "parallel"))]
+        {
+            r.into_iter()
+        }
+    }
+    .enumerate()
+    .map(|(i, y)| y * roots_of_unity[i % vals.len()]);
+
+    {
+        #[cfg(feature = "parallel")]
+        {
+            l.into_par_iter()
+        }
+        #[cfg(not(feature = "parallel"))]
+        {
+            l.into_iter()
+        }
+    }
+    .zip(y_times_root)
+    .enumerate()
+    .map(|(i, (x, y_times_root))| {
+        if i < vals.len() / 2 {
+            x + y_times_root
+        } else {
+            x - y_times_root
+        }
+    })
+    .collect()
+}
+
 #[cfg(test)]
 mod test {
+    use crate::common::compute_roots_of_unity;
     use crate::fft::{fft_g1, ifft_g1};
     use ark_bls12_381::{Fr, G1Affine};
     use ark_ec::{AffineRepr, CurveGroup};
@@ -116,10 +199,8 @@ mod test {
     #[test]
     fn test_fft_ifft_g1() {
         for size in [16usize, 32, 64, 128, 256, 512, 1024, 2048, 4096] {
-            let primitive_root = <Fr as FftField>::get_root_of_unity(size as u64).unwrap();
-            let roots_of_unity: Vec<_> = (1..=size)
-                .map(|i| primitive_root.pow::<ark_ff::BigInt<4>>(BigInt::from(i as u64)))
-                .collect();
+            let roots_of_unity = compute_roots_of_unity(size);
+            let r: Vec<_> = roots_of_unity.iter().map(|a| a.to_string()).collect();
             let buff: Vec<G1Affine> = (0..size)
                 .map(|i| {
                     G1Affine::identity()

nomos-da/kzgrs/src/fk20.rs

@@ -1,20 +1,120 @@
-use crate::{GlobalParameters, Polynomial, Proof};
-use ark_bls12_381::{Bls12_381, Fq, Fr, G1Affine, G1Projective};
-use ark_ec::mnt4::G1Prepared;
+use ark_bls12_381::{Fr, G1Affine, G1Projective};
 use ark_ec::{AffineRepr, CurveGroup};
-
-use ark_poly::univariate::DensePolynomial;
-use ark_poly::{EvaluationDomain, Radix2EvaluationDomain};
+use ark_ff::{FftField, Field};
+use ark_poly::{EvaluationDomain, GeneralEvaluationDomain, Polynomial as _};
 use num_traits::Zero;
-fn toeplitz1(global_parameters: &GlobalParameters, polynomial_degree: usize) -> Vec<G1Projective> {
-    debug_assert_eq!(global_parameters.powers_of_g.len(), polynomial_degree);
+
+use crate::common::compute_roots_of_unity;
+use crate::fft::{fft_fr, fft_g1, ifft_g1};
+use crate::{GlobalParameters, Polynomial, Proof};
+
+fn toeplitz1(global_parameters: &[G1Affine], polynomial_degree: usize) -> Vec<G1Affine> {
+    debug_assert_eq!(global_parameters.len(), polynomial_degree);
     debug_assert!(polynomial_degree.is_power_of_two());
-    unimplemented!()
+    let roots_of_unity = compute_roots_of_unity(polynomial_degree * 2);
+    let vector_extended: Vec<G1Affine> = global_parameters
+        .iter()
+        .copied()
+        .chain(std::iter::repeat_with(G1Affine::zero).take(polynomial_degree))
+        .collect();
+    fft_g1(&vector_extended, &roots_of_unity)
+}
+
+fn toeplitz2(coefficients: &[Fr], extended_vector: &[G1Affine]) -> Vec<G1Affine> {
+    debug_assert!(coefficients.len().is_power_of_two());
+    // let domain: GeneralEvaluationDomain<Fr> =
+    //     GeneralEvaluationDomain::new(coefficients.len()).expect("Domain should be able to build");
+    // let toeplitz_coefficients_fft = domain.fft(coefficients);
+    let roots_of_unity = compute_roots_of_unity(coefficients.len());
+    let toeplitz_coefficients_fft = fft_fr(coefficients, &roots_of_unity);
+    extended_vector
+        .iter()
+        .copied()
+        .zip(toeplitz_coefficients_fft)
+        .map(|(v, c)| (v * c).into_affine())
+        .collect()
+}
+
+fn toeplitz3(h_extended_fft: &[G1Affine], polynomial_degree: usize) -> Vec<G1Affine> {
+    let roots_of_unity: Vec<Fr> = compute_roots_of_unity(h_extended_fft.len());
+    ifft_g1(h_extended_fft, &roots_of_unity)
+        .into_iter()
+        .take(polynomial_degree)
+        .collect()
 }
 
 pub fn fk20_batch_generate_elements_proofs(
-    polynomial: &DensePolynomial<Fr>,
+    polynomial: &Polynomial,
     global_parameters: &GlobalParameters,
 ) -> Vec<Proof> {
-    unimplemented!()
+    let polynomial_degree = polynomial.len();
+    debug_assert!(polynomial_degree <= global_parameters.powers_of_g.len());
+    debug_assert!(polynomial_degree.is_power_of_two());
+    let roots_of_unity: Vec<Fr> = compute_roots_of_unity(polynomial_degree);
+    let global_parameters: Vec<G1Affine> = global_parameters
+        .powers_of_g
+        .iter()
+        .copied()
+        .take(polynomial_degree)
+        .rev()
+        .collect();
+
+    let extended_vector = toeplitz1(&global_parameters, polynomial_degree);
+    let toeplitz_coefficients: Vec<Fr> = std::iter::repeat(Fr::ZERO)
+        .take(polynomial_degree)
+        .chain(polynomial.coeffs.iter().copied())
+        .collect();
+    let h_extended_vector = toeplitz2(&toeplitz_coefficients, &extended_vector);
+    let h_vector = toeplitz3(&h_extended_vector, polynomial_degree);
+    fft_g1(&h_vector, &roots_of_unity)
+        .into_iter()
+        .map(|g1| Proof {
+            w: g1,
+            random_v: None,
+        })
+        .collect()
+}
+
+#[cfg(test)]
+mod test {
+    use crate::common::compute_roots_of_unity;
+    use crate::fk20::fk20_batch_generate_elements_proofs;
+    use crate::{
+        common::bytes_to_polynomial, kzg::generate_element_proof, GlobalParameters, Proof,
+        BYTES_PER_FIELD_ELEMENT,
+    };
+    use ark_bls12_381::{Bls12_381, Fr};
+    use ark_ff::FftField;
+    use ark_poly::univariate::DensePolynomial;
+    use ark_poly::{EvaluationDomain, GeneralEvaluationDomain};
+    use ark_poly_commit::kzg10::KZG10;
+    use once_cell::sync::Lazy;
+    use rand::SeedableRng;
+
+    static GLOBAL_PARAMETERS: Lazy<GlobalParameters> = Lazy::new(|| {
+        let mut rng = rand::rngs::StdRng::seed_from_u64(1987);
+        KZG10::<Bls12_381, DensePolynomial<Fr>>::setup(4096, true, &mut rng).unwrap()
+    });
+
+    #[test]
+    fn test_generate_proofs() {
+        for size in [16, 32, 64, 128, 256] {
+            let mut buff: Vec<_> = (0..BYTES_PER_FIELD_ELEMENT * size)
+                .map(|i| (i % 255) as u8)
+                .rev()
+                .collect();
+            let domain = GeneralEvaluationDomain::new(size).unwrap();
+            let (evals, poly) =
+                bytes_to_polynomial::<BYTES_PER_FIELD_ELEMENT>(&buff, domain).unwrap();
+            let polynomial_degree = poly.len();
+            let roots_of_unity: Vec<Fr> = compute_roots_of_unity(size);
+            let slow_proofs: Vec<Proof> = (0..polynomial_degree)
+                .map(|i| {
+                    generate_element_proof(i, &poly, &evals, &GLOBAL_PARAMETERS, domain).unwrap()
+                })
+                .collect();
+            let fk20_proofs = fk20_batch_generate_elements_proofs(&poly, &GLOBAL_PARAMETERS);
+            assert_eq!(slow_proofs, fk20_proofs);
+        }
+    }
 }