perf: tweak peak related calculation

This commit is contained in:
quake
2023-02-24 20:18:32 +09:00
committed by jjy
parent 7a831abae2
commit c5bfa3d2a5
3 changed files with 117 additions and 72 deletions
+76 -53
View File
@@ -1,3 +1,4 @@
use crate::vec;
use crate::vec::Vec;
pub fn leaf_index_to_pos(index: u64) -> u64 {
@@ -15,75 +16,97 @@ pub fn leaf_index_to_mmr_size(index: u64) -> u64 {
2 * leaves_count - peak_count
}
pub fn pos_height_in_tree(mut pos: u64) -> u32 {
pos += 1;
fn all_ones(num: u64) -> bool {
num != 0 && num.count_zeros() == num.leading_zeros()
}
fn jump_left(pos: u64) -> u64 {
let bit_length = 64 - pos.leading_zeros();
let most_significant_bits = 1 << (bit_length - 1);
pos - (most_significant_bits - 1)
pub fn pos_height_in_tree(mut pos: u64) -> u8 {
if pos == 0 {
return 0;
}
while !all_ones(pos) {
pos = jump_left(pos)
let mut peak_size = u64::MAX >> pos.leading_zeros();
while peak_size > 0 {
if pos >= peak_size {
pos -= peak_size;
}
peak_size >>= 1;
}
64 - pos.leading_zeros() - 1
pos as u8
}
pub fn parent_offset(height: u32) -> u64 {
pub fn parent_offset(height: u8) -> u64 {
2 << height
}
pub fn sibling_offset(height: u32) -> u64 {
pub fn sibling_offset(height: u8) -> u64 {
(2 << height) - 1
}
pub fn get_peaks(mmr_size: u64) -> Vec<u64> {
let mut pos_s = Vec::new();
let (mut height, mut pos) = left_peak_height_pos(mmr_size);
pos_s.push(pos);
while height > 0 {
let peak = match get_right_peak(height, pos, mmr_size) {
Some(peak) => peak,
None => break,
};
height = peak.0;
pos = peak.1;
pos_s.push(pos);
/// Returns the height of the peaks in the mmr, presented by a bitmap.
/// for example, for a mmr with 11 leaves, the mmr_size is 19, it will return 0b1010.
/// 0b1011 indicates that the left peaks are at height 0, 1 and 3.
/// 14
/// / \
/// 6 13
/// / \ / \
/// 2 5 9 12 17
/// / \ / \ / \ / \ / \
/// 0 1 3 4 7 8 10 11 15 16 18
///
/// please note that when the mmr_size is invalid, it will return the bitmap of the last valid mmr.
/// in the below example, the mmr_size is 6, but it's not a valid mmr, it will return 0b11.
/// 2 5
/// / \ / \
/// 0 1 3 4
pub fn get_peak_map(mmr_size: u64) -> u64 {
if mmr_size == 0 {
return 0;
}
pos_s
}
fn get_right_peak(mut height: u32, mut pos: u64, mmr_size: u64) -> Option<(u32, u64)> {
// move to right sibling pos
pos += sibling_offset(height);
// loop until we find a pos in mmr
while pos > mmr_size - 1 {
if height == 0 {
return None;
let mut pos = mmr_size;
let mut peak_size = u64::MAX >> pos.leading_zeros();
let mut peak_map = 0;
while peak_size > 0 {
peak_map <<= 1;
if pos >= peak_size {
pos -= peak_size;
peak_map |= 1;
}
// move to left child
pos -= parent_offset(height - 1);
height -= 1;
peak_size >>= 1;
}
Some((height, pos))
peak_map
}
fn get_peak_pos_by_height(height: u32) -> u64 {
(1 << (height + 1)) - 2
}
fn left_peak_height_pos(mmr_size: u64) -> (u32, u64) {
let mut height = 1;
let mut prev_pos = 0;
let mut pos = get_peak_pos_by_height(height);
while pos < mmr_size {
height += 1;
prev_pos = pos;
pos = get_peak_pos_by_height(height);
/// Returns the pos of the peaks in the mmr.
/// for example, for a mmr with 11 leaves, the mmr_size is 19, it will return [14, 17, 18].
/// 14
/// / \
/// 6 13
/// / \ / \
/// 2 5 9 12 17
/// / \ / \ / \ / \ / \
/// 0 1 3 4 7 8 10 11 15 16 18
///
/// please note that when the mmr_size is invalid, it will return the peaks of the last valid mmr.
/// in the below example, the mmr_size is 6, but it's not a valid mmr, it will return [2, 3].
/// 2 5
/// / \ / \
/// 0 1 3 4
pub fn get_peaks(mmr_size: u64) -> Vec<u64> {
if mmr_size == 0 {
return vec![];
}
(height - 1, prev_pos)
let leading_zeros = mmr_size.leading_zeros();
let mut pos = mmr_size;
let mut peak_size = u64::MAX >> leading_zeros;
let mut peaks = Vec::with_capacity(64 - leading_zeros as usize);
let mut peaks_sum = 0;
while peak_size > 0 {
if pos >= peak_size {
pos -= peak_size;
peaks.push(peaks_sum + peak_size - 1);
peaks_sum += peak_size;
}
peak_size >>= 1;
}
peaks
}
+12 -15
View File
@@ -6,7 +6,7 @@
use crate::borrow::Cow;
use crate::collections::VecDeque;
use crate::helper::{get_peaks, parent_offset, pos_height_in_tree, sibling_offset};
use crate::helper::{get_peak_map, get_peaks, parent_offset, pos_height_in_tree, sibling_offset};
use crate::mmr_store::{MMRBatch, MMRStoreReadOps, MMRStoreWriteOps};
use crate::vec;
use crate::vec::Vec;
@@ -52,22 +52,19 @@ impl<T: Clone + PartialEq, M: Merge<Item = T>, S: MMRStoreReadOps<T>> MMR<T, M,
// push a element and return position
pub fn push(&mut self, elem: T) -> Result<u64> {
let mut elems: Vec<T> = Vec::new();
// position of new elem
let mut elems = vec![elem];
let elem_pos = self.mmr_size;
elems.push(elem);
let mut height = 0u32;
let mut pos = elem_pos;
// continue to merge tree node if next pos heigher than current
while pos_height_in_tree(pos + 1) > height {
let peak_map = get_peak_map(self.mmr_size);
let mut pos = self.mmr_size;
let mut peak = 1;
while (peak_map & peak) != 0 {
peak <<= 1;
pos += 1;
let left_pos = pos - parent_offset(height);
let right_pos = left_pos + sibling_offset(height);
let left_pos = pos - peak;
let left_elem = self.find_elem(left_pos, &elems)?;
let right_elem = self.find_elem(right_pos, &elems)?;
let parent_elem = M::merge(&left_elem, &right_elem)?;
let right_elem = elems.last().expect("checked");
let parent_elem = M::merge(&left_elem, right_elem)?;
elems.push(parent_elem);
height += 1
}
// store hashes
self.batch.append(elem_pos, elems);
@@ -129,7 +126,7 @@ impl<T: Clone + PartialEq, M: Merge<Item = T>, S: MMRStoreReadOps<T>> MMR<T, M,
return Ok(());
}
let mut queue: VecDeque<_> = pos_list.into_iter().map(|pos| (pos, 0u32)).collect();
let mut queue: VecDeque<_> = pos_list.into_iter().map(|pos| (pos, 0)).collect();
// Generate sub-tree merkle proof for positions
while let Some((pos, height)) = queue.pop_front() {
@@ -298,7 +295,7 @@ fn calculate_peak_root<'a, T: 'a + Clone, M: Merge<Item = T>, I: Iterator<Item =
let mut queue: VecDeque<_> = leaves
.into_iter()
.map(|(pos, item)| (pos, item, 0u32))
.map(|(pos, item)| (pos, item, 0))
.collect();
// calculate tree root from each items
+29 -4
View File
@@ -1,6 +1,6 @@
use super::{MergeNumberHash, NumberHash};
use crate::{
helper::{get_peaks, pos_height_in_tree},
helper::{get_peak_map, get_peaks, pos_height_in_tree},
leaf_index_to_mmr_size, leaf_index_to_pos,
util::MemStore,
MMR,
@@ -55,16 +55,41 @@ fn test_pos_height_in_tree() {
assert_eq!(pos_height_in_tree(7), 0);
}
#[test]
fn test_get_peak_map() {
assert_eq!(get_peak_map(0), 0b0);
assert_eq!(get_peak_map(1), 0b1);
assert_eq!(get_peak_map(3), 0b10);
assert_eq!(get_peak_map(4), 0b11);
// 5 and 6 are not valid mmr_size, it will return the bitmap of the last valid mmr (size 4)
assert_eq!(get_peak_map(5), 0b11);
assert_eq!(get_peak_map(6), 0b11);
assert_eq!(get_peak_map(7), 0b100);
assert_eq!(get_peak_map(8), 0b101);
// 9 is not valid mmr_size, it will return the bitmap of the last valid mmr (size 8)
assert_eq!(get_peak_map(9), 0b101);
assert_eq!(get_peak_map(15), 0b1000);
assert_eq!(get_peak_map(16), 0b1001);
assert_eq!(get_peak_map(18), 0b1010);
assert_eq!(get_peak_map(19), 0b1011);
}
#[test]
fn test_get_peaks() {
assert_eq!(get_peaks(0), vec![0]);
assert_eq!(get_peaks(0), vec![]);
assert_eq!(get_peaks(1), vec![0]);
assert_eq!(get_peaks(2), vec![0]);
assert_eq!(get_peaks(3), vec![2]);
assert_eq!(get_peaks(4), vec![2, 3]);
// 5 and 6 are not valid mmr_size, it will return the peaks of the last valid mmr (size 4)
assert_eq!(get_peaks(5), vec![2, 3]);
assert_eq!(get_peaks(6), vec![2, 5]);
assert_eq!(get_peaks(6), vec![2, 3]);
assert_eq!(get_peaks(7), vec![6]);
assert_eq!(get_peaks(8), vec![6, 7]);
// 9 is not valid mmr_size, it will return the peaks of the last valid mmr (size 8)
assert_eq!(get_peaks(9), vec![6, 7]);
assert_eq!(get_peaks(15), vec![14]);
assert_eq!(get_peaks(16), vec![14, 15]);
assert_eq!(get_peaks(18), vec![14, 17]);
assert_eq!(get_peaks(19), vec![14, 17, 18]);
}