Merge pull request #70 from mir-protocol/optimize_polynomials

Optimize some polynomial operations

src/field/interpolation.rs

@@ -65,11 +65,23 @@ pub fn barycentric_weights<F: Field>(points: &[(F, F)]) -> Vec<F> {
     )
 }
 
+/// Interpolate the linear polynomial passing through `points` on `x`.
+pub fn interpolate2<F: Field>(points: [(F, F); 2], x: F) -> F {
+    // a0 -> a1
+    // b0 -> b1
+    // x  -> a1 + (x-a0)*(b1-a1)/(b0-a0)
+    let (a0, a1) = points[0];
+    let (b0, b1) = points[1];
+    assert_ne!(a0, b0);
+    a1 + (x - a0) * (b1 - a1) / (b0 - a0)
+}
+
 #[cfg(test)]
 mod tests {
+    use super::*;
     use crate::field::crandall_field::CrandallField;
+    use crate::field::extension_field::quartic::QuarticCrandallField;
     use crate::field::field::Field;
-    use crate::field::lagrange::interpolant;
     use crate::polynomial::polynomial::PolynomialCoeffs;
 
     #[test]
@@ -120,4 +132,18 @@ mod tests {
     fn eval_naive<F: Field>(coeffs: &PolynomialCoeffs<F>, domain: &[F]) -> Vec<(F, F)> {
         domain.iter().map(|&x| (x, coeffs.eval(x))).collect()
     }
+
+    #[test]
+    fn test_interpolate2() {
+        type F = QuarticCrandallField;
+        let points = [(F::rand(), F::rand()), (F::rand(), F::rand())];
+        let x = F::rand();
+
+        let ev0 = interpolant(&points).eval(x);
+        let ev1 = interpolate(&points, x, &barycentric_weights(&points));
+        let ev2 = interpolate2(points, x);
+
+        assert_eq!(ev0, ev1);
+        assert_eq!(ev0, ev2);
+    }
 }

src/field/mod.rs

@@ -3,7 +3,7 @@ pub mod crandall_field;
 pub mod extension_field;
 pub mod fft;
 pub mod field;
-pub(crate) mod lagrange;
+pub(crate) mod interpolation;
 
 #[cfg(test)]
 mod field_testing;

src/fri/verifier.rs

@@ -2,7 +2,7 @@ use anyhow::{ensure, Result};
 
 use crate::field::extension_field::{flatten, Extendable, FieldExtension, Frobenius};
 use crate::field::field::Field;
-use crate::field::lagrange::{barycentric_weights, interpolant, interpolate};
+use crate::field::interpolation::{barycentric_weights, interpolate, interpolate2};
 use crate::fri::FriConfig;
 use crate::hash::hash_n_to_1;
 use crate::merkle_proofs::verify_merkle_proof;
@@ -185,14 +185,17 @@ fn fri_combine_initial<F: Field + Extendable<D>, const D: usize>(
         .map(|&e| F::Extension::from_basefield(e));
     let zs_composition_eval = reduce_with_iter(zs_evals, alpha_powers.clone());
     let zeta_right = F::Extension::primitive_root_of_unity(degree_log) * zeta;
-    let zs_interpol = interpolant(&[
-        (zeta, reduce_with_iter(&os.plonk_zs, alpha_powers.clone())),
-        (
-            zeta_right,
-            reduce_with_iter(&os.plonk_zs_right, &mut alpha_powers),
-        ),
-    ]);
-    let zs_numerator = zs_composition_eval - zs_interpol.eval(subgroup_x);
+    let zs_interpol = interpolate2(
+        [
+            (zeta, reduce_with_iter(&os.plonk_zs, alpha_powers.clone())),
+            (
+                zeta_right,
+                reduce_with_iter(&os.plonk_zs_right, &mut alpha_powers),
+            ),
+        ],
+        subgroup_x,
+    );
+    let zs_numerator = zs_composition_eval - zs_interpol;
     let zs_denominator = (subgroup_x - zeta) * (subgroup_x - zeta_right);
     sum += zs_numerator / zs_denominator;
 
@@ -211,8 +214,8 @@ fn fri_combine_initial<F: Field + Extendable<D>, const D: usize>(
     // and one call at the end of the sum.
     let alpha_powers_frob = alpha_powers.repeated_frobenius(D - 1);
     let wire_eval_frob = reduce_with_iter(&os.wires, alpha_powers_frob).frobenius();
-    let wire_interpol = interpolant(&[(zeta, wire_eval), (zeta_frob, wire_eval_frob)]);
-    let wire_numerator = wire_composition_eval - wire_interpol.eval(subgroup_x);
+    let wire_interpol = interpolate2([(zeta, wire_eval), (zeta_frob, wire_eval_frob)], subgroup_x);
+    let wire_numerator = wire_composition_eval - wire_interpol;
     let wire_denominator = (subgroup_x - zeta) * (subgroup_x - zeta_frob);
     sum += wire_numerator / wire_denominator;
 

src/gadgets/interpolation.rs

@@ -1,9 +1,10 @@
+use std::marker::PhantomData;
+
 use crate::circuit_builder::CircuitBuilder;
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::Extendable;
 use crate::gates::interpolation::InterpolationGate;
 use crate::target::Target;
-use std::marker::PhantomData;
 
 impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     /// Interpolate two points. No need for an `InterpolationGate` since the coefficients
@@ -56,15 +57,16 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
 #[cfg(test)]
 mod tests {
+    use std::convert::TryInto;
+
     use super::*;
     use crate::circuit_data::CircuitConfig;
     use crate::field::crandall_field::CrandallField;
     use crate::field::extension_field::quartic::QuarticCrandallField;
     use crate::field::extension_field::FieldExtension;
     use crate::field::field::Field;
-    use crate::field::lagrange::{interpolant, interpolate};
+    use crate::field::interpolation::{interpolant, interpolate};
     use crate::witness::PartialWitness;
-    use std::convert::TryInto;
 
     #[test]
     fn test_interpolate() {

src/gates/interpolation.rs

@@ -6,7 +6,7 @@ use crate::circuit_builder::CircuitBuilder;
 use crate::field::extension_field::algebra::PolynomialCoeffsAlgebra;
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::{Extendable, FieldExtension};
-use crate::field::lagrange::interpolant;
+use crate::field::interpolation::interpolant;
 use crate::gadgets::polynomial::PolynomialCoeffsExtAlgebraTarget;
 use crate::gates::gate::{Gate, GateRef};
 use crate::generator::{SimpleGenerator, WitnessGenerator};

src/polynomial/commitment.rs

@@ -4,11 +4,10 @@ use rayon::prelude::*;
 use crate::field::extension_field::Extendable;
 use crate::field::extension_field::{FieldExtension, Frobenius};
 use crate::field::field::Field;
-use crate::field::lagrange::interpolant;
 use crate::fri::{prover::fri_proof, verifier::verify_fri_proof, FriConfig};
 use crate::merkle_tree::MerkleTree;
 use crate::plonk_challenger::Challenger;
-use crate::plonk_common::{reduce_polys_with_iter, reduce_with_iter};
+use crate::plonk_common::reduce_polys_with_iter;
 use crate::polynomial::polynomial::PolynomialCoeffs;
 use crate::proof::{FriProof, FriProofTarget, Hash, OpeningSet};
 use crate::timed;
@@ -120,50 +119,27 @@ impl<F: Field> ListPolynomialCommitment<F> {
             .iter()
             .flat_map(|&i| &commitments[i].polynomials)
             .map(|p| p.to_extension());
-        let single_os = [&os.constants, &os.plonk_s_sigmas, &os.quotient_polys];
-        let single_evals = single_os.iter().flat_map(|v| v.iter());
-        let single_composition_poly = reduce_polys_with_iter(single_polys, alpha_powers.clone());
-        let single_composition_eval = reduce_with_iter(single_evals, &mut alpha_powers);
+        let single_composition_poly = reduce_polys_with_iter(single_polys, &mut alpha_powers);
 
-        let single_quotient = Self::compute_quotient(
-            &[zeta],
-            &[single_composition_eval],
-            &single_composition_poly,
-        );
-        final_poly = &final_poly + &single_quotient;
+        let single_quotient = Self::compute_quotient([zeta], single_composition_poly);
+        final_poly += single_quotient;
 
         // Zs polynomials are opened at `zeta` and `g*zeta`.
         let zs_polys = commitments[3].polynomials.iter().map(|p| p.to_extension());
-        let zs_composition_poly = reduce_polys_with_iter(zs_polys, alpha_powers.clone());
-        let zs_composition_evals = [
-            reduce_with_iter(&os.plonk_zs, alpha_powers.clone()),
-            reduce_with_iter(&os.plonk_zs_right, &mut alpha_powers),
-        ];
+        let zs_composition_poly = reduce_polys_with_iter(zs_polys, &mut alpha_powers);
 
-        let zs_quotient = Self::compute_quotient(
-            &[zeta, g * zeta],
-            &zs_composition_evals,
-            &zs_composition_poly,
-        );
-        final_poly = &final_poly + &zs_quotient;
+        let zs_quotient = Self::compute_quotient([zeta, g * zeta], zs_composition_poly);
+        final_poly += zs_quotient;
 
         // When working in an extension field, need to check that wires are in the base field.
         // Check this by opening the wires polynomials at `zeta` and `zeta.frobenius()` and using the fact that
         // a polynomial `f` is over the base field iff `f(z).frobenius()=f(z.frobenius())` with high probability.
         let wire_polys = commitments[2].polynomials.iter().map(|p| p.to_extension());
-        let wire_composition_poly = reduce_polys_with_iter(wire_polys, alpha_powers.clone());
-        let wire_evals_frob = os.wires.iter().map(|e| e.frobenius()).collect::<Vec<_>>();
-        let wire_composition_evals = [
-            reduce_with_iter(&os.wires, alpha_powers.clone()),
-            reduce_with_iter(&wire_evals_frob, alpha_powers),
-        ];
+        let wire_composition_poly = reduce_polys_with_iter(wire_polys, &mut alpha_powers);
 
-        let wires_quotient = Self::compute_quotient(
-            &[zeta, zeta.frobenius()],
-            &wire_composition_evals,
-            &wire_composition_poly,
-        );
-        final_poly = &final_poly + &wires_quotient;
+        let wires_quotient =
+            Self::compute_quotient([zeta, zeta.frobenius()], wire_composition_poly);
+        final_poly += wires_quotient;
 
         let lde_final_poly = final_poly.lde(config.rate_bits);
         let lde_final_values = lde_final_poly
@@ -194,28 +170,27 @@ impl<F: Field> ListPolynomialCommitment<F> {
 
     /// Given `points=(x_i)`, `evals=(y_i)` and `poly=P` with `P(x_i)=y_i`, computes the polynomial
     /// `Q=(P-I)/Z` where `I` interpolates `(x_i, y_i)` and `Z` is the vanishing polynomial on `(x_i)`.
-    fn compute_quotient<const D: usize>(
-        points: &[F::Extension],
-        evals: &[F::Extension],
-        poly: &PolynomialCoeffs<F::Extension>,
+    fn compute_quotient<const D: usize, const N: usize>(
+        points: [F::Extension; N],
+        poly: PolynomialCoeffs<F::Extension>,
     ) -> PolynomialCoeffs<F::Extension>
     where
         F: Extendable<D>,
     {
-        let pairs = points
-            .iter()
-            .zip(evals)
-            .map(|(&x, &e)| (x, e))
-            .collect::<Vec<_>>();
-        debug_assert!(pairs.iter().all(|&(x, e)| poly.eval(x) == e));
-
-        let interpolant = interpolant(&pairs);
-        let denominator = points.iter().fold(PolynomialCoeffs::one(), |acc, &x| {
-            &acc * &PolynomialCoeffs::new(vec![-x, F::Extension::ONE])
-        });
-        let numerator = poly - &interpolant;
-        let (quotient, rem) = numerator.div_rem(&denominator);
-        debug_assert!(rem.is_zero());
+        let quotient = if N == 1 {
+            poly.divide_by_linear(points[0]).0
+        } else if N == 2 {
+            // The denominator is `(X - p0)(X - p1) = p0 p1 - (p0 + p1) X + X^2`.
+            let denominator = vec![
+                points[0] * points[1],
+                -points[0] - points[1],
+                F::Extension::ONE,
+            ]
+            .into();
+            poly.div_rem_long_division(&denominator).0 // Could also use `divide_by_linear` twice.
+        } else {
+            unreachable!("This shouldn't happen. Plonk should open polynomials at 1 or 2 points.")
+        };
 
         quotient.padded(quotient.degree_plus_one().next_power_of_two())
     }

src/polynomial/division.rs

@@ -26,7 +26,7 @@ impl<F: Field> PolynomialCoeffs<F> {
                 .to_vec()
                 .into();
             let mut q = rev_q.rev();
-            let mut qb = &q * b;
+            let qb = &q * b;
             let mut r = self - &qb;
             q.trim();
             r.trim();
@@ -59,8 +59,7 @@ impl<F: Field> PolynomialCoeffs<F> {
                 quotient.coeffs[cur_q_degree] = cur_q_coeff;
 
                 for (i, &div_coeff) in b.coeffs.iter().enumerate() {
-                    remainder.coeffs[cur_q_degree + i] =
-                        remainder.coeffs[cur_q_degree + i] - (cur_q_coeff * div_coeff);
+                    remainder.coeffs[cur_q_degree + i] -= cur_q_coeff * div_coeff;
                 }
                 remainder.trim();
             }
@@ -97,7 +96,7 @@ impl<F: Field> PolynomialCoeffs<F> {
         let denominators = (0..a_eval.len())
             .map(|i| {
                 if i != 0 {
-                    root_pow = root_pow * root_n;
+                    root_pow *= root_n;
                 }
                 denominator_g * root_pow - F::ONE
             })
@@ -125,8 +124,25 @@ impl<F: Field> PolynomialCoeffs<F> {
         p
     }
 
+    /// Let `self=p(X)`, this returns `(p(X)-p(z))/(X-z)` and `p(z)`.
+    /// See https://en.wikipedia.org/wiki/Horner%27s_method
+    pub(crate) fn divide_by_linear(&self, z: F) -> (PolynomialCoeffs<F>, F) {
+        let mut bs = self
+            .coeffs
+            .iter()
+            .rev()
+            .scan(F::ZERO, |acc, &c| {
+                *acc = *acc * z + c;
+                Some(*acc)
+            })
+            .collect::<Vec<_>>();
+        let ev = bs.pop().unwrap_or(F::ZERO);
+        bs.reverse();
+        (Self { coeffs: bs }, ev)
+    }
+
     /// Computes the inverse of `self` modulo `x^n`.
-    pub(crate) fn inv_mod_xn(&self, n: usize) -> Self {
+    pub fn inv_mod_xn(&self, n: usize) -> Self {
         assert!(self.coeffs[0].is_nonzero(), "Inverse doesn't exist.");
 
         let h = if self.len() < n {
@@ -166,7 +182,10 @@ impl<F: Field> PolynomialCoeffs<F> {
 
 #[cfg(test)]
 mod tests {
+    use std::time::Instant;
+
     use crate::field::crandall_field::CrandallField;
+    use crate::field::extension_field::quartic::QuarticCrandallField;
     use crate::field::field::Field;
     use crate::polynomial::polynomial::PolynomialCoeffs;
 
@@ -199,4 +218,49 @@ mod tests {
         let computed_q = a.divide_by_z_h(4);
         assert_eq!(computed_q, q);
     }
+
+    #[test]
+    #[ignore]
+    fn test_division_by_linear() {
+        type F = QuarticCrandallField;
+        let n = 1_000_000;
+        let poly = PolynomialCoeffs::new(F::rand_vec(n));
+        let z = F::rand();
+        let ev = poly.eval(z);
+
+        let timer = Instant::now();
+        let (quotient, ev2) = poly.div_rem(&PolynomialCoeffs::new(vec![-z, F::ONE]));
+        println!("{:.3}s for usual", timer.elapsed().as_secs_f32());
+        assert_eq!(ev2.trimmed().coeffs, vec![ev]);
+
+        let timer = Instant::now();
+        let (quotient, ev3) = poly.div_rem_long_division(&PolynomialCoeffs::new(vec![-z, F::ONE]));
+        println!("{:.3}s for long division", timer.elapsed().as_secs_f32());
+        assert_eq!(ev3.trimmed().coeffs, vec![ev]);
+
+        let timer = Instant::now();
+        let horn = poly.divide_by_linear(z);
+        println!("{:.3}s for Horner", timer.elapsed().as_secs_f32());
+        assert_eq!((quotient, ev), horn);
+    }
+
+    #[test]
+    #[ignore]
+    fn test_division_by_quadratic() {
+        type F = QuarticCrandallField;
+        let n = 1_000_000;
+        let poly = PolynomialCoeffs::new(F::rand_vec(n));
+        let quad = PolynomialCoeffs::new(F::rand_vec(2));
+
+        let timer = Instant::now();
+        let (quotient0, rem0) = poly.div_rem(&quad);
+        println!("{:.3}s for usual", timer.elapsed().as_secs_f32());
+
+        let timer = Instant::now();
+        let (quotient1, rem1) = poly.div_rem_long_division(&quad);
+        println!("{:.3}s for long division", timer.elapsed().as_secs_f32());
+
+        assert_eq!(quotient0.trimmed(), quotient1.trimmed());
+        assert_eq!(rem0.trimmed(), rem1.trimmed());
+    }
 }

src/polynomial/polynomial.rs

@@ -1,6 +1,6 @@
 use std::cmp::max;
 use std::iter::Sum;
-use std::ops::{Add, Mul, Sub};
+use std::ops::{Add, AddAssign, Mul, Sub, SubAssign};
 
 use crate::field::extension_field::Extendable;
 use crate::field::fft::{fft, ifft};
@@ -243,6 +243,26 @@ impl<F: Field> Sub for &PolynomialCoeffs<F> {
     }
 }
 
+impl<F: Field> AddAssign for PolynomialCoeffs<F> {
+    fn add_assign(&mut self, rhs: Self) {
+        let len = max(self.len(), rhs.len());
+        self.coeffs.resize(len, F::ZERO);
+        for (l, r) in self.coeffs.iter_mut().zip(rhs.coeffs) {
+            *l += r;
+        }
+    }
+}
+
+impl<F: Field> SubAssign for PolynomialCoeffs<F> {
+    fn sub_assign(&mut self, rhs: Self) {
+        let len = max(self.len(), rhs.len());
+        self.coeffs.resize(len, F::ZERO);
+        for (l, r) in self.coeffs.iter_mut().zip(rhs.coeffs) {
+            *l -= r;
+        }
+    }
+}
+
 impl<F: Field> Mul<F> for &PolynomialCoeffs<F> {
     type Output = PolynomialCoeffs<F>;