From e84c3b96d8bb6d09b35546d062b2f2b0b0ecbd02 Mon Sep 17 00:00:00 2001 From: ZK Clay Date: Mon, 22 Apr 2024 22:11:01 -0400 Subject: [PATCH 1/3] perf: optimize get_peaks_and_heights for archival MMR. --- src/util_types/mutator_set/archival_mmr.rs | 42 ++++++---------------- 1 file changed, 10 insertions(+), 32 deletions(-) diff --git a/src/util_types/mutator_set/archival_mmr.rs b/src/util_types/mutator_set/archival_mmr.rs index 60a177a0..79ba2463 100644 --- a/src/util_types/mutator_set/archival_mmr.rs +++ b/src/util_types/mutator_set/archival_mmr.rs @@ -16,7 +16,7 @@ use itertools::Itertools; use twenty_first::util_types::mmr::{ mmr_accumulator::MmrAccumulator, mmr_membership_proof::MmrMembershipProof, mmr_trait::Mmr, - shared_advanced, shared_basic, + shared_advanced, }; /// A Merkle Mountain Range is a datastructure for storing a list of hashes. @@ -228,39 +228,17 @@ impl> ArchivalMmr { /// Return a list of tuples (peaks, height) pub async fn get_peaks_with_heights_async(&self) -> Vec<(Digest, u32)> { - if self.is_empty().await { - return vec![]; - } - - // 1. Find top peak - // 2. Jump to right sibling (will not be included) - // 3. Take left child of sibling, continue until a node in tree is found - // 4. Once new node is found, jump to right sibling (will not be included) - // 5. Take left child of sibling, continue until a node in tree is found - let mut peaks_and_heights: Vec<(Digest, u32)> = vec![]; - let (mut top_peak, mut top_height) = - shared_advanced::leftmost_ancestor(self.digests.len().await - 1); - if top_peak > self.digests.len().await - 1 { - top_peak = shared_basic::left_child(top_peak, top_height); - top_height -= 1; - } + let leaf_count = self.count_leaves().await; + let (peak_heights, peak_node_indices) = get_peak_heights_and_peak_node_indices(leaf_count); + let peaks = self.digests.get_many(&peak_node_indices).await; - peaks_and_heights.push((self.digests.get(top_peak).await, top_height)); - let mut height = top_height; - let mut candidate = shared_advanced::right_sibling(top_peak, height); - 'outer: while height > 0 { - '_inner: while candidate > self.digests.len().await && height > 0 { - candidate = shared_basic::left_child(candidate, height); - height -= 1; - if candidate < self.digests.len().await { - peaks_and_heights.push((self.digests.get(candidate).await, height)); - candidate = shared_advanced::right_sibling(candidate, height); - continue 'outer; - } - } - } + let peak_heights_and_values: Vec<_> = peaks + .iter() + .zip(peak_heights.iter()) + .map(|(&x, &y)| (x, y)) + .collect(); - peaks_and_heights + peak_heights_and_values } /// Remove the last leaf from the archival MMR From adf778f3df5ed0d0af8a6f54b0e719a5e15ac498 Mon Sep 17 00:00:00 2001 From: ZK Clay Date: Tue, 23 Apr 2024 11:57:35 -0400 Subject: [PATCH 2/3] perf: optimize get_peaks_and_heights for archival MMR. --- src/util_types/mutator_set/archival_mmr.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/util_types/mutator_set/archival_mmr.rs b/src/util_types/mutator_set/archival_mmr.rs index 79ba2463..8b9a6f43 100644 --- a/src/util_types/mutator_set/archival_mmr.rs +++ b/src/util_types/mutator_set/archival_mmr.rs @@ -232,13 +232,13 @@ impl> ArchivalMmr { let (peak_heights, peak_node_indices) = get_peak_heights_and_peak_node_indices(leaf_count); let peaks = self.digests.get_many(&peak_node_indices).await; - let peak_heights_and_values: Vec<_> = peaks + let peaks_and_heights: Vec<_> = peaks .iter() .zip(peak_heights.iter()) .map(|(&x, &y)| (x, y)) .collect(); - peak_heights_and_values + peaks_and_heights } /// Remove the last leaf from the archival MMR From 377543cb67e638c0cc6e39660e5404827be45f99 Mon Sep 17 00:00:00 2001 From: ZK Clay Date: Thu, 25 Apr 2024 16:37:03 -0400 Subject: [PATCH 3/3] review feedback: remove get_peaks_and_heights_async completely. --- src/util_types/mutator_set/archival_mmr.rs | 56 +++++----------------- 1 file changed, 13 insertions(+), 43 deletions(-) diff --git a/src/util_types/mutator_set/archival_mmr.rs b/src/util_types/mutator_set/archival_mmr.rs index 8b9a6f43..625f25b8 100644 --- a/src/util_types/mutator_set/archival_mmr.rs +++ b/src/util_types/mutator_set/archival_mmr.rs @@ -41,8 +41,9 @@ where /// Return the digests of the peaks of the MMR pub async fn get_peaks(&self) -> Vec { - let peaks_and_heights = self.get_peaks_with_heights_async().await; - peaks_and_heights.into_iter().map(|x| x.0).collect() + let leaf_count = self.count_leaves().await; + let (_, peak_node_indices) = get_peak_heights_and_peak_node_indices(leaf_count); + self.digests.get_many(&peak_node_indices).await } /// Whether the MMR is empty. Note that since indexing starts at @@ -226,21 +227,6 @@ impl> ArchivalMmr { MmrMembershipProof::new(leaf_index, authentication_path) } - /// Return a list of tuples (peaks, height) - pub async fn get_peaks_with_heights_async(&self) -> Vec<(Digest, u32)> { - let leaf_count = self.count_leaves().await; - let (peak_heights, peak_node_indices) = get_peak_heights_and_peak_node_indices(leaf_count); - let peaks = self.digests.get_many(&peak_node_indices).await; - - let peaks_and_heights: Vec<_> = peaks - .iter() - .zip(peak_heights.iter()) - .map(|(&x, &y)| (x, y)) - .collect(); - - peaks_and_heights - } - /// Remove the last leaf from the archival MMR pub async fn remove_last_leaf_async(&mut self) -> Option { if self.is_empty().await { @@ -280,7 +266,6 @@ pub(crate) mod mmr_test { use twenty_first::util_types::merkle_tree_maker::MerkleTreeMaker; use twenty_first::util_types::mmr::mmr_accumulator::MmrAccumulator; use twenty_first::util_types::mmr::shared_advanced::get_peak_heights; - use twenty_first::util_types::mmr::shared_advanced::get_peak_heights_and_peak_node_indices; type Storage = OrdinaryVec; @@ -774,10 +759,8 @@ pub(crate) mod mmr_test { assert_eq!(1, mmr.count_leaves().await); assert_eq!(1, mmr.count_nodes().await); - let original_peaks_and_heights: Vec<(Digest, u32)> = - mmr.get_peaks_with_heights_async().await; - assert_eq!(1, original_peaks_and_heights.len()); - assert_eq!(0, original_peaks_and_heights[0].1); + let original_peaks: Vec = mmr.get_peaks().await; + assert_eq!(1, original_peaks.len()); { let leaf_index = 0; @@ -791,11 +774,9 @@ pub(crate) mod mmr_test { assert_eq!(2, mmr.count_leaves().await); assert_eq!(3, mmr.count_nodes().await); - let new_peaks_and_heights = mmr.get_peaks_with_heights_async().await; - assert_eq!(1, new_peaks_and_heights.len()); - assert_eq!(1, new_peaks_and_heights[0].1); + let new_peaks = mmr.get_peaks().await; + assert_eq!(1, new_peaks.len()); - let new_peaks: Vec = new_peaks_and_heights.iter().map(|x| x.0).collect(); assert!( original_mmr .verify_batch_update(&new_peaks, &[new_input_hash], &[]) @@ -845,10 +826,9 @@ pub(crate) mod mmr_test { assert_eq!(num_leaves, mmr.count_leaves().await); assert_eq!(1 + num_leaves, mmr.count_nodes().await); - let original_peaks_and_heights: Vec<(Digest, u32)> = - mmr.get_peaks_with_heights_async().await; + let original_peaks: Vec = mmr.get_peaks().await; let expected_peaks = 2; - assert_eq!(expected_peaks, original_peaks_and_heights.len()); + assert_eq!(expected_peaks, original_peaks.len()); { let leaf_index = 0; @@ -904,13 +884,8 @@ pub(crate) mod mmr_test { assert_eq!(leaf_count, mmr.count_leaves().await); assert_eq!(node_count, mmr.count_nodes().await); - let original_peaks_and_heights = mmr.get_peaks_with_heights_async().await; - let peak_heights_1: Vec = original_peaks_and_heights.iter().map(|x| x.1).collect(); - - let (peak_heights_2, _) = get_peak_heights_and_peak_node_indices(leaf_count); - assert_eq!(peak_heights_1, peak_heights_2); - - let actual_peak_count = original_peaks_and_heights.len() as u64; + let original_peaks = mmr.get_peaks().await; + let actual_peak_count = original_peaks.len() as u64; assert_eq!(peak_count, actual_peak_count); // Verify that MMR root from odd number of digests and MMR bagged peaks agree @@ -1031,12 +1006,8 @@ pub(crate) mod mmr_test { mock::get_ammr_from_digests::(input_digests.clone()).await; assert_eq!(size, mmr.count_leaves().await); assert_eq!(node_count, mmr.count_nodes().await); - let original_peaks_and_heights: Vec<(Digest, u32)> = - mmr.get_peaks_with_heights_async().await; - let peak_heights_1: Vec = original_peaks_and_heights.iter().map(|x| x.1).collect(); - let (peak_heights_2, _) = get_peak_heights_and_peak_node_indices(size); - assert_eq!(peak_heights_1, peak_heights_2); - assert_eq!(peak_count, original_peaks_and_heights.len() as u64); + let original_peaks: Vec = mmr.get_peaks().await; + assert_eq!(peak_count, original_peaks.len() as u64); // Verify that MMR root from odd number of digests and MMR bagged peaks agree let mmra_root = mmr.bag_peaks().await; @@ -1054,7 +1025,6 @@ pub(crate) mod mmr_test { let valid_res = membership_proof.verify(&peaks, input_digests[leaf_index as usize], size); assert!(valid_res); - let new_leaf: Digest = random(); // The below verify_modify tests should only fail if `wrong_leaf_index` is