Skip to content

Commit

Permalink
refactor: Don't return peaks when getting MMR auth path from AMMR
Browse files Browse the repository at this point in the history
No real noticable in performance, but this interface was not well liked,
so we remove the peaks from the return value. In theory this should be
slightly faster, as a step of the calculation is skipped, and since the
peak digests no longer need to be read from the database.
  • Loading branch information
Sword-Smith committed Apr 7, 2024
1 parent a5c6418 commit fa6de86
Show file tree
Hide file tree
Showing 6 changed files with 37 additions and 48 deletions.
2 changes: 1 addition & 1 deletion benches/archival_mmr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -216,7 +216,7 @@ mod batch_mutate_leaf_and_update_mps {

let mut mps = leaf_indices_for_mps_to_preserve
.iter()
.map(|i| rt.block_on(async { ammr.prove_membership_async(*i).await.0 }))
.map(|i| rt.block_on(async { ammr.prove_membership_async(*i).await }))
.collect_vec();

bencher.bench_local(|| {
Expand Down
2 changes: 1 addition & 1 deletion src/models/blockchain/block/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -793,7 +793,7 @@ mod block_tests {

let index = thread_rng().gen_range(0..blocks.len() - 1);
let block_digest = blocks[index].hash();
let (membership_proof, _) = ammr.prove_membership_async(index as u64).await;
let membership_proof = ammr.prove_membership_async(index as u64).await;
let (v, _) = membership_proof.verify(
&last_block_mmra.get_peaks(),
block_digest,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,8 @@ mod tests {

let leaf_index = rng.next_u64() % num_leafs;
let leaf = mmr.get_leaf_async(leaf_index).await;
let (mmr_mp, peaks) = mmr.prove_membership_async(leaf_index).await;
let peaks = mmr.get_peaks().await;
let mmr_mp = mmr.prove_membership_async(leaf_index).await;
let mut msmp = pseudorandom_mutator_set_membership_proof(rng.gen());
msmp.auth_path_aocl = mmr_mp;

Expand Down
56 changes: 25 additions & 31 deletions src/util_types/mutator_set/archival_mmr.rs
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,7 @@ where

let mut modified_mps: Vec<usize> = vec![];
for (i, mp) in membership_proofs.iter_mut().enumerate() {
let new_mp = self.prove_membership_async(mp.leaf_index).await.0;
let new_mp = self.prove_membership_async(mp.leaf_index).await;
if new_mp != **mp {
modified_mps.push(i);
}
Expand Down Expand Up @@ -197,10 +197,7 @@ impl<H: AlgebraicHasher, Storage: StorageVec<Digest>> ArchivalMmr<H, Storage> {
}

/// Return (membership_proof, peaks)
pub async fn prove_membership_async(
&self,
leaf_index: u64,
) -> (MmrMembershipProof<H>, Vec<Digest>) {
pub async fn prove_membership_async(&self, leaf_index: u64) -> MmrMembershipProof<H> {
// A proof consists of an authentication path
// and a list of peaks
assert!(
Expand Down Expand Up @@ -243,16 +240,7 @@ impl<H: AlgebraicHasher, Storage: StorageVec<Digest>> ArchivalMmr<H, Storage> {
index_height = next_index_info.1;
}

let peaks: Vec<Digest> = self
.get_peaks_with_heights_async()
.await
.iter()
.map(|x| x.0)
.collect();

let membership_proof = MmrMembershipProof::new(leaf_index, authentication_path);

(membership_proof, peaks)
MmrMembershipProof::new(leaf_index, authentication_path)
}

/// Return a list of tuples (peaks, height)
Expand Down Expand Up @@ -499,8 +487,8 @@ pub(crate) mod mmr_test {

// let archival_mmr = ArchivalMmr::<Hasher>::new(leaf_hashes.clone());
let archival_mmr = mock::get_ammr_from_digests::<H>(leaf_hashes.clone()).await;
let (mut membership_proof, peaks): (MmrMembershipProof<H>, Vec<Digest>) =
archival_mmr.prove_membership_async(0).await;
let peaks = archival_mmr.get_peaks().await;
let mut membership_proof = archival_mmr.prove_membership_async(0).await;

// Verify that the accumulated hash in the verifier is compared against the **correct** hash,
// not just **any** hash in the peaks list.
Expand Down Expand Up @@ -543,8 +531,8 @@ pub(crate) mod mmr_test {
let mut archival_mmr = mock::get_ammr_from_digests::<H>(leaf_hashes.clone()).await;

let leaf_index: u64 = 2;
let (mp1, old_peaks): (MmrMembershipProof<H>, Vec<Digest>) =
archival_mmr.prove_membership_async(leaf_index).await;
let old_peaks = archival_mmr.get_peaks().await;
let mp1 = archival_mmr.prove_membership_async(leaf_index).await;

// Verify single leaf

Expand All @@ -560,7 +548,7 @@ pub(crate) mod mmr_test {
let mut other_archival_mmr: ArchivalMmr<H, Storage> =
mock::get_ammr_from_digests::<H>(leaf_hashes.clone()).await;

let (mp2, _acc_hash_2) = other_archival_mmr.prove_membership_async(leaf_index).await;
let mp2 = other_archival_mmr.prove_membership_async(leaf_index).await;

// Mutate leaf + mutate leaf raw, assert that they're equivalent

Expand Down Expand Up @@ -658,7 +646,7 @@ pub(crate) mod mmr_test {
mock::get_ammr_from_digests::<H>(vec![new_leaf; size]).await;
for i in 0..size {
let i = i as u64;
let (mp, _archival_peaks) = archival.prove_membership_async(i).await;
let mp = archival.prove_membership_async(i).await;
assert_eq!(i, mp.leaf_index);
acc.mutate_leaf(&mp, new_leaf);
archival.mutate_leaf(i, new_leaf).await;
Expand All @@ -685,7 +673,8 @@ pub(crate) mod mmr_test {
mock::get_ammr_from_digests::<H>(vec![new_leaf; size]).await;
for i in 0..size {
let i = i as u64;
let (mp, peaks_before_update) = archival.prove_membership_async(i).await;
let peaks_before_update = archival.get_peaks().await;
let mp = archival.prove_membership_async(i).await;
assert_eq!(archival.get_peaks().await, peaks_before_update);

// Verify the update operation using the batch verifier
Expand Down Expand Up @@ -753,8 +742,7 @@ pub(crate) mod mmr_test {
// Verify that membership proofs are the same as generating them from an archival MMR
let archival_membership_proof_direct = archival_iterative
.prove_membership_async(leaf_index as u64)
.await
.0;
.await;
assert_eq!(archival_membership_proof_direct, archival_membership_proof);
}

Expand Down Expand Up @@ -805,7 +793,8 @@ pub(crate) mod mmr_test {

{
let leaf_index = 0;
let (membership_proof, peaks) = mmr.prove_membership_async(leaf_index).await;
let peaks = mmr.get_peaks().await;
let membership_proof = mmr.prove_membership_async(leaf_index).await;
let valid_res = membership_proof.verify(&peaks, input_hash, 1);
assert!(valid_res.0);
assert!(valid_res.1.is_some());
Expand Down Expand Up @@ -837,7 +826,7 @@ pub(crate) mod mmr_test {
futures::future::join_all((0..2).map(|i| mmr_after_append.prove_membership_async(i)))
.await
.into_iter()
.map(|p| (new_leaf, p.0))
.map(|p| (new_leaf, p))
.collect();

for leaf_index in [0u64, 1] {
Expand Down Expand Up @@ -877,7 +866,8 @@ pub(crate) mod mmr_test {
{
let leaf_index = 0;
let input_digest = input_digests[leaf_index];
let (mut membership_proof, peaks) = mmr.prove_membership_async(leaf_index as u64).await;
let peaks = mmr.get_peaks().await;
let mut membership_proof = mmr.prove_membership_async(leaf_index as u64).await;

let (mp_verifies_1, acc_hash_1) =
membership_proof.verify(&peaks, input_digest, num_leaves);
Expand Down Expand Up @@ -952,7 +942,8 @@ pub(crate) mod mmr_test {
// Get an authentication path for **all** values in MMR,
// verify that it is valid
for index in 0..leaf_count {
let (membership_proof, peaks) = mmr.prove_membership_async(index).await;
let peaks = mmr.get_peaks().await;
let membership_proof = mmr.prove_membership_async(index).await;
let valid_res =
membership_proof.verify(&peaks, input_hashes[index as usize], leaf_count);

Expand Down Expand Up @@ -1077,7 +1068,8 @@ pub(crate) mod mmr_test {
// Get an authentication path for **all** values in MMR,
// verify that it is valid
for leaf_index in 0..size {
let (mut membership_proof, peaks) = mmr.prove_membership_async(leaf_index).await;
let peaks = mmr.get_peaks().await;
let mut membership_proof = mmr.prove_membership_async(leaf_index).await;
let valid_res =
membership_proof.verify(&peaks, input_digests[leaf_index as usize], size);
assert!(valid_res.0);
Expand All @@ -1099,13 +1091,15 @@ pub(crate) mod mmr_test {
let old_leaf = input_digests[leaf_index as usize];
mmr.mutate_leaf(leaf_index, new_leaf).await;

let (new_mp, new_peaks) = mmr.prove_membership_async(leaf_index).await;
let new_peaks = mmr.get_peaks().await;
let new_mp = mmr.prove_membership_async(leaf_index).await;
assert!(new_mp.verify(&new_peaks, new_leaf, size).0);
assert!(!new_mp.verify(&new_peaks, old_leaf, size).0);

// Return the element to its former value and run prove/verify for membership
mmr.mutate_leaf(leaf_index, old_leaf).await;
let (old_mp, old_peaks) = mmr.prove_membership_async(leaf_index).await;
let old_peaks = mmr.get_peaks().await;
let old_mp = mmr.prove_membership_async(leaf_index).await;
assert!(!old_mp.verify(&old_peaks, new_leaf, size).0);
assert!(old_mp.verify(&old_peaks, old_leaf, size).0);
}
Expand Down
16 changes: 5 additions & 11 deletions src/util_types/mutator_set/archival_mutator_set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,7 @@ where
)));
}

Ok(self.aocl.prove_membership_async(index).await.0)
Ok(self.aocl.prove_membership_async(index).await)
}

/// Returns an authentication path for a chunk in the sliding window Bloom filter
Expand All @@ -184,11 +184,8 @@ where
)));
}

let chunk_auth_path: mmr::mmr_membership_proof::MmrMembershipProof<Hash> = self
.swbf_inactive
.prove_membership_async(chunk_index)
.await
.0;
let chunk_auth_path: mmr::mmr_membership_proof::MmrMembershipProof<Hash> =
self.swbf_inactive.prove_membership_async(chunk_index).await;

// This check should never fail. It would mean that chunks are missing but that the
// archival MMR has the membership proof for the chunk. That would be a programming
Expand Down Expand Up @@ -237,11 +234,8 @@ where
self.chunks.len().await > chunk_index,
"Chunks must be known if its authentication path is known."
);
let chunk_membership_proof: mmr::mmr_membership_proof::MmrMembershipProof<Hash> = self
.swbf_inactive
.prove_membership_async(chunk_index)
.await
.0;
let chunk_membership_proof: mmr::mmr_membership_proof::MmrMembershipProof<Hash> =
self.swbf_inactive.prove_membership_async(chunk_index).await;
target_chunks
.dictionary
.insert(chunk_index, (chunk_membership_proof, chunk.to_owned()));
Expand Down
6 changes: 3 additions & 3 deletions src/util_types/mutator_set/chunk_dictionary.rs
Original file line number Diff line number Diff line change
Expand Up @@ -141,7 +141,7 @@ mod chunk_dict_tests {
let archival_mmr = mock::get_ammr_from_digests::<H>(leaf_hashes).await;

let key1: u64 = 898989;
let mp1: MmrMembershipProof<H> = archival_mmr.prove_membership_async(1).await.0;
let mp1: MmrMembershipProof<H> = archival_mmr.prove_membership_async(1).await;
let chunk1: Chunk = {
Chunk {
relative_indices: (0..CHUNK_SIZE).collect(),
Expand All @@ -153,7 +153,7 @@ mod chunk_dict_tests {
// Insert two more element and verify that the hash is deterministic which implies that the
// elements in the preimage are sorted deterministically.
let key2: u64 = 8989;
let mp2: MmrMembershipProof<H> = archival_mmr.prove_membership_async(2).await.0;
let mp2: MmrMembershipProof<H> = archival_mmr.prove_membership_async(2).await;
let mut chunk2 = Chunk::empty_chunk();
chunk2.insert(CHUNK_SIZE / 2 + 1);
let value2 = (mp2, chunk2);
Expand Down Expand Up @@ -215,7 +215,7 @@ mod chunk_dict_tests {
let key: u64 = 898989;
let leaf_hashes: Vec<Digest> = random_elements(3);
let archival_mmr = mock::get_ammr_from_digests::<H>(leaf_hashes).await;
let mp: MmrMembershipProof<H> = archival_mmr.prove_membership_async(1).await.0;
let mp: MmrMembershipProof<H> = archival_mmr.prove_membership_async(1).await;
let chunk = Chunk {
relative_indices: (0..CHUNK_SIZE).collect(),
};
Expand Down

0 comments on commit fa6de86

Please sign in to comment.