Skip to content

Commit

Permalink
perf: faster count_leaves and prove_membership_async for MMRA
Browse files Browse the repository at this point in the history
With this commit the DB is only accessed once when counting leaves, whereas
earlier, a helper function that read multiple digests from the MMRA was used.

Also changes index arithmetic for getting the relevant node indices for an
MMR authentication path slightly, and then reads all the digests in one
go instead of in a loop.

Good speedup seen from this commit:
- batch-mutation benchmark is 10 % faster, regular mutation is 20 % faster.
  • Loading branch information
Sword-Smith committed Apr 7, 2024
1 parent fa6de86 commit 4d568b2
Showing 1 changed file with 20 additions and 40 deletions.
60 changes: 20 additions & 40 deletions src/util_types/mutator_set/archival_mmr.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
use crate::database::storage::storage_vec::traits::*;
use crate::prelude::twenty_first;

use tasm_lib::twenty_first::util_types::mmr::shared_advanced::get_authentication_path_node_indices;
use tasm_lib::twenty_first::util_types::mmr::shared_advanced::get_peak_heights_and_peak_node_indices;
use tasm_lib::twenty_first::util_types::mmr::shared_advanced::node_index_to_leaf_index;
use tasm_lib::twenty_first::util_types::mmr::shared_basic::leaf_index_to_mt_index_and_peak_index;
use tasm_lib::twenty_first::util_types::mmr::shared_basic::right_lineage_length_from_leaf_index;
use twenty_first::shared_math::digest::Digest;
use twenty_first::util_types::algebraic_hasher::AlgebraicHasher;
Expand Down Expand Up @@ -51,13 +54,10 @@ where

/// Return the number of leaves in the tree
pub async fn count_leaves(&self) -> u64 {
let peaks_and_heights: Vec<(_, u32)> = self.get_peaks_with_heights_async().await;
let mut acc = 0;
for (_, height) in peaks_and_heights {
acc += 1 << height
}
let leaf_index_last_insertion = node_index_to_leaf_index(self.digests.len().await).unwrap();

acc
// Add one to last index, as leaf-indices are 0-indexed
leaf_index_last_insertion + 1
}

/// Append an element to the archival MMR, return the membership proof of the newly added leaf.
Expand Down Expand Up @@ -200,45 +200,25 @@ impl<H: AlgebraicHasher, Storage: StorageVec<Digest>> ArchivalMmr<H, Storage> {
pub async fn prove_membership_async(&self, leaf_index: u64) -> MmrMembershipProof<H> {
// A proof consists of an authentication path
// and a list of peaks
let num_leafs = self.count_leaves().await;
assert!(
leaf_index < self.count_leaves().await,
leaf_index < num_leafs,
"Cannot prove membership of leaf outside of range. Got leaf_index {leaf_index}. Leaf count is {}", self.count_leaves().await
);

// Find out how long the authentication path is
let node_index = shared_advanced::leaf_index_to_node_index(leaf_index);
let mut top_height: i32 = -1;
let mut parent_index = node_index;
while parent_index < self.digests.len().await {
parent_index = shared_advanced::parent(parent_index);
top_height += 1;
}

// Build the authentication path
let mut authentication_path: Vec<Digest> = vec![];
let mut index = node_index;
let (mut right_ancestor_count, mut index_height): (u32, u32) =
shared_advanced::right_lineage_length_and_own_height(index);
while index_height < top_height as u32 {
if right_ancestor_count != 0 {
// index is right child
let left_sibling_index = shared_advanced::left_sibling(index, index_height);
authentication_path.push(self.digests.get(left_sibling_index).await);

// parent of right child is index + 1
index += 1;
} else {
// index is left child
let right_sibling_index = shared_advanced::right_sibling(index, index_height);
authentication_path.push(self.digests.get(right_sibling_index).await);

// parent of left child:
index += 1 << (index_height + 1);
}
let next_index_info = shared_advanced::right_lineage_length_and_own_height(index);
right_ancestor_count = next_index_info.0;
index_height = next_index_info.1;
}
let (_, own_index_into_peaks_list) =
leaf_index_to_mt_index_and_peak_index(leaf_index, num_leafs);
let (_, peak_indices) = get_peak_heights_and_peak_node_indices(num_leafs);
let num_nodes = self.digests.len().await;
let sibling_indices = get_authentication_path_node_indices(
node_index,
peak_indices[own_index_into_peaks_list as usize],
num_nodes,
)
.unwrap();

let authentication_path = self.digests.get_many(&sibling_indices).await;

MmrMembershipProof::new(leaf_index, authentication_path)
}
Expand Down

0 comments on commit 4d568b2

Please sign in to comment.