From 3c1d8c3f8b12ae9e8c3bf51f9c1effcda1aa0d5f Mon Sep 17 00:00:00 2001
From: Gary Yu <gairy.yu@gmail.com>
Date: Fri, 14 Sep 2018 00:59:49 +0800
Subject: [PATCH] Small cleanup and perf improvement for `peak_map_height`
 (#1521)

---
 core/src/core/pmmr.rs |  7 ++++---
 core/tests/pmmr.rs    | 24 ++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/core/src/core/pmmr.rs b/core/src/core/pmmr.rs
index 810b0edc1..496ddfb8d 100644
--- a/core/src/core/pmmr.rs
+++ b/core/src/core/pmmr.rs
@@ -37,6 +37,7 @@
 //! either be a simple Vec or a database.
 
 use std::marker;
+use std::u64;
 
 use croaring::Bitmap;
 
@@ -573,10 +574,10 @@ pub fn peak_sizes_height(size: u64) -> (Vec<u64>, u64) {
 /// since the path turns left (resp. right) if-and-only-if
 /// a peak at that height is absent (resp. present)
 pub fn peak_map_height(mut pos: u64) -> (u64, u64) {
-	let mut peak_size = 1;
-	while peak_size <= pos {
-		peak_size = peak_size << 1 | 1;
+	if pos == 0 {
+		return (0, 0);
 	}
+	let mut peak_size = u64::MAX >> 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 e294aacfc..4262b6a24 100644
--- a/core/tests/pmmr.rs
+++ b/core/tests/pmmr.rs
@@ -15,10 +15,14 @@
 //! PMMR tests
 #[macro_use]
 extern crate grin_core as core;
+extern crate chrono;
 extern crate croaring;
 
 mod vec_backend;
 
+use chrono::prelude::Utc;
+use std::u64;
+
 use core::core::hash::Hash;
 use core::core::pmmr::{self, PMMR};
 use core::ser::PMMRIndexHashable;
@@ -34,6 +38,26 @@ fn some_peak_map() {
 	assert_eq!(pmmr::peak_map_height(5), (0b11, 1));
 	assert_eq!(pmmr::peak_map_height(6), (0b11, 2));
 	assert_eq!(pmmr::peak_map_height(7), (0b100, 0));
+	assert_eq!(pmmr::peak_map_height(u64::MAX), ((u64::MAX >> 1) + 1, 0));
+	assert_eq!(pmmr::peak_map_height(u64::MAX - 1), (u64::MAX >> 1, 63));
+}
+
+#[ignore]
+#[test]
+fn bench_peak_map() {
+	let nano_to_millis = 1.0 / 1_000_000.0;
+
+	let increments = vec![1000_000u64, 10_000_000u64, 100_000_000u64];
+
+	for v in increments {
+		let start = Utc::now().timestamp_nanos();
+		for i in 0..v {
+			let _ = pmmr::peak_map_height(i);
+		}
+		let fin = Utc::now().timestamp_nanos();
+		let dur_ms = (fin - start) as f64 * nano_to_millis;
+		println!("{:9?} peak_map_height() in {:9.3?}ms", v, dur_ms);
+	}
 }
 
 #[test]