Lagrange Form Bi kzg update (#130)

* tmp

* clean up

* fix clippy

* fix clippy

* add a bit more tests

* minor

---------

Co-authored-by: zhenfei <zhenfei.zhang@hotmail.com>

bi-kzg/src/bi_fft.rs

@@ -53,7 +53,9 @@ fn mul_assign_vec<F: Field>(a: &mut [F], b: &F) {
 //
 /// Performs a radix-$2$ Fast-Fourier Transformation (FFT) on a vector of size
 /// $n = 2^k$, when provided `log_n` = $k$ and an element of multiplicative
-/// order $n$ called `omega` ($\omega$). The result is that the vector `a`, when
+/// order $n$ called `omega` ($\omega$).
+///
+/// The result is that the vector `a`, when
 /// interpreted as the coefficients of a polynomial of degree $n - 1$, is
 /// transformed into the evaluations of this polynomial at each of the $n$
 /// distinct powers of $\omega$. This transformation is invertible by providing

bi-kzg/src/coeff_form_bi_kzg.rs

@@ -12,14 +12,14 @@ use halo2curves::CurveAffine;
 use itertools::Itertools;
 use rand::RngCore;
 
-use crate::poly::{lagrange_coefficients, univariate_quotient};
-use crate::structs::BivariateLagrangePolynomial;
-use crate::structs::BivariatePolynomial;
+use crate::poly::{
+    lagrange_coefficients, univariate_quotient, BivariateLagrangePolynomial, BivariatePolynomial,
+};
 use crate::util::parallelize;
 use crate::{
     pcs::PolynomialCommitmentScheme,
     util::{powers_of_field_elements, tensor_product_parallel},
-    BiKZGCommitment, BiKZGProof, BiKZGSRS, BiKZGVerifierParam,
+    BiKZGCommitment, BiKZGProof, BiKZGVerifierParam, CoefFormBiKZGSRS,
 };
 
 /// Commit to the bi-variate polynomial in its coefficient form.
@@ -32,8 +32,8 @@ impl<E: MultiMillerLoop> PolynomialCommitmentScheme for CoeffFormBiKZG<E>
 where
     E::G1Affine: CurveAffine<ScalarExt = E::Fr, CurveExt = E::G1>,
 {
-    type SRS = BiKZGSRS<E>;
-    type ProverParam = BiKZGSRS<E>;
+    type SRS = CoefFormBiKZGSRS<E>;
+    type ProverParam = CoefFormBiKZGSRS<E>;
     type VerifierParam = BiKZGVerifierParam<E>;
     type Polynomial = BivariatePolynomial<E::Fr>;
     type Commitment = BiKZGCommitment<E>;
@@ -123,7 +123,7 @@ where
             affine_bases
         };
 
-        BiKZGSRS {
+        CoefFormBiKZGSRS {
             powers_of_g: coeff_bases,
             powers_of_g_lagrange_over_both_roots: lagrange_bases,
             h: E::G2Affine::generator(),

bi-kzg/src/lagrange_form_bi_kzg.rs

@@ -1,10 +1,7 @@
-//! We don't need this file for now. We will use the `CoeffFormBiKZG`.
-
 use std::{borrow::Borrow, marker::PhantomData};
 
 use ark_std::{end_timer, start_timer};
 use halo2curves::ff::Field;
-use halo2curves::ff::PrimeField;
 use halo2curves::group::prime::PrimeCurveAffine;
 use halo2curves::group::Curve;
 use halo2curves::group::Group;
@@ -14,13 +11,13 @@ use halo2curves::CurveAffine;
 use itertools::Itertools;
 use rand::RngCore;
 
-use crate::poly::{lagrange_coefficients, univariate_quotient};
-use crate::structs::BivariateLagrangePolynomial;
-use crate::util::parallelize;
+use crate::parallelize;
+use crate::poly::{lagrange_coefficients, BivariateLagrangePolynomial};
+use crate::primitive_root_of_unity;
 use crate::{
     pcs::PolynomialCommitmentScheme,
     util::{powers_of_field_elements, tensor_product_parallel},
-    BiKZGCommitment, BiKZGProof, BiKZGSRS, BiKZGVerifierParam,
+    BiKZGCommitment, BiKZGProof, BiKZGVerifierParam, LagrangeFormBiKZGSRS,
 };
 
 /// Commit to the bi-variate polynomial in its lagrange form.
@@ -33,8 +30,8 @@ impl<E: MultiMillerLoop> PolynomialCommitmentScheme for LagrangeFormBiKZG<E>
 where
     E::G1Affine: CurveAffine<ScalarExt = E::Fr, CurveExt = E::G1>,
 {
-    type SRS = BiKZGSRS<E>;
-    type ProverParam = BiKZGSRS<E>;
+    type SRS = LagrangeFormBiKZGSRS<E>;
+    type ProverParam = LagrangeFormBiKZGSRS<E>;
     type VerifierParam = BiKZGVerifierParam<E>;
     type Polynomial = BivariateLagrangePolynomial<E::Fr>;
     type Commitment = BiKZGCommitment<E>;
@@ -53,52 +50,37 @@ where
 
         let tau_0 = E::Fr::random(&mut rng);
         let tau_1 = E::Fr::random(&mut rng);
+
         let g1 = E::G1Affine::generator();
 
         // roots of unity for supported_n and supported_m
-        let (omega_0, omega_1) = {
-            let omega = E::Fr::ROOT_OF_UNITY;
-            let omega_0 = omega.pow_vartime(&[(1 << E::Fr::S) / supported_n as u64]);
-            let omega_1 = omega.pow_vartime(&[(1 << E::Fr::S) / supported_m as u64]);
-
-            assert!(
-                omega_0.pow_vartime(&[supported_n as u64]) == E::Fr::ONE,
-                "omega_0 is not root of unity for supported_n"
-            );
-            assert!(
-                omega_1.pow_vartime(&[supported_m as u64]) == E::Fr::ONE,
-                "omega_1 is not root of unity for supported_m"
-            );
-            (omega_0, omega_1)
-        };
+        let omega_0 = primitive_root_of_unity(supported_n);
+        let omega_1 = primitive_root_of_unity(supported_m);
 
         // computes the vector of L_i^N(tau_0) * L_j^M(tau_1) for i in 0..supported_n and j in 0..supported_m
-        let (scalars, lagrange_scalars) = {
+        let (lagrange_tau_0, lagrange_scalars) = {
             let powers_of_omega_0 = powers_of_field_elements(&omega_0, supported_n);
-            let powers_of_tau_0 = powers_of_field_elements(&tau_0, supported_n);
             let lagrange_tau_0 = lagrange_coefficients(&powers_of_omega_0, &tau_0);
             let powers_of_omega_1 = powers_of_field_elements(&omega_1, supported_m);
-            let powers_of_tau_1 = powers_of_field_elements(&tau_1, supported_m);
             let lagrange_tau_1 = lagrange_coefficients(&powers_of_omega_1, &tau_1);
-            let scalars = tensor_product_parallel(&powers_of_tau_0, &powers_of_tau_1);
             let lagrange_scalars = tensor_product_parallel(&lagrange_tau_0, &lagrange_tau_1);
 
-            (scalars, lagrange_scalars)
+            (lagrange_tau_0, lagrange_scalars)
         };
 
         let g1_prog = g1.to_curve();
-        let coeff_bases = {
-            let mut proj_bases = vec![E::G1::identity(); supported_n * supported_m];
+        let lagrange_x_bases = {
+            let mut proj_bases = vec![E::G1::identity(); supported_n];
             parallelize(&mut proj_bases, |g, start| {
                 for (idx, g) in g.iter_mut().enumerate() {
                     let offset = start + idx;
-                    *g = g1_prog * scalars[offset];
+                    *g = g1_prog * lagrange_tau_0[offset];
                 }
             });
 
-            let mut g_bases = vec![E::G1Affine::identity(); supported_n * supported_m];
-            parallelize(&mut g_bases, |g, starts| {
-                E::G1::batch_normalize(&proj_bases[starts..(starts + g.len())], g);
+            let mut g_bases = vec![E::G1Affine::identity(); supported_n];
+            parallelize(&mut g_bases, |g, start| {
+                E::G1::batch_normalize(&proj_bases[start..start + g.len()], g);
             });
             drop(proj_bases);
             g_bases
@@ -114,18 +96,16 @@ where
             });
 
             let mut affine_bases = vec![E::G1Affine::identity(); supported_n * supported_m];
-            parallelize(&mut affine_bases, |affine_bases, starts| {
-                E::G1::batch_normalize(
-                    &proj_bases[starts..(starts + affine_bases.len())],
-                    affine_bases,
-                );
+            parallelize(&mut affine_bases, |g, start| {
+                E::G1::batch_normalize(&proj_bases[start..start + g.len()], g);
             });
             drop(proj_bases);
             affine_bases
         };
 
-        BiKZGSRS {
-            powers_of_g: coeff_bases,
+        LagrangeFormBiKZGSRS {
+            g: g1,
+            powers_of_g_lagrange_over_x: lagrange_x_bases,
             powers_of_g_lagrange_over_both_roots: lagrange_bases,
             h: E::G2Affine::generator(),
             tau_0_h: (E::G2Affine::generator() * tau_0).into(),
@@ -179,55 +159,52 @@ where
 
         let timer2 = start_timer!(|| "Computing the proof pi0");
         let (pi_0, f_x_b) = {
-            let f_x_b = polynomial.evaluate_y(&point.1);
-            let mut q_0_x_b = f_x_b.clone();
-            q_0_x_b[0] -= u;
-            let q_0_x_b = univariate_quotient(&q_0_x_b, &point.0);
+            let f_x_b = polynomial.evaluate_at_y(&point.1);
+
+            let omega_0 = primitive_root_of_unity(polynomial.degree_0);
+            let powers_of_omega_0 =
+                powers_of_field_elements::<E::Fr>(&omega_0, polynomial.degree_0);
+            // todo use batch inversion
+            let powers_of_omega_0_minus_x_inv = powers_of_omega_0
+                .iter()
+                .map(|w| (*w - point.0).invert().unwrap())
+                .collect::<Vec<_>>();
+
+            let q_0_x_b = f_x_b
+                .iter()
+                .zip(powers_of_omega_0_minus_x_inv)
+                .map(|(v0, v1)| (*v0 - u) * v1)
+                .collect::<Vec<_>>();
 
             let pi_0 = best_multiexp(
                 &q_0_x_b,
-                prover_param.borrow().powers_of_g[..polynomial.degree_0].as_ref(),
+                prover_param.borrow().powers_of_g_lagrange_over_x.as_ref(),
             )
             .to_affine();
             (pi_0, f_x_b)
         };
         end_timer!(timer2);
 
+        // f(X, Y) = qx(X)(X - a) + qy(X, Y)(Y - b) + u
         let timer2 = start_timer!(|| "Computing the proof pi1");
         let pi_1 = {
-            let mut t = polynomial.clone();
-            t.coefficients
-                .iter_mut()
-                .take(polynomial.degree_0)
-                .zip_eq(f_x_b.iter())
-                .for_each(|(c, f)| *c -= f);
-            let coeffs = t.lagrange_coeffs();
-
-            let mut divisor = vec![E::Fr::from(0); polynomial.degree_0 * polynomial.degree_1];
-            divisor[0] = -point.1;
-            divisor[polynomial.degree_0] = E::Fr::ONE;
-            let divisor =
-                BivariatePolynomial::new(divisor, polynomial.degree_0, polynomial.degree_1);
-
-            let divisor = divisor.lagrange_coeffs();
-
-            // todo: batch invert
-            let y_minus_a_inv_lag = divisor
-                .iter()
-                .map(|o| {
-                    if o.is_zero_vartime() {
-                        panic!("not invertible")
-                    } else {
-                        o.invert().unwrap()
-                    }
+            let omega_1 = primitive_root_of_unity(polynomial.degree_1);
+            let powers_of_omega_1 =
+                powers_of_field_elements::<E::Fr>(&omega_1, polynomial.degree_1);
+
+            // todo use batch inversion
+            let q_1_x_y = polynomial
+                .coefficients
+                .chunks_exact(polynomial.degree_0)
+                .zip_eq(powers_of_omega_1)
+                .flat_map(|(coeffs_i, w_y_i)| {
+                    coeffs_i
+                        .iter()
+                        .zip(f_x_b.iter())
+                        .map(|(coeff, v)| (*coeff - v) * (w_y_i - point.1).invert().unwrap())
+                        .collect::<Vec<E::Fr>>()
                 })
-                .collect::<Vec<_>>();
-
-            let q_1_x_y = coeffs
-                .iter()
-                .zip_eq(y_minus_a_inv_lag.iter())
-                .map(|(c, y)| (*c) * *y)
-                .collect::<Vec<_>>();
+                .collect::<Vec<E::Fr>>();
 
             best_multiexp(
                 &q_1_x_y,
@@ -260,24 +237,17 @@ where
     {
         let pi0_a_pi1_b_g1_cmu = best_multiexp(
             &[point.0, point.1, E::Fr::ONE, -*value],
-            &[
-                proof.pi0,
-                proof.pi1,
-                commitment.com.into(),
-                verifier_param.g.into(),
-            ],
+            &[proof.pi0, proof.pi1, commitment.com, verifier_param.g],
         );
         let pi0_a_pi1_b_g1_cmu = (-pi0_a_pi1_b_g1_cmu).to_affine();
         let res = E::multi_miller_loop(&[
             (&proof.pi0, &verifier_param.tau_0_h.into()),
             (&proof.pi1, &verifier_param.tau_1_h.into()),
             (&pi0_a_pi1_b_g1_cmu, &verifier_param.h.into()),
         ]);
-        let res = res.final_exponentiation().is_identity().into();
 
-        res
+        res.final_exponentiation().is_identity().into()
     }
 
-
     // TODO: implement multi-opening and batch verification
 }

bi-kzg/src/lib.rs

@@ -1,18 +1,23 @@
 mod bi_fft;
+pub use bi_fft::*;
+
 mod coeff_form_bi_kzg;
+pub use coeff_form_bi_kzg::*;
+
+mod lagrange_form_bi_kzg;
+pub use lagrange_form_bi_kzg::*;
+
 mod pcs;
+pub use pcs::*;
+
 mod poly;
+pub use poly::*;
+
 mod structs;
-mod util;
+pub use structs::*;
 
-// mod lagrange_form_bi_kzg;
+mod util;
+pub use util::*;
 
 #[cfg(test)]
 mod tests;
-
-pub use coeff_form_bi_kzg::CoeffFormBiKZG;
-pub use pcs::PolynomialCommitmentScheme;
-pub use structs::BivariatePolynomial;
-pub use structs::{BiKZGCommitment, BiKZGProof, BiKZGSRS, BiKZGVerifierParam};
-
-// pub use lagrange_form_bi_kzg::LagrangeFormBiKZG;