FRI tweaks (#111)

- Call `exp_power_of_2` instead of manual squaring
- Replace `evaluations[i]` with `evals`

src/fri/recursive_verifier.rs

@@ -246,7 +246,6 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     ) {
         let config = &common_data.config.fri_config;
         let n_log = log2_strict(n);
-        let mut evaluations: Vec<Vec<ExtensionTarget<D>>> = Vec::new();
         // TODO: Do we need to range check `x_index` to a target smaller than `p`?
         let mut x_index = challenger.get_challenge(self);
         x_index = self.split_low_high(x_index, n_log, 64).0;
@@ -273,6 +272,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             self.mul(g, phi)
         });
 
+        let mut evaluations: Vec<Vec<ExtensionTarget<D>>> = Vec::new();
         for (i, &arity_bits) in config.reduction_arity_bits.iter().enumerate() {
             let next_domain_size = domain_size >> arity_bits;
             let e_x = if i == 0 {
@@ -308,23 +308,21 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             let (low_x_index, high_x_index) =
                 self.split_low_high(x_index, arity_bits, x_index_num_bits);
             evals = self.insert(low_x_index, e_x, evals);
-            evaluations.push(evals);
             context!(
                 self,
                 "verify FRI round Merkle proof.",
                 self.verify_merkle_proof(
-                    flatten_target(&evaluations[i]),
+                    flatten_target(&evals),
                     high_x_index,
                     proof.commit_phase_merkle_roots[i],
                     &round_proof.steps[i].merkle_proof,
                 )
             );
+            evaluations.push(evals);
 
             if i > 0 {
                 // Update the point x to x^arity.
-                for _ in 0..config.reduction_arity_bits[i - 1] {
+                subgroup_x = self.exp_power_of_2(subgroup_x, config.reduction_arity_bits[i - 1]);
-                    subgroup_x = self.square(subgroup_x);
-                }
             }
             domain_size = next_domain_size;
             old_x_index = low_x_index;
@@ -345,9 +343,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
                 *betas.last().unwrap(),
             )
         );
-        for _ in 0..final_arity_bits {
+        subgroup_x = self.exp_power_of_2(subgroup_x, final_arity_bits);
-            subgroup_x = self.square(subgroup_x);
-        }
 
         // Final check of FRI. After all the reductions, we check that the final polynomial is equal
         // to the one sent by the prover.

src/fri/verifier.rs

@@ -248,7 +248,6 @@ fn fri_verifier_query_round<F: Field + Extendable<D>, const D: usize>(
     common_data: &CommonCircuitData<F, D>,
 ) -> Result<()> {
     let config = &common_data.config.fri_config;
-    let mut evaluations: Vec<Vec<F::Extension>> = Vec::new();
     let x = challenger.get_challenge();
     let mut domain_size = n;
     let mut x_index = x.to_canonical_u64() as usize % n;
@@ -262,6 +261,8 @@ fn fri_verifier_query_round<F: Field + Extendable<D>, const D: usize>(
     let log_n = log2_strict(n);
     let mut subgroup_x = F::MULTIPLICATIVE_GROUP_GENERATOR
         * F::primitive_root_of_unity(log_n).exp(reverse_bits(x_index, log_n) as u64);
+
+    let mut evaluations: Vec<Vec<F::Extension>> = Vec::new();
     for (i, &arity_bits) in config.reduction_arity_bits.iter().enumerate() {
         let arity = 1 << arity_bits;
         let next_domain_size = domain_size >> arity_bits;
@@ -288,20 +289,18 @@ fn fri_verifier_query_round<F: Field + Extendable<D>, const D: usize>(
         let mut evals = round_proof.steps[i].evals.clone();
         // Insert P(y) into the evaluation vector, since it wasn't included by the prover.
         evals.insert(x_index & (arity - 1), e_x);
-        evaluations.push(evals);
         verify_merkle_proof(
-            flatten(&evaluations[i]),
+            flatten(&evals),
             x_index >> arity_bits,
             proof.commit_phase_merkle_roots[i],
             &round_proof.steps[i].merkle_proof,
             false,
         )?;
+        evaluations.push(evals);
 
         if i > 0 {
             // Update the point x to x^arity.
-            for _ in 0..config.reduction_arity_bits[i - 1] {
+            subgroup_x = subgroup_x.exp_power_of_2(config.reduction_arity_bits[i - 1]);
-                subgroup_x = subgroup_x.square();
-            }
         }
         domain_size = next_domain_size;
         old_x_index = x_index & (arity - 1);
@@ -317,9 +316,7 @@ fn fri_verifier_query_round<F: Field + Extendable<D>, const D: usize>(
         last_evals,
         *betas.last().unwrap(),
     );
-    for _ in 0..final_arity_bits {
+    subgroup_x = subgroup_x.exp_power_of_2(final_arity_bits);
-        subgroup_x = subgroup_x.square();
-    }
 
     // Final check of FRI. After all the reductions, we check that the final polynomial is equal
     // to the one sent by the prover.

src/gadgets/arithmetic.rs

@@ -153,6 +153,15 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         product
     }
 
+    /// Exponentiate `base` to the power of `2^power_log`.
+    // TODO: Test
+    pub fn exp_power_of_2(&mut self, mut base: Target, power_log: usize) -> Target {
+        for _ in 0..power_log {
+            base = self.square(base);
+        }
+        base
+    }
+
     // TODO: Optimize this, maybe with a new gate.
     // TODO: Test
     /// Exponentiate `base` to the power of `exponent`, where `exponent < 2^num_bits`.

src/gadgets/arithmetic_extension.rs

@@ -292,7 +292,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
 
     /// Exponentiate `base` to the power of `2^power_log`.
     // TODO: Test
-    pub fn exp_power_of_2(
+    pub fn exp_power_of_2_extension(
         &mut self,
         mut base: ExtensionTarget<D>,
         power_log: usize,

src/recursive_verifier.rs

@@ -59,7 +59,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         let s_sigmas = &proof.openings.plonk_sigmas;
         let partial_products = &proof.openings.partial_products;
 
-        let zeta_pow_deg = self.exp_power_of_2(zeta, inner_common_data.degree_bits);
+        let zeta_pow_deg = self.exp_power_of_2_extension(zeta, inner_common_data.degree_bits);
         let vanishing_polys_zeta = context!(
             self,
             "evaluate the vanishing polynomial at our challenge point, zeta.",