diff --git a/src/fri/proof.rs b/src/fri/proof.rs index ccac24eb..eef39ad6 100644 --- a/src/fri/proof.rs +++ b/src/fri/proof.rs @@ -1,3 +1,4 @@ +use itertools::izip; use serde::{Deserialize, Serialize}; use crate::field::extension_field::target::ExtensionTarget; @@ -158,7 +159,7 @@ impl, const D: usize> FriProof { .map(|(is, ps)| compress_merkle_proofs(cap_height, is, &ps)) .collect::>(); - // Replace the query round proofs with the decompressed versions. + // Replace the query round proofs with the compressed versions. for (i, qrp) in query_round_proofs.iter_mut().enumerate() { qrp.initial_trees_proof = FriInitialTreeProof { evals_proofs: (0..num_initial_trees) @@ -241,19 +242,16 @@ impl, const D: usize> FriProof { } } - // Compress all Merkle proofs. - let initial_trees_proofs = initial_trees_leaves - .iter() - .zip(&initial_trees_indices) - .zip(initial_trees_proofs) - .map(|((ls, is), ps)| decompress_merkle_proofs(&ls, is, &ps, height, cap_height)) - .collect::>(); - let steps_proofs = steps_evals - .iter() - .zip(&steps_indices) - .zip(steps_proofs) - .zip(heights) - .map(|(((ls, is), ps), h)| decompress_merkle_proofs(ls, is, &ps, h, cap_height)) + // Decompress all Merkle proofs. + let initial_trees_proofs = izip!( + &initial_trees_leaves, + &initial_trees_indices, + initial_trees_proofs + ) + .map(|(ls, is, ps)| decompress_merkle_proofs(&ls, is, &ps, height, cap_height)) + .collect::>(); + let steps_proofs = izip!(&steps_evals, &steps_indices, steps_proofs, heights) + .map(|(ls, is, ps, h)| decompress_merkle_proofs(ls, is, &ps, h, cap_height)) .collect::>(); // Replace the query round proofs with the decompressed versions. diff --git a/src/hash/path_compression.rs b/src/hash/path_compression.rs index bf1b164b..b492e69b 100644 --- a/src/hash/path_compression.rs +++ b/src/hash/path_compression.rs @@ -22,6 +22,8 @@ pub(crate) fn compress_merkle_proofs( let num_leaves = 1 << height; let mut compressed_proofs = Vec::with_capacity(proofs.len()); // Holds the known nodes in the tree at a given time. The root is at index 1. + // Valid indices are 1 through n, and each element at index `i` has + // children at indices `2i` and `2i +1` its parent at index `floor(i ∕ 2)`. let mut known = vec![false; 2 * num_leaves]; for &i in indices { // The leaves are known. @@ -70,10 +72,9 @@ pub(crate) fn decompress_merkle_proofs( seen.insert(i + num_leaves, hash_or_noop(v.to_vec())); } // For every index, go up the tree by querying `seen` to get node values, or if they are unknown - // pop them from the compressed proof. - for (&i, p) in leaves_indices.iter().zip(compressed_proofs.iter_mut()) { - // We'll use `pop` to run through the Merkle paths, so we reverse it first. - p.siblings.reverse(); + // get them from the compressed proof. + for (&i, p) in leaves_indices.iter().zip(compressed_proofs) { + let mut compressed_siblings = p.siblings.into_iter(); let mut decompressed_proof = MerkleProof { siblings: Vec::new(), }; @@ -81,10 +82,10 @@ pub(crate) fn decompress_merkle_proofs( let mut current_digest = seen[&index]; for _ in 0..height - cap_height { let sibling_index = index ^ 1; - // Get the value of the sibling node by querying it or popping it from the proof. + // Get the value of the sibling node by querying it or getting it from the proof. let h = *seen .entry(sibling_index) - .or_insert_with(|| p.siblings.pop().unwrap()); + .or_insert_with(|| compressed_siblings.next().unwrap()); decompressed_proof.siblings.push(h); // Update the current digest to the value of the parent. current_digest = if index.is_even() {