Move inferred elements to a new struct

src/fri/proof.rs

@@ -14,7 +14,7 @@ use crate::hash::path_compression::{compress_merkle_proofs, decompress_merkle_pr
 use crate::iop::target::Target;
 use crate::plonk::circuit_data::CommonCircuitData;
 use crate::plonk::plonk_common::PolynomialsIndexBlinding;
-use crate::plonk::proof::ProofChallenges;
+use crate::plonk::proof::{FriInferredElements, ProofChallenges};
 use crate::polynomial::polynomial::PolynomialCoeffs;
 
 /// Evaluations and Merkle proof produced by the prover in a FRI query step.
@@ -236,6 +236,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedFriProof<F, D> {
     pub(crate) fn decompress(
         self,
         challenges: &ProofChallenges<F, D>,
+        fri_inferred_elements: FriInferredElements<F, D>,
         common_data: &CommonCircuitData<F, D>,
     ) -> FriProof<F, D> {
         let CompressedFriProof {
@@ -247,14 +248,9 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedFriProof<F, D> {
         } = self;
         let ProofChallenges {
             fri_query_indices: indices,
-            fri_query_inferred_elements,
             ..
         } = challenges;
-        let mut fri_query_inferred_elements = if let Some(v) = fri_query_inferred_elements {
+        let mut fri_inferred_elements = fri_inferred_elements.0.into_iter();
-            v.iter().copied()
-        } else {
-            panic!("Proof challenges must be computed with `CompressedProofWithPublicInputs::get_challenges()`.")
-        };
         let cap_height = common_data.config.cap_height;
         let reduction_arity_bits = &common_data.fri_params.reduction_arity_bits;
         let num_reductions = reduction_arity_bits.len();
@@ -309,10 +305,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedFriProof<F, D> {
                     evals = v.to_vec();
                 } else {
                     // Otherwise insert the next inferred element.
-                    evals.insert(
+                    evals.insert(index_within_coset, fri_inferred_elements.next().unwrap());
-                        index_within_coset,
-                        fri_query_inferred_elements.next().unwrap(),
-                    );
                     evals_by_depth[i].insert(index, evals.clone());
                 }
                 steps_evals[i].push(flatten(&evals));

src/plonk/get_challenges.rs

@@ -7,7 +7,7 @@ use crate::hash::hashing::hash_n_to_1;
 use crate::iop::challenger::Challenger;
 use crate::plonk::circuit_data::CommonCircuitData;
 use crate::plonk::proof::{
-    CompressedProofWithPublicInputs, ProofChallenges, ProofWithPublicInputs,
+    CompressedProofWithPublicInputs, FriInferredElements, ProofChallenges, ProofWithPublicInputs,
 };
 use crate::util::reverse_bits;
 
@@ -87,7 +87,6 @@ impl<F: RichField + Extendable<D>, const D: usize> ProofWithPublicInputs<F, D> {
             fri_betas,
             fri_pow_response,
             fri_query_indices,
-            fri_query_inferred_elements: None,
         })
     }
 }
@@ -152,16 +151,39 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInpu
             .map(|_| challenger.get_challenge().to_canonical_u64() as usize % lde_size)
             .collect::<Vec<_>>();
 
-        let mut fri_query_inferred_elements = Vec::new();
+        Ok(ProofChallenges {
+            plonk_betas,
+            plonk_gammas,
+            plonk_alphas,
+            plonk_zeta,
+            fri_alpha,
+            fri_betas,
+            fri_pow_response,
+            fri_query_indices,
+        })
+    }
+    pub(crate) fn get_inferred_elements(
+        &self,
+        challenges: &ProofChallenges<F, D>,
+        common_data: &CommonCircuitData<F, D>,
+    ) -> FriInferredElements<F, D> {
+        let ProofChallenges {
+            plonk_zeta,
+            fri_alpha,
+            fri_betas,
+            fri_query_indices,
+            ..
+        } = challenges;
+        let mut fri_inferred_elements = Vec::new();
         // Holds the indices that have already been seen at each reduction depth.
         let mut seen_indices_by_depth =
             vec![HashSet::new(); common_data.fri_params.reduction_arity_bits.len()];
         let precomputed_reduced_evals =
-            PrecomputedReducedEvals::from_os_and_alpha(&self.proof.openings, fri_alpha);
+            PrecomputedReducedEvals::from_os_and_alpha(&self.proof.openings, *fri_alpha);
         let log_n = common_data.degree_bits + common_data.config.rate_bits;
         // Simulate the proof verification and collect the inferred elements.
         // The content of the loop is basically the same as the `fri_verifier_query_round` function.
-        for &(mut x_index) in &fri_query_indices {
+        for &(mut x_index) in fri_query_indices {
             let mut subgroup_x = F::MULTIPLICATIVE_GROUP_GENERATOR
                 * F::primitive_root_of_unity(log_n).exp_u64(reverse_bits(x_index, log_n) as u64);
             let mut old_eval = fri_combine_initial(
@@ -170,8 +192,8 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInpu
                     .opening_proof
                     .query_round_proofs
                     .initial_trees_proofs[&x_index],
-                fri_alpha,
+                *fri_alpha,
-                plonk_zeta,
+                *plonk_zeta,
                 subgroup_x,
                 precomputed_reduced_evals,
                 common_data,
@@ -186,7 +208,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInpu
                     // If this index has already been seen, we can skip the rest of the reductions.
                     break;
                 }
-                fri_query_inferred_elements.push(old_eval);
+                fri_inferred_elements.push(old_eval);
                 let arity = 1 << arity_bits;
                 let mut evals = self.proof.opening_proof.query_round_proofs.steps[i]
                     [&(x_index >> arity_bits)]
@@ -205,17 +227,6 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInpu
                 x_index >>= arity_bits;
             }
         }
-
+        FriInferredElements(fri_inferred_elements)
-        Ok(ProofChallenges {
-            plonk_betas,
-            plonk_gammas,
-            plonk_alphas,
-            plonk_zeta,
-            fri_alpha,
-            fri_betas,
-            fri_pow_response,
-            fri_query_indices,
-            fri_query_inferred_elements: Some(fri_query_inferred_elements),
-        })
     }
 }

src/plonk/proof.rs

@@ -122,6 +122,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProof<F, D> {
     pub(crate) fn decompress(
         self,
         challenges: &ProofChallenges<F, D>,
+        fri_inferred_elements: FriInferredElements<F, D>,
         common_data: &CommonCircuitData<F, D>,
     ) -> Proof<F, D> {
         let CompressedProof {
@@ -137,7 +138,11 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProof<F, D> {
             plonk_zs_partial_products_cap,
             quotient_polys_cap,
             openings,
-            opening_proof: opening_proof.decompress(challenges, common_data),
+            opening_proof: opening_proof.decompress(
+                &challenges,
+                fri_inferred_elements,
+                common_data,
+            ),
         }
     }
 }
@@ -155,7 +160,10 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInpu
         common_data: &CommonCircuitData<F, D>,
     ) -> anyhow::Result<ProofWithPublicInputs<F, D>> {
         let challenges = self.get_challenges(common_data)?;
-        let compressed_proof = self.proof.decompress(&challenges, common_data);
+        let fri_inferred_elements = self.get_inferred_elements(&challenges, common_data);
+        let compressed_proof =
+            self.proof
+                .decompress(&challenges, fri_inferred_elements, common_data);
         Ok(ProofWithPublicInputs {
             public_inputs: self.public_inputs,
             proof: compressed_proof,
@@ -168,7 +176,10 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInpu
         common_data: &CommonCircuitData<F, D>,
     ) -> anyhow::Result<()> {
         let challenges = self.get_challenges(common_data)?;
-        let compressed_proof = self.proof.decompress(&challenges, common_data);
+        let fri_inferred_elements = self.get_inferred_elements(&challenges, common_data);
+        let compressed_proof =
+            self.proof
+                .decompress(&challenges, fri_inferred_elements, common_data);
         verify_with_challenges(
             ProofWithPublicInputs {
                 public_inputs: self.public_inputs,
@@ -223,12 +234,13 @@ pub(crate) struct ProofChallenges<F: RichField + Extendable<D>, const D: usize>
 
     // Indices at which the oracle is queried in FRI.
     pub fri_query_indices: Vec<usize>,
-
-    // Coset element that can be inferred in the FRI reduction step.
-    // Is typically set to None iff the challenges are computed from a non-compressed proof.
-    pub fri_query_inferred_elements: Option<Vec<F::Extension>>,
 }
 
+/// Coset element that can be inferred in the FRI reduction step.
+pub(crate) struct FriInferredElements<F: RichField + Extendable<D>, const D: usize>(
+    pub Vec<F::Extension>,
+);
+
 pub struct ProofWithPublicInputsTarget<const D: usize> {
     pub proof: ProofTarget<D>,
     pub public_inputs: Vec<Target>,