From 2c712fefff287f46b30fb5246031e2ac375f7229 Mon Sep 17 00:00:00 2001
From: Gary Yu <gairy.yu@gmail.com>
Date: Sat, 15 Sep 2018 15:28:41 +0800
Subject: [PATCH] Small cleanup and perf improvement for `peak_sizes_height`
 and `peaks` (#1526)

* Small cleanup and perf improvement for  and

* rustfmt

* define a const 'ALL_ONES' for u64::MAX
---
 core/src/core/pmmr.rs | 19 +++++++++++--------
 core/tests/pmmr.rs    | 16 ++++++++++++++++
 2 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/core/src/core/pmmr.rs b/core/src/core/pmmr.rs
index 496ddfb8d..0d3fef8ac 100644
--- a/core/src/core/pmmr.rs
+++ b/core/src/core/pmmr.rs
@@ -47,6 +47,9 @@ use core::BlockHeader;
 use ser::{PMMRIndexHashable, PMMRable};
 use util::LOGGER;
 
+/// 64 bits all ones: 0b11111111...1
+const ALL_ONES: u64 = u64::MAX;
+
 /// Storage backend for the MMR, just needs to be indexed by order of insertion.
 /// The PMMR itself does not need the Backend to be accurate on the existence
 /// of an element (i.e. remove could be a no-op) but layers above can
@@ -501,10 +504,10 @@ where
 /// side of the range, and navigates toward lower siblings toward the right
 /// of the range.
 pub fn peaks(num: u64) -> Vec<u64> {
-	let mut peak_size = 1;
-	while peak_size < num {
-		peak_size = peak_size << 1 | 1;
+	if num == 0 {
+		return vec![];
 	}
+	let mut peak_size = ALL_ONES >> num.leading_zeros();
 	let mut num_left = num;
 	let mut sum_prev_peaks = 0;
 	let mut peaks = vec![];
@@ -546,10 +549,10 @@ pub fn insertion_to_pmmr_index(mut sz: u64) -> u64 {
 ///   / \
 ///  0   1   3   4
 pub fn peak_sizes_height(size: u64) -> (Vec<u64>, u64) {
-	let mut peak_size = 1; // start at arbitrary 2-power minus 1
-	while peak_size < size {
-		peak_size = 2 * peak_size + 1;
+	if size == 0 {
+		return (vec![], 0);
 	}
+	let mut peak_size = ALL_ONES >> size.leading_zeros();
 	let mut sizes = vec![];
 	let mut size_left = size;
 	while peak_size != 0 {
@@ -557,7 +560,7 @@ pub fn peak_sizes_height(size: u64) -> (Vec<u64>, u64) {
 			sizes.push(peak_size);
 			size_left -= peak_size;
 		}
-		peak_size /= 2;
+		peak_size >>= 1;
 	}
 	(sizes, size_left)
 }
@@ -577,7 +580,7 @@ pub fn peak_map_height(mut pos: u64) -> (u64, u64) {
 	if pos == 0 {
 		return (0, 0);
 	}
-	let mut peak_size = u64::MAX >> pos.leading_zeros();
+	let mut peak_size = ALL_ONES >> pos.leading_zeros();
 	let mut bitmap = 0;
 	while peak_size != 0 {
 		bitmap = bitmap << 1;
diff --git a/core/tests/pmmr.rs b/core/tests/pmmr.rs
index 4262b6a24..3860a831e 100644
--- a/core/tests/pmmr.rs
+++ b/core/tests/pmmr.rs
@@ -60,6 +60,22 @@ fn bench_peak_map() {
 	}
 }
 
+#[test]
+fn some_peak_size() {
+	assert_eq!(pmmr::peak_sizes_height(0), (vec![], 0));
+	assert_eq!(pmmr::peak_sizes_height(1), (vec![1], 0));
+	assert_eq!(pmmr::peak_sizes_height(2), (vec![1], 1));
+	assert_eq!(pmmr::peak_sizes_height(3), (vec![3], 0));
+	assert_eq!(pmmr::peak_sizes_height(4), (vec![3, 1], 0));
+	assert_eq!(pmmr::peak_sizes_height(5), (vec![3, 1], 1));
+	assert_eq!(pmmr::peak_sizes_height(6), (vec![3, 1], 2));
+	assert_eq!(pmmr::peak_sizes_height(7), (vec![7], 0));
+	assert_eq!(pmmr::peak_sizes_height(u64::MAX), (vec![u64::MAX], 0));
+
+	let size_of_peaks = (1..64).map(|i| u64::MAX >> i).collect::<Vec<u64>>();
+	assert_eq!(pmmr::peak_sizes_height(u64::MAX - 1), (size_of_peaks, 63));
+}
+
 #[test]
 #[allow(unused_variables)]
 fn first_100_mmr_heights() {