perf: tweak peak related calculation

src/helper.rs

@@ -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 {
+pub fn pos_height_in_tree(mut pos: u64) -> u8 {
-    pos += 1;
+    if pos == 0 {
-    fn all_ones(num: u64) -> bool {
+        return 0;
-        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)
     }
 
-    while !all_ones(pos) {
+    let mut peak_size = u64::MAX >> pos.leading_zeros();
-        pos = jump_left(pos)
+    while peak_size > 0 {
+        if pos >= peak_size {
+            pos -= peak_size;
+        }
+        peak_size >>= 1;
     }
-
+    pos as u8
-    64 - pos.leading_zeros() - 1
 }
 
-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> {
+/// Returns the height of the peaks in the mmr, presented by a bitmap.
-    let mut pos_s = Vec::new();
+/// for example, for a mmr with 11 leaves, the mmr_size is 19, it will return 0b1010.
-    let (mut height, mut pos) = left_peak_height_pos(mmr_size);
+/// 0b1011 indicates that the left peaks are at height 0, 1 and 3.
-    pos_s.push(pos);
+///           14
-    while height > 0 {
+///        /       \
-        let peak = match get_right_peak(height, pos, mmr_size) {
+///      6          13
-            Some(peak) => peak,
+///    /   \       /   \
-            None => break,
+///   2     5     9     12     17
-        };
+///  / \   /  \  / \   /  \   /  \
-        height = peak.0;
+/// 0   1 3   4 7   8 10  11 15  16 18
-        pos = peak.1;
+///
-        pos_s.push(pos);
+/// 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)> {
+    let mut pos = mmr_size;
-    // move to right sibling pos
+    let mut peak_size = u64::MAX >> pos.leading_zeros();
-    pos += sibling_offset(height);
+    let mut peak_map = 0;
-    // loop until we find a pos in mmr
+    while peak_size > 0 {
-    while pos > mmr_size - 1 {
+        peak_map <<= 1;
-        if height == 0 {
+        if pos >= peak_size {
-            return None;
+            pos -= peak_size;
+            peak_map |= 1;
         }
-        // move to left child
+        peak_size >>= 1;
-        pos -= parent_offset(height - 1);
-        height -= 1;
     }
-    Some((height, pos))
+
+    peak_map
 }
 
-fn get_peak_pos_by_height(height: u32) -> u64 {
+/// Returns the pos of the peaks in the mmr.
-    (1 << (height + 1)) - 2
+/// for example, for a mmr with 11 leaves, the mmr_size is 19, it will return [14, 17, 18].
-}
+///           14
-
+///        /       \
-fn left_peak_height_pos(mmr_size: u64) -> (u32, u64) {
+///      6          13
-    let mut height = 1;
+///    /   \       /   \
-    let mut prev_pos = 0;
+///   2     5     9     12     17
-    let mut pos = get_peak_pos_by_height(height);
+///  / \   /  \  / \   /  \   /  \
-    while pos < mmr_size {
+/// 0   1 3   4 7   8 10  11 15  16 18
-        height += 1;
+///
-        prev_pos = pos;
+/// please note that when the mmr_size is invalid, it will return the peaks of the last valid mmr.
-        pos = get_peak_pos_by_height(height);
+/// 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
 }

src/mmr.rs

@@ -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();
+        let mut elems = vec![elem];
-        // position of new elem
         let elem_pos = self.mmr_size;
-        elems.push(elem);
+        let peak_map = get_peak_map(self.mmr_size);
-        let mut height = 0u32;
+        let mut pos = self.mmr_size;
-        let mut pos = elem_pos;
+        let mut peak = 1;
-        // continue to merge tree node if next pos heigher than current
+        while (peak_map & peak) != 0 {
-        while pos_height_in_tree(pos + 1) > height {
+            peak <<= 1;
             pos += 1;
-            let left_pos = pos - parent_offset(height);
+            let left_pos = pos - peak;
-            let right_pos = left_pos + sibling_offset(height);
             let left_elem = self.find_elem(left_pos, &elems)?;
-            let right_elem = self.find_elem(right_pos, &elems)?;
+            let right_elem = elems.last().expect("checked");
-            let parent_elem = M::merge(&left_elem, &right_elem)?;
+            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

src/tests/test_helper.rs

@@ -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]);
 }