Merge pull request #298 from mir-protocol/remove_inferred_elmt

Remove inferred element in compressed proof

src/fri/proof.rs

@@ -14,6 +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::{FriInferredElements, ProofChallenges};
 use crate::polynomial::polynomial::PolynomialCoeffs;
 
 /// Evaluations and Merkle proof produced by the prover in a FRI query step.
@@ -164,9 +165,13 @@ impl<F: RichField + Extendable<D>, const D: usize> FriProof<F, D> {
                 initial_trees_proofs[i].push(proof);
             }
             for (i, query_step) in steps.into_iter().enumerate() {
+                let index_within_coset = index & ((1 << reduction_arity_bits[i]) - 1);
                 index >>= reduction_arity_bits[i];
                 steps_indices[i].push(index);
-                steps_evals[i].push(query_step.evals);
+                let mut evals = query_step.evals;
+                // Remove the element that can be inferred.
+                evals.remove(index_within_coset);
+                steps_evals[i].push(evals);
                 steps_proofs[i].push(query_step.merkle_proof);
             }
         }
@@ -228,9 +233,10 @@ impl<F: RichField + Extendable<D>, const D: usize> FriProof<F, D> {
 
 impl<F: RichField + Extendable<D>, const D: usize> CompressedFriProof<F, D> {
     /// Decompress all the Merkle paths in the FRI proof and reinsert duplicate indices.
-    pub fn decompress(
+    pub(crate) fn decompress(
         self,
-        indices: &[usize],
+        challenges: &ProofChallenges<F, D>,
+        fri_inferred_elements: FriInferredElements<F, D>,
         common_data: &CommonCircuitData<F, D>,
     ) -> FriProof<F, D> {
         let CompressedFriProof {
@@ -240,6 +246,11 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedFriProof<F, D> {
             pow_witness,
             ..
         } = self;
+        let ProofChallenges {
+            fri_query_indices: indices,
+            ..
+        } = challenges;
+        let mut fri_inferred_elements = fri_inferred_elements.0.into_iter();
         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();
@@ -267,7 +278,12 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedFriProof<F, D> {
             })
             .collect::<Vec<_>>();
 
-        for mut index in indices.iter().copied() {
+        // Holds the `evals` vectors that have already been reconstructed at each reduction depth.
+        let mut evals_by_depth = vec![
+            HashMap::<usize, Vec<_>>::new();
+            common_data.fri_params.reduction_arity_bits.len()
+        ];
+        for &(mut index) in indices {
             let initial_trees_proof = query_round_proofs.initial_trees_proofs[&index].clone();
             for (i, (leaves_data, proof)) in
                 initial_trees_proof.evals_proofs.into_iter().enumerate()
@@ -277,11 +293,23 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedFriProof<F, D> {
                 initial_trees_proofs[i].push(proof);
             }
             for i in 0..num_reductions {
+                let index_within_coset = index & ((1 << reduction_arity_bits[i]) - 1);
                 index >>= reduction_arity_bits[i];
-                let query_step = query_round_proofs.steps[i][&index].clone();
+                let FriQueryStep {
+                    mut evals,
+                    merkle_proof,
+                } = query_round_proofs.steps[i][&index].clone();
                 steps_indices[i].push(index);
-                steps_evals[i].push(flatten(&query_step.evals));
-                steps_proofs[i].push(query_step.merkle_proof);
+                if let Some(v) = evals_by_depth[i].get(&index) {
+                    // If this index has already been seen, get `evals` from the `HashMap`.
+                    evals = v.to_vec();
+                } else {
+                    // Otherwise insert the next inferred element.
+                    evals.insert(index_within_coset, fri_inferred_elements.next().unwrap());
+                    evals_by_depth[i].insert(index, evals.clone());
+                }
+                steps_evals[i].push(flatten(&evals));
+                steps_proofs[i].push(merkle_proof);
             }
         }
 

src/fri/verifier.rs

@@ -15,7 +15,7 @@ use crate::util::{log2_strict, reverse_bits, reverse_index_bits_in_place};
 
 /// Computes P'(x^arity) from {P(x*g^i)}_(i=0..arity), where g is a `arity`-th root of unity
 /// and P' is the FRI reduced polynomial.
-fn compute_evaluation<F: Field + Extendable<D>, const D: usize>(
+pub(crate) fn compute_evaluation<F: Field + Extendable<D>, const D: usize>(
     x: F,
     x_index_within_coset: usize,
     arity_bits: usize,
@@ -92,7 +92,7 @@ pub(crate) fn verify_fri_proof<F: RichField + Extendable<D>, const D: usize>(
             challenges,
             precomputed_reduced_evals,
             initial_merkle_caps,
-            &proof,
+            proof,
             x_index,
             n,
             round_proof,
@@ -118,14 +118,14 @@ fn fri_verify_initial_proof<F: RichField>(
 /// Holds the reduced (by `alpha`) evaluations at `zeta` for the polynomial opened just at
 /// zeta, for `Z` at zeta and for `Z` at `g*zeta`.
 #[derive(Copy, Clone, Debug)]
-struct PrecomputedReducedEvals<F: Extendable<D>, const D: usize> {
+pub(crate) struct PrecomputedReducedEvals<F: Extendable<D>, const D: usize> {
     pub single: F::Extension,
     pub zs: F::Extension,
     pub zs_right: F::Extension,
 }
 
 impl<F: Extendable<D>, const D: usize> PrecomputedReducedEvals<F, D> {
-    fn from_os_and_alpha(os: &OpeningSet<F, D>, alpha: F::Extension) -> Self {
+    pub(crate) fn from_os_and_alpha(os: &OpeningSet<F, D>, alpha: F::Extension) -> Self {
         let mut alpha = ReducingFactor::new(alpha);
         let single = alpha.reduce(
             os.constants
@@ -146,7 +146,7 @@ impl<F: Extendable<D>, const D: usize> PrecomputedReducedEvals<F, D> {
     }
 }
 
-fn fri_combine_initial<F: RichField + Extendable<D>, const D: usize>(
+pub(crate) fn fri_combine_initial<F: RichField + Extendable<D>, const D: usize>(
     proof: &FriInitialTreeProof<F>,
     alpha: F::Extension,
     zeta: F::Extension,
@@ -157,10 +157,6 @@ fn fri_combine_initial<F: RichField + Extendable<D>, const D: usize>(
     let config = &common_data.config;
     assert!(D > 1, "Not implemented for D=1.");
     let degree_log = common_data.degree_bits;
-    debug_assert_eq!(
-        degree_log,
-        common_data.config.cap_height + proof.evals_proofs[0].1.siblings.len() - config.rate_bits
-    );
     let subgroup_x = F::Extension::from_basefield(subgroup_x);
     let mut alpha = ReducingFactor::new(alpha);
     let mut sum = F::Extension::ZERO;

src/plonk/get_challenges.rs

@@ -1,11 +1,15 @@
+use std::collections::HashSet;
+
 use crate::field::extension_field::Extendable;
 use crate::field::field_types::RichField;
+use crate::fri::verifier::{compute_evaluation, fri_combine_initial, PrecomputedReducedEvals};
 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;
 
 impl<F: RichField + Extendable<D>, const D: usize> ProofWithPublicInputs<F, D> {
     pub(crate) fn fri_query_indices(
@@ -88,13 +92,6 @@ impl<F: RichField + Extendable<D>, const D: usize> ProofWithPublicInputs<F, D> {
 }
 
 impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInputs<F, D> {
-    pub(crate) fn fri_query_indices(
-        &self,
-        common_data: &CommonCircuitData<F, D>,
-    ) -> anyhow::Result<Vec<usize>> {
-        Ok(self.get_challenges(common_data)?.fri_query_indices)
-    }
-
     pub(crate) fn get_challenges(
         &self,
         common_data: &CommonCircuitData<F, D>,
@@ -135,7 +132,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInpu
                 challenger.observe_cap(cap);
                 challenger.get_extension_challenge()
             })
-            .collect();
+            .collect::<Vec<_>>();
 
         challenger.observe_extension_elements(&self.proof.opening_proof.final_poly.coeffs);
 
@@ -152,7 +149,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInpu
 
         let fri_query_indices = (0..num_fri_queries)
             .map(|_| challenger.get_challenge().to_canonical_u64() as usize % lde_size)
-            .collect();
+            .collect::<Vec<_>>();
 
         Ok(ProofChallenges {
             plonk_betas,
@@ -165,4 +162,71 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInpu
             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);
+        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 {
+            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(
+                &self
+                    .proof
+                    .opening_proof
+                    .query_round_proofs
+                    .initial_trees_proofs[&x_index],
+                *fri_alpha,
+                *plonk_zeta,
+                subgroup_x,
+                precomputed_reduced_evals,
+                common_data,
+            );
+            for (i, &arity_bits) in common_data
+                .fri_params
+                .reduction_arity_bits
+                .iter()
+                .enumerate()
+            {
+                let coset_index = x_index >> arity_bits;
+                if !seen_indices_by_depth[i].insert(coset_index) {
+                    // If this index has already been seen, we can skip the rest of the reductions.
+                    break;
+                }
+                fri_inferred_elements.push(old_eval);
+                let arity = 1 << arity_bits;
+                let mut evals = self.proof.opening_proof.query_round_proofs.steps[i][&coset_index]
+                    .evals
+                    .clone();
+                let x_index_within_coset = x_index & (arity - 1);
+                evals.insert(x_index_within_coset, old_eval);
+                old_eval = compute_evaluation(
+                    subgroup_x,
+                    x_index_within_coset,
+                    arity_bits,
+                    &evals,
+                    fri_betas[i],
+                );
+                subgroup_x = subgroup_x.exp_power_of_2(arity_bits);
+                x_index = coset_index;
+            }
+        }
+        FriInferredElements(fri_inferred_elements)
+    }
 }

src/plonk/proof.rs

@@ -10,7 +10,8 @@ use crate::hash::hash_types::{HashOut, MerkleCapTarget};
 use crate::hash::hashing::hash_n_to_hash;
 use crate::hash::merkle_tree::MerkleCap;
 use crate::iop::target::Target;
-use crate::plonk::circuit_data::CommonCircuitData;
+use crate::plonk::circuit_data::{CommonCircuitData, VerifierOnlyCircuitData};
+use crate::plonk::verifier::verify_with_challenges;
 use crate::util::serialization::Buffer;
 
 #[derive(Serialize, Deserialize, Clone, Debug, Eq, PartialEq)]
@@ -118,9 +119,10 @@ pub struct CompressedProof<F: Extendable<D>, const D: usize> {
 
 impl<F: RichField + Extendable<D>, const D: usize> CompressedProof<F, D> {
     /// Decompress the proof.
-    pub fn decompress(
+    pub(crate) fn decompress(
         self,
-        indices: &[usize],
+        challenges: &ProofChallenges<F, D>,
+        fri_inferred_elements: FriInferredElements<F, D>,
         common_data: &CommonCircuitData<F, D>,
     ) -> Proof<F, D> {
         let CompressedProof {
@@ -136,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(indices, common_data),
+            opening_proof: opening_proof.decompress(
+                &challenges,
+                fri_inferred_elements,
+                common_data,
+            ),
         }
     }
 }
@@ -153,14 +159,38 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedProofWithPublicInpu
         self,
         common_data: &CommonCircuitData<F, D>,
     ) -> anyhow::Result<ProofWithPublicInputs<F, D>> {
-        let indices = self.fri_query_indices(common_data)?;
-        let compressed_proof = self.proof.decompress(&indices, common_data);
+        let challenges = self.get_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,
         })
     }
 
+    pub(crate) fn verify(
+        self,
+        verifier_data: &VerifierOnlyCircuitData<F>,
+        common_data: &CommonCircuitData<F, D>,
+    ) -> anyhow::Result<()> {
+        let challenges = self.get_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,
+                proof: compressed_proof,
+            },
+            challenges,
+            verifier_data,
+            common_data,
+        )
+    }
+
     pub(crate) fn get_public_inputs_hash(&self) -> HashOut<F> {
         hash_n_to_hash(self.public_inputs.clone(), true)
     }
@@ -202,9 +232,15 @@ pub(crate) struct ProofChallenges<F: RichField + Extendable<D>, const D: usize>
 
     pub fri_pow_response: F,
 
+    // 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.
+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>,
@@ -273,6 +309,7 @@ mod tests {
 
     use crate::field::crandall_field::CrandallField;
     use crate::field::field_types::Field;
+    use crate::fri::reduction_strategies::FriReductionStrategy;
     use crate::iop::witness::PartialWitness;
     use crate::plonk::circuit_builder::CircuitBuilder;
     use crate::plonk::circuit_data::CircuitConfig;
@@ -284,6 +321,7 @@ mod tests {
         const D: usize = 4;
 
         let mut config = CircuitConfig::large_config();
+        config.fri_config.reduction_strategy = FriReductionStrategy::Fixed(vec![2, 1]);
         config.fri_config.num_query_rounds = 50;
 
         let pw = PartialWitness::new();
@@ -300,12 +338,14 @@ mod tests {
         builder.connect(zt, comp_zt);
         let data = builder.build();
         let proof = data.prove(pw)?;
+        verify(proof.clone(), &data.verifier_only, &data.common)?;
 
         // Verify that `decompress ∘ compress = identity`.
         let compressed_proof = proof.clone().compress(&data.common)?;
-        let decompressed_compressed_proof = compressed_proof.decompress(&data.common)?;
+        let decompressed_compressed_proof = compressed_proof.clone().decompress(&data.common)?;
         assert_eq!(proof, decompressed_compressed_proof);
 
-        verify(proof, &data.verifier_only, &data.common)
+        verify(proof, &data.verifier_only, &data.common)?;
+        compressed_proof.verify(&data.verifier_only, &data.common)
     }
 }

src/plonk/verifier.rs

@@ -5,7 +5,7 @@ use crate::field::field_types::{Field, RichField};
 use crate::fri::verifier::verify_fri_proof;
 use crate::plonk::circuit_data::{CommonCircuitData, VerifierOnlyCircuitData};
 use crate::plonk::plonk_common::reduce_with_powers;
-use crate::plonk::proof::ProofWithPublicInputs;
+use crate::plonk::proof::{ProofChallenges, ProofWithPublicInputs};
 use crate::plonk::vanishing_poly::eval_vanishing_poly;
 use crate::plonk::vars::EvaluationVars;
 
@@ -14,9 +14,17 @@ pub(crate) fn verify<F: RichField + Extendable<D>, const D: usize>(
     verifier_data: &VerifierOnlyCircuitData<F>,
     common_data: &CommonCircuitData<F, D>,
 ) -> Result<()> {
-    let public_inputs_hash = &proof_with_pis.get_public_inputs_hash();
-
     let challenges = proof_with_pis.get_challenges(common_data)?;
+    verify_with_challenges(proof_with_pis, challenges, verifier_data, common_data)
+}
+
+pub(crate) fn verify_with_challenges<F: RichField + Extendable<D>, const D: usize>(
+    proof_with_pis: ProofWithPublicInputs<F, D>,
+    challenges: ProofChallenges<F, D>,
+    verifier_data: &VerifierOnlyCircuitData<F>,
+    common_data: &CommonCircuitData<F, D>,
+) -> Result<()> {
+    let public_inputs_hash = &proof_with_pis.get_public_inputs_hash();
 
     let ProofWithPublicInputs { proof, .. } = proof_with_pis;
 

src/util/serialization.rs

@@ -248,8 +248,9 @@ impl Buffer {
     fn read_fri_query_step<F: Extendable<D>, const D: usize>(
         &mut self,
         arity: usize,
+        compressed: bool,
     ) -> Result<FriQueryStep<F, D>> {
-        let evals = self.read_field_ext_vec::<F, D>(arity)?;
+        let evals = self.read_field_ext_vec::<F, D>(arity - if compressed { 1 } else { 0 })?;
         let merkle_proof = self.read_merkle_proof()?;
         Ok(FriQueryStep {
             evals,
@@ -281,7 +282,7 @@ impl Buffer {
                 .fri_params
                 .reduction_arity_bits
                 .iter()
-                .map(|&ar| self.read_fri_query_step(1 << ar))
+                .map(|&ar| self.read_fri_query_step(1 << ar, false))
                 .collect::<Result<_>>()?;
             fqrs.push(FriQueryRound {
                 initial_trees_proof,
@@ -424,7 +425,7 @@ impl Buffer {
             });
             indices.dedup();
             let query_steps = (0..indices.len())
-                .map(|_| self.read_fri_query_step(1 << a))
+                .map(|_| self.read_fri_query_step(1 << a, true))
                 .collect::<Result<Vec<_>>>()?;
             steps.push(
                 indices