diff --git a/src/helper.rs b/src/helper.rs index d997ef9..7cfb61a 100644 --- a/src/helper.rs +++ b/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 { - 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 { - 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 { + 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 } diff --git a/src/mmr.rs b/src/mmr.rs index b9d5b94..f32883b 100644 --- a/src/mmr.rs +++ b/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, S: MMRStoreReadOps> MMR Result { - let mut elems: Vec = 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, S: MMRStoreReadOps> MMR = 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, I: Iterator = leaves .into_iter() - .map(|(pos, item)| (pos, item, 0u32)) + .map(|(pos, item)| (pos, item, 0)) .collect(); // calculate tree root from each items diff --git a/src/tests/test_helper.rs b/src/tests/test_helper.rs index dbc08f0..167f6bd 100644 --- a/src/tests/test_helper.rs +++ b/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]); }