Merge branch 'duplicate_indices' into custom_serializer

src/fri/proof.rs

@@ -1,4 +1,4 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 
 use itertools::izip;
 use serde::{Deserialize, Serialize};
@@ -87,7 +87,7 @@ pub struct CompressedFriQueryRounds<F: Extendable<D>, const D: usize> {
     pub indices: Vec<usize>,
     /// Map from initial indices `i` to the `FriInitialProof` for the `i`th leaf.
     pub initial_trees_proofs: HashMap<usize, FriInitialTreeProof<F>>,
-    /// For each FRI query step, a map from initial indices `i` to the `FriInitialProof` for the `i`th leaf.
+    /// For each FRI query step, a map from indices `i` to the `FriQueryStep` for the `i`th leaf.
     pub steps: Vec<HashMap<usize, FriQueryStep<F, D>>>,
 }
 
@@ -203,14 +203,17 @@ impl<F: RichField + Extendable<D>, const D: usize> FriProof<F, D> {
             };
             compressed_query_proofs
                 .initial_trees_proofs
-                .insert(index, initial_proof);
+                .entry(index)
+                .or_insert(initial_proof);
             for j in 0..num_reductions {
                 index >>= reduction_arity_bits[j];
                 let query_step = FriQueryStep {
                     evals: steps_evals[j][i].clone(),
                     merkle_proof: steps_proofs[j][i].clone(),
                 };
-                compressed_query_proofs.steps[j].insert(index, query_step);
+                compressed_query_proofs.steps[j]
+                    .entry(index)
+                    .or_insert(query_step);
             }
         }
 
@@ -264,8 +267,17 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedFriProof<F, D> {
             })
             .collect::<Vec<_>>();
 
+        let mut seen_indices = vec![HashSet::new(); num_reductions + 1];
         for mut index in indices.iter().copied() {
-            let initial_trees_proof = query_round_proofs.initial_trees_proofs[&index].clone();
+            let mut initial_trees_proof = query_round_proofs.initial_trees_proofs[&index].clone();
+            if !seen_indices[0].insert(index) {
+                initial_trees_proof
+                    .evals_proofs
+                    .iter_mut()
+                    .for_each(|(_, p)| {
+                        p.siblings = vec![];
+                    });
+            }
             for (i, (leaves_data, proof)) in
                 initial_trees_proof.evals_proofs.into_iter().enumerate()
             {
@@ -275,7 +287,10 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedFriProof<F, D> {
             }
             for i in 0..num_reductions {
                 index >>= reduction_arity_bits[i];
-                let query_step = query_round_proofs.steps[i][&index].clone();
+                let mut query_step = query_round_proofs.steps[i][&index].clone();
+                if !seen_indices[1 + i].insert(index) {
+                    query_step.merkle_proof.siblings = vec![];
+                }
                 steps_indices[i].push(index);
                 steps_evals[i].push(flatten(&query_step.evals));
                 steps_proofs[i].push(query_step.merkle_proof);
@@ -295,7 +310,7 @@ impl<F: RichField + Extendable<D>, const D: usize> CompressedFriProof<F, D> {
             .collect::<Vec<_>>();
 
         let mut decompressed_query_proofs = Vec::with_capacity(num_reductions);
-        for i in 0..num_reductions {
+        for i in 0..indices.len() {
             let initial_trees_proof = FriInitialTreeProof {
                 evals_proofs: (0..num_initial_trees)
                     .map(|j| {

src/plonk/get_challenges.rs

@@ -1,6 +1,5 @@
 use crate::field::extension_field::Extendable;
 use crate::field::field_types::RichField;
-use crate::fri::proof::FriProof;
 use crate::hash::hashing::hash_n_to_1;
 use crate::iop::challenger::Challenger;
 use crate::plonk::circuit_data::CommonCircuitData;

src/plonk/proof.rs

@@ -283,7 +283,8 @@ mod tests {
         type F = CrandallField;
         const D: usize = 4;
 
-        let config = CircuitConfig::large_config();
+        let mut config = CircuitConfig::large_config();
+        config.fri_config.num_query_rounds = 50;
 
         let pw = PartialWitness::new();
         let mut builder = CircuitBuilder::<F, D>::new(config);

src/plonk/recursive_verifier.rs

@@ -490,6 +490,10 @@ mod tests {
         assert_eq!(recursive_proof, proof_from_bytes);
         let now = std::time::Instant::now();
         let compressed_recursive_proof = recursive_proof.clone().compress(&data.common)?;
+        let decompressed_compressed_proof = compressed_recursive_proof
+            .clone()
+            .decompress(&data.common)?;
+        assert_eq!(recursive_proof, decompressed_compressed_proof);
         info!("{:.4} to compress proof", now.elapsed().as_secs_f64());
         let compressed_proof_bytes = compressed_recursive_proof.to_bytes()?;
         info!(