Optimize evaluation of Z_H on coset.

src/field/field.rs

@@ -104,10 +104,7 @@ pub trait Field:
     fn primitive_root_of_unity(n_log: usize) -> Self {
         assert!(n_log <= Self::TWO_ADICITY);
         let mut base = Self::POWER_OF_TWO_GENERATOR;
-        for _ in n_log..Self::TWO_ADICITY {
-            base = base.square();
-        }
-        base
+        base.exp_power_of_2(Self::TWO_ADICITY - n_log)
     }
 
     /// Computes a multiplicative subgroup whose order is known in advance.
@@ -158,6 +155,14 @@ pub trait Field:
         bits_u64(self.to_canonical_u64())
     }
 
+    fn exp_power_of_2(&self, power_log: usize) -> Self {
+        let mut res = *self;
+        for _ in 0..power_log {
+            res = res.square();
+        }
+        res
+    }
+
     fn exp(&self, power: u64) -> Self {
         let mut current = *self;
         let mut product = Self::ONE;

src/gates/gate.rs

@@ -1,6 +1,3 @@
-use std::borrow::Borrow;
-use std::collections::HashMap;
-use std::fmt::{Debug, Error, Formatter};
 use std::hash::{Hash, Hasher};
 use std::iter::FromIterator;
 use std::ops::Index;
@@ -10,7 +7,6 @@ use crate::circuit_builder::CircuitBuilder;
 use crate::field::extension_field::target::ExtensionTarget;
 use crate::field::extension_field::{Extendable, FieldExtension};
 use crate::field::field::Field;
-use crate::gates::gate_tree::Tree;
 use crate::generator::WitnessGenerator;
 use crate::vars::{EvaluationTargets, EvaluationVars, EvaluationVarsBase};
 

src/plonk_common.rs

@@ -116,6 +116,7 @@ pub(crate) fn eval_vanishing_poly<F: Extendable<D>, const D: usize>(
 /// Like `eval_vanishing_poly`, but specialized for base field points.
 pub(crate) fn eval_vanishing_poly_base<F: Extendable<D>, const D: usize>(
     common_data: &CommonCircuitData<F, D>,
+    i: usize,
     x: F,
     vars: EvaluationVarsBase<F>,
     local_plonk_zs: &[F],
@@ -124,6 +125,7 @@ pub(crate) fn eval_vanishing_poly_base<F: Extendable<D>, const D: usize>(
     betas: &[F],
     gammas: &[F],
     alphas: &[F],
+    z_h_on_coset: &ZeroPolyOnCoset<F>,
 ) -> Vec<F> {
     let constraint_terms =
         evaluate_gate_constraints_base(&common_data.gates, common_data.num_gate_constraints, vars);
@@ -136,7 +138,7 @@ pub(crate) fn eval_vanishing_poly_base<F: Extendable<D>, const D: usize>(
     for i in 0..common_data.config.num_challenges {
         let z_x = local_plonk_zs[i];
         let z_gz = next_plonk_zs[i];
-        vanishing_z_1_terms.push(eval_l_1(common_data.degree(), x) * (z_x - F::ONE));
+        vanishing_z_1_terms.push(z_h_on_coset.eval_l1(i, x) * (z_x - F::ONE));
 
         let mut f_prime = F::ONE;
         let mut g_prime = F::ONE;
@@ -226,6 +228,51 @@ pub(crate) fn eval_zero_poly<F: Field>(n: usize, x: F) -> F {
     x.exp(n as u64) - F::ONE
 }
 
+/// Precomputations of the evaluation of `Z_H(X) = X^n - 1` on a coset `gK` with `H <= K`.
+pub(crate) struct ZeroPolyOnCoset<F: Field> {
+    /// `n = |H|`.
+    n: F,
+    /// `rate = |K|/|H|`.
+    rate: usize,
+    /// Holds `g^n * (w^n)^i - 1 = g^n * v^i - 1` for `i in 0..rate`, with `w` a generator of `K` and `v` a
+    /// `rate`-primitive root of unity.
+    evals: Vec<F>,
+    /// Holds the multiplicative inverses of `evals`.
+    inverses: Vec<F>,
+}
+impl<F: Field> ZeroPolyOnCoset<F> {
+    pub fn new(n_log: usize, rate_bits: usize) -> Self {
+        let g_pow_n = F::coset_shift().exp_power_of_2(n_log);
+        let evals = F::two_adic_subgroup(rate_bits)
+            .into_iter()
+            .map(|x| g_pow_n * x - F::ONE)
+            .collect::<Vec<_>>();
+        let inverses = F::batch_multiplicative_inverse(&evals);
+        Self {
+            n: F::from_canonical_usize(1 << n_log),
+            rate: 1 << rate_bits,
+            evals,
+            inverses,
+        }
+    }
+
+    /// Returns `Z_H(g * w^i)`.
+    pub fn eval(&self, i: usize) -> F {
+        self.evals[i % self.rate]
+    }
+
+    /// Returns `Z_H(g * w^i)`.
+    pub fn eval_inverse(&self, i: usize) -> F {
+        self.inverses[i % self.rate]
+    }
+
+    /// Returns `L_1(x) = Z_H(x)/(n * (x - 1))` with `x = w^i`.
+    pub fn eval_l1(&self, i: usize, x: F) -> F {
+        // Could also precompute the inverses using Montgomery.
+        self.eval(i) * (self.n * (x - F::ONE)).inverse()
+    }
+}
+
 /// Evaluate the Lagrange basis `L_1` with `L_1(1) = 1`, and `L_1(x) = 0` for other members of an
 /// order `n` multiplicative subgroup.
 pub(crate) fn eval_l_1<F: Field>(n: usize, x: F) -> F {

src/prover.rs

@@ -8,7 +8,7 @@ use crate::field::extension_field::Extendable;
 use crate::field::fft::ifft;
 use crate::generator::generate_partial_witness;
 use crate::plonk_challenger::Challenger;
-use crate::plonk_common::eval_vanishing_poly_base;
+use crate::plonk_common::{eval_vanishing_poly_base, ZeroPolyOnCoset};
 use crate::polynomial::commitment::ListPolynomialCommitment;
 use crate::polynomial::polynomial::{PolynomialCoeffs, PolynomialValues};
 use crate::proof::Proof;
@@ -234,6 +234,11 @@ fn compute_quotient_polys<'a, F: Extendable<D>, const D: usize>(
         comm.get_lde_values(i * step)
     };
 
+    let z_h_on_coset = ZeroPolyOnCoset::new(
+        common_data.degree_bits,
+        common_data.max_filtered_constraint_degree_bits,
+    );
+
     let quotient_values: Vec<Vec<F>> = points
         .into_par_iter()
         .enumerate()
@@ -255,6 +260,7 @@ fn compute_quotient_polys<'a, F: Extendable<D>, const D: usize>(
             };
             let mut quotient_values = eval_vanishing_poly_base(
                 common_data,
+                i,
                 shifted_x,
                 vars,
                 local_plonk_zs,
@@ -263,9 +269,9 @@ fn compute_quotient_polys<'a, F: Extendable<D>, const D: usize>(
                 betas,
                 gammas,
                 alphas,
+                &z_h_on_coset,
             );
-            // TODO: We can avoid computing the exp.
-            let denominator_inv = (shifted_x.exp(common_data.degree() as u64) - F::ONE).inverse();
+            let denominator_inv = z_h_on_coset.eval_inverse(i);
             quotient_values
                 .iter_mut()
                 .for_each(|v| *v *= denominator_inv);