FRI refactor (#172)

I sort of "shifted" the loop in `fri_verifier_query_round` so that `fri_combine_initial` is called before the loop, and all `compute_evaluation` calls are in the loop (rather than the final one being outside). This lines up with my mental model of FRI, and I think it's more natural as it results in a loop with no branches, no `i - 1`s, and less state stored between iterations. Also added some comments etc.

Should be functionally equivalent to the old version.

src/fri/prover.rs

@@ -74,15 +74,12 @@ fn fri_committed_trees<F: Field + Extendable<D>, const D: usize>(
         let arity = 1 << config.reduction_arity_bits[i];
 
         reverse_index_bits_in_place(&mut values.values);
-        let tree = MerkleTree::new(
-            values
-                .values
-                .par_chunks(arity)
-                .map(|chunk: &[F::Extension]| flatten(chunk))
-                .collect(),
-            config.cap_height,
-            false,
-        );
+        let chunked_values = values
+            .values
+            .par_chunks(arity)
+            .map(|chunk: &[F::Extension]| flatten(chunk))
+            .collect();
+        let tree = MerkleTree::new(chunked_values, config.cap_height, false);
 
         challenger.observe_cap(&tree.cap);
         trees.push(tree);

src/fri/recursive_verifier.rs

@@ -20,7 +20,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     fn compute_evaluation(
         &mut self,
         x: Target,
-        old_x_index_bits: &[Target],
+        x_index_within_coset_bits: &[Target],
         arity_bits: usize,
         last_evals: &[ExtensionTarget<D>],
         beta: ExtensionTarget<D>,
@@ -35,8 +35,9 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         // The evaluation vector needs to be reordered first.
         let mut evals = last_evals.to_vec();
         reverse_index_bits_in_place(&mut evals);
-        // Want `g^(arity - rev_old_x_index)` as in the out-of-circuit version. Compute it as `(g^-1)^rev_old_x_index`.
-        let start = self.exp_from_bits(g_inv_t, old_x_index_bits.iter().rev());
+        // Want `g^(arity - rev_x_index_within_coset)` as in the out-of-circuit version. Compute it
+        // as `(g^-1)^rev_x_index_within_coset`.
+        let start = self.exp_from_bits(g_inv_t, x_index_within_coset_bits.iter().rev());
         let coset_start = self.mul(start, x);
 
         // The answer is gotten by interpolating {(x*g^i, P(x*g^i))} and evaluating at beta.
@@ -294,7 +295,6 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             x_index_bits[x_index_bits.len() - common_data.config.fri_config.cap_height..]
                 .into_iter(),
         );
-        let mut domain_size = n;
         with_context!(
             self,
             "check FRI initial proof",
@@ -305,7 +305,6 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
                 cap_index
             )
         );
-        let mut old_x_index_bits = Vec::new();
 
         // `subgroup_x` is `subgroup[x_index]`, i.e., the actual field element in the domain.
         let mut subgroup_x = with_context!(self, "compute x from its index", {
@@ -316,79 +315,63 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             self.mul(g, phi)
         });
 
-        let mut evaluations: Vec<Vec<ExtensionTarget<D>>> = Vec::new();
+        // old_eval is the last derived evaluation; it will be checked for consistency with its
+        // committed "parent" value in the next iteration.
+        let mut old_eval = with_context!(
+            self,
+            "combine initial oracles",
+            self.fri_combine_initial(
+                &round_proof.initial_trees_proof,
+                alpha,
+                zeta,
+                subgroup_x,
+                precomputed_reduced_evals,
+                common_data,
+            )
+        );
+
         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 {
-                with_context!(
-                    self,
-                    "combine initial oracles",
-                    self.fri_combine_initial(
-                        &round_proof.initial_trees_proof,
-                        alpha,
-                        zeta,
-                        subgroup_x,
-                        precomputed_reduced_evals,
-                        common_data,
-                    )
+            let evals = &round_proof.steps[i].evals;
+
+            // Split x_index into the index of the coset x is in, and the index of x within that coset.
+            let coset_index_bits = x_index_bits[arity_bits..].to_vec();
+            let x_index_within_coset_bits = &x_index_bits[..arity_bits];
+            let x_index_within_coset = self.le_sum(x_index_within_coset_bits.iter());
+
+            // Check consistency with our old evaluation from the previous round.
+            self.random_access(x_index_within_coset, old_eval, evals.clone());
+
+            // Infer P(y) from {P(x)}_{x^arity=y}.
+            old_eval = with_context!(
+                self,
+                "infer evaluation using interpolation",
+                self.compute_evaluation(
+                    subgroup_x,
+                    &x_index_within_coset_bits,
+                    arity_bits,
+                    evals,
+                    betas[i],
                 )
-            } else {
-                let last_evals = &evaluations[i - 1];
-                // Infer P(y) from {P(x)}_{x^arity=y}.
-                with_context!(
-                    self,
-                    "infer evaluation using interpolation",
-                    self.compute_evaluation(
-                        subgroup_x,
-                        &old_x_index_bits,
-                        config.reduction_arity_bits[i - 1],
-                        last_evals,
-                        betas[i - 1],
-                    )
-                )
-            };
-            let evals = round_proof.steps[i].evals.clone();
-            // Insert P(y) into the evaluation vector, since it wasn't included by the prover.
-            let high_x_index_bits = x_index_bits.split_off(arity_bits);
-            old_x_index_bits = x_index_bits;
-            let low_x_index = self.le_sum(old_x_index_bits.iter());
-            self.random_access(low_x_index, e_x, evals.clone());
+            );
+
             with_context!(
                 self,
                 "verify FRI round Merkle proof.",
                 self.verify_merkle_proof_with_cap_index(
-                    flatten_target(&evals),
-                    &high_x_index_bits,
+                    flatten_target(evals),
+                    &coset_index_bits,
                     cap_index,
                     &proof.commit_phase_merkle_caps[i],
                     &round_proof.steps[i].merkle_proof,
                 )
             );
-            evaluations.push(evals);
 
-            if i > 0 {
-                // Update the point x to x^arity.
-                subgroup_x = self.exp_power_of_2(subgroup_x, config.reduction_arity_bits[i - 1]);
-            }
-            domain_size = next_domain_size;
-            x_index_bits = high_x_index_bits;
+            // Update the point x to x^arity.
+            subgroup_x = self.exp_power_of_2(subgroup_x, arity_bits);
+
+            x_index_bits = coset_index_bits;
         }
 
-        let last_evals = evaluations.last().unwrap();
-        let final_arity_bits = *config.reduction_arity_bits.last().unwrap();
-        let purported_eval = with_context!(
-            self,
-            "infer final evaluation using interpolation",
-            self.compute_evaluation(
-                subgroup_x,
-                &old_x_index_bits,
-                final_arity_bits,
-                last_evals,
-                *betas.last().unwrap(),
-            )
-        );
-        subgroup_x = self.exp_power_of_2(subgroup_x, final_arity_bits);
-
         // Final check of FRI. After all the reductions, we check that the final polynomial is equal
         // to the one sent by the prover.
         let eval = with_context!(
@@ -396,7 +379,7 @@ impl<F: Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             "evaluate final polynomial",
             proof.final_poly.eval_scalar(self, subgroup_x)
         );
-        self.assert_equal_extension(eval, purported_eval);
+        self.assert_equal_extension(eval, old_eval);
     }
 }
 

src/fri/verifier.rs

@@ -19,7 +19,7 @@ use crate::util::{log2_strict, reverse_bits, reverse_index_bits_in_place};
 /// and P' is the FRI reduced polynomial.
 fn compute_evaluation<F: Field + Extendable<D>, const D: usize>(
     x: F,
-    old_x_index: usize,
+    x_index_within_coset: usize,
     arity_bits: usize,
     last_evals: &[F::Extension],
     beta: F::Extension,
@@ -32,8 +32,8 @@ fn compute_evaluation<F: Field + Extendable<D>, const D: usize>(
     // The evaluation vector needs to be reordered first.
     let mut evals = last_evals.to_vec();
     reverse_index_bits_in_place(&mut evals);
-    let rev_old_x_index = reverse_bits(old_x_index, arity_bits);
-    let coset_start = x * g.exp((arity - rev_old_x_index) as u64);
+    let rev_x_index_within_coset = reverse_bits(x_index_within_coset, arity_bits);
+    let coset_start = x * g.exp((arity - rev_x_index_within_coset) as u64);
     // The answer is gotten by interpolating {(x*g^i, P(x*g^i))} and evaluating at beta.
     let points = g
         .powers()
@@ -258,79 +258,66 @@ fn fri_verifier_query_round<F: Field + Extendable<D>, const D: usize>(
 ) -> Result<()> {
     let config = &common_data.config.fri_config;
     let x = challenger.get_challenge();
-    let mut domain_size = n;
     let mut x_index = x.to_canonical_u64() as usize % n;
     fri_verify_initial_proof(
         x_index,
         &round_proof.initial_trees_proof,
         initial_merkle_caps,
     )?;
-    let mut old_x_index = 0;
     // `subgroup_x` is `subgroup[x_index]`, i.e., the actual field element in the domain.
     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();
+    // old_eval is the last derived evaluation; it will be checked for consistency with its
+    // committed "parent" value in the next iteration.
+    let mut old_eval = fri_combine_initial(
+        &round_proof.initial_trees_proof,
+        alpha,
+        zeta,
+        subgroup_x,
+        precomputed_reduced_evals,
+        common_data,
+    );
+
     for (i, &arity_bits) in config.reduction_arity_bits.iter().enumerate() {
         let arity = 1 << arity_bits;
-        let next_domain_size = domain_size >> arity_bits;
-        let e_x = if i == 0 {
-            fri_combine_initial(
-                &round_proof.initial_trees_proof,
-                alpha,
-                zeta,
-                subgroup_x,
-                precomputed_reduced_evals,
-                common_data,
-            )
-        } else {
-            let last_evals = &evaluations[i - 1];
-            // Infer P(y) from {P(x)}_{x^arity=y}.
-            compute_evaluation(
-                subgroup_x,
-                old_x_index,
-                config.reduction_arity_bits[i - 1],
-                last_evals,
-                betas[i - 1],
-            )
-        };
         let evals = &round_proof.steps[i].evals;
-        // Insert P(y) into the evaluation vector, since it wasn't included by the prover.
-        ensure!(evals[x_index & (arity - 1)] == e_x);
+
+        // Split x_index into the index of the coset x is in, and the index of x within that coset.
+        let coset_index = x_index >> arity_bits;
+        let x_index_within_coset = x_index & (arity - 1);
+
+        // Check consistency with our old evaluation from the previous round.
+        ensure!(evals[x_index_within_coset] == old_eval);
+
+        // Infer P(y) from {P(x)}_{x^arity=y}.
+        old_eval = compute_evaluation(
+            subgroup_x,
+            x_index_within_coset,
+            arity_bits,
+            evals,
+            betas[i],
+        );
+
         verify_merkle_proof(
             flatten(evals),
-            x_index >> arity_bits,
+            coset_index,
             &proof.commit_phase_merkle_caps[i],
             &round_proof.steps[i].merkle_proof,
             false,
         )?;
-        evaluations.push(evals.to_vec());
 
-        if i > 0 {
-            // Update the point x to x^arity.
-            subgroup_x = subgroup_x.exp_power_of_2(config.reduction_arity_bits[i - 1]);
-        }
-        domain_size = next_domain_size;
-        old_x_index = x_index & (arity - 1);
-        x_index >>= arity_bits;
+        // Update the point x to x^arity.
+        subgroup_x = subgroup_x.exp_power_of_2(arity_bits);
+
+        x_index = coset_index;
     }
 
-    let last_evals = evaluations.last().unwrap();
-    let final_arity_bits = *config.reduction_arity_bits.last().unwrap();
-    let purported_eval = compute_evaluation(
-        subgroup_x,
-        old_x_index,
-        final_arity_bits,
-        last_evals,
-        *betas.last().unwrap(),
-    );
-    subgroup_x = subgroup_x.exp_power_of_2(final_arity_bits);
-
     // Final check of FRI. After all the reductions, we check that the final polynomial is equal
     // to the one sent by the prover.
     ensure!(
-        proof.final_poly.eval(subgroup_x.into()) == purported_eval,
+        proof.final_poly.eval(subgroup_x.into()) == old_eval,
         "Final polynomial evaluation is invalid."
     );