From 26b411e79e3a959f97ff3e2fa4ad37f3dbf333a3 Mon Sep 17 00:00:00 2001 From: hashmap Date: Sun, 24 May 2020 17:50:27 +0200 Subject: [PATCH] Reduce memory allocations in PMMR (#3328) We have a pattern in the code - return Vec, turn it into an iterator, filter and collect to another Vec. The idea was to return iterator where possible to avoid allocating intermediate vecs. --- core/src/core/pmmr/pmmr.rs | 71 ++++++++++++++++++++++++++------------ core/tests/merkle_proof.rs | 4 +-- core/tests/pmmr.rs | 25 +++++++------- store/src/prune_list.rs | 3 +- 4 files changed, 63 insertions(+), 40 deletions(-) diff --git a/core/src/core/pmmr/pmmr.rs b/core/src/core/pmmr/pmmr.rs index 904d082b6..52e6e5b96 100644 --- a/core/src/core/pmmr/pmmr.rs +++ b/core/src/core/pmmr/pmmr.rs @@ -90,16 +90,13 @@ where } /// Returns a vec of the peaks of this MMR. - pub fn peaks(&self) -> Vec { + pub fn peaks(&self) -> impl DoubleEndedIterator + '_ { let peaks_pos = peaks(self.last_pos); - peaks_pos - .into_iter() - .filter_map(|pi| { - // here we want to get from underlying hash file - // as the pos *may* have been "removed" - self.backend.get_from_file(pi) - }) - .collect() + peaks_pos.into_iter().filter_map(move |pi| { + // here we want to get from underlying hash file + // as the pos *may* have been "removed" + self.backend.get_from_file(pi) + }) } fn peak_path(&self, peak_pos: u64) -> Vec { @@ -125,11 +122,10 @@ where let rhs = peaks(self.last_pos) .into_iter() .filter(|x| *x > peak_pos) - .filter_map(|x| self.backend.get_from_file(x)) - .collect::>(); + .filter_map(|x| self.backend.get_from_file(x)); let mut res = None; - for peak in rhs.into_iter().rev() { + for peak in rhs.rev() { res = match res { None => Some(peak), Some(rhash) => Some((peak, rhash).hash_with_index(self.unpruned_size())), @@ -145,7 +141,7 @@ where return Ok(ZERO_HASH); } let mut res = None; - for peak in self.peaks().into_iter().rev() { + for peak in self.peaks().rev() { res = match res { None => Some(peak), Some(rhash) => Some((peak, rhash).hash_with_index(self.unpruned_size())), @@ -538,17 +534,46 @@ pub fn is_left_sibling(pos: u64) -> bool { /// corresponding peak in the MMR. /// The size (and therefore the set of peaks) of the MMR /// is defined by last_pos. -pub fn path(pos: u64, last_pos: u64) -> Vec { - let (peak_map, height) = peak_map_height(pos - 1); - let mut peak = 1 << height; - let mut path = vec![]; - let mut current = pos; - while current <= last_pos { - path.push(current); - current += if (peak_map & peak) != 0 { 1 } else { 2 * peak }; - peak <<= 1; +pub fn path(pos: u64, last_pos: u64) -> impl Iterator { + Path::new(pos, last_pos) +} + +struct Path { + current: u64, + last_pos: u64, + peak: u64, + peak_map: u64, +} + +impl Path { + fn new(pos: u64, last_pos: u64) -> Self { + let (peak_map, height) = peak_map_height(pos - 1); + Path { + current: pos, + peak: 1 << height, + peak_map, + last_pos, + } + } +} + +impl Iterator for Path { + type Item = u64; + + fn next(&mut self) -> Option { + if self.current > self.last_pos { + return None; + } + + let next = Some(self.current); + self.current += if (self.peak_map & self.peak) != 0 { + 1 + } else { + 2 * self.peak + }; + self.peak <<= 1; + next } - path } /// For a given starting position calculate the parent and sibling positions diff --git a/core/tests/merkle_proof.rs b/core/tests/merkle_proof.rs index 293fea3cc..dfcf3432c 100644 --- a/core/tests/merkle_proof.rs +++ b/core/tests/merkle_proof.rs @@ -90,7 +90,7 @@ fn pmmr_merkle_proof() { assert_eq!(pmmr.get_hash(3).unwrap(), pos_2); assert_eq!(pmmr.root().unwrap(), pos_2); - assert_eq!(pmmr.peaks(), [pos_2]); + assert_eq!(pmmr.peaks().collect::>(), [pos_2]); // single peak, path with single sibling let proof = pmmr.merkle_proof(1).unwrap(); @@ -107,7 +107,7 @@ fn pmmr_merkle_proof() { assert_eq!(pmmr.get_hash(4).unwrap(), pos_3); assert_eq!(pmmr.root().unwrap(), (pos_2, pos_3).hash_with_index(4)); - assert_eq!(pmmr.peaks(), [pos_2, pos_3]); + assert_eq!(pmmr.peaks().collect::>(), [pos_2, pos_3]); let proof = pmmr.merkle_proof(1).unwrap(); assert_eq!(proof.path, vec![pos_1, pos_3]); diff --git a/core/tests/pmmr.rs b/core/tests/pmmr.rs index 4ad84c555..1e02ea135 100644 --- a/core/tests/pmmr.rs +++ b/core/tests/pmmr.rs @@ -150,10 +150,9 @@ fn various_families() { #[test] fn test_paths() { - assert_eq!(pmmr::path(1, 1), [1]); - assert_eq!(pmmr::path(1, 3), [1, 3]); - assert_eq!(pmmr::path(2, 3), [2, 3]); - assert_eq!(pmmr::path(4, 16), [4, 6, 7, 15]); + assert_eq!(pmmr::path(1, 3).collect::>(), [1, 3]); + assert_eq!(pmmr::path(2, 3).collect::>(), [2, 3]); + assert_eq!(pmmr::path(4, 16).collect::>(), [4, 6, 7, 15]); } #[test] @@ -279,7 +278,7 @@ fn pmmr_push_root() { pmmr.push(&elems[0]).unwrap(); pmmr.dump(false); let pos_0 = elems[0].hash_with_index(0); - assert_eq!(pmmr.peaks(), vec![pos_0]); + assert_eq!(pmmr.peaks().collect::>(), vec![pos_0]); assert_eq!(pmmr.root().unwrap(), pos_0); assert_eq!(pmmr.unpruned_size(), 1); @@ -288,7 +287,7 @@ fn pmmr_push_root() { pmmr.dump(false); let pos_1 = elems[1].hash_with_index(1); let pos_2 = (pos_0, pos_1).hash_with_index(2); - assert_eq!(pmmr.peaks(), vec![pos_2]); + assert_eq!(pmmr.peaks().collect::>(), vec![pos_2]); assert_eq!(pmmr.root().unwrap(), pos_2); assert_eq!(pmmr.unpruned_size(), 3); @@ -296,7 +295,7 @@ fn pmmr_push_root() { pmmr.push(&elems[2]).unwrap(); pmmr.dump(false); let pos_3 = elems[2].hash_with_index(3); - assert_eq!(pmmr.peaks(), vec![pos_2, pos_3]); + assert_eq!(pmmr.peaks().collect::>(), vec![pos_2, pos_3]); assert_eq!(pmmr.root().unwrap(), (pos_2, pos_3).hash_with_index(4)); assert_eq!(pmmr.unpruned_size(), 4); @@ -306,7 +305,7 @@ fn pmmr_push_root() { let pos_4 = elems[3].hash_with_index(4); let pos_5 = (pos_3, pos_4).hash_with_index(5); let pos_6 = (pos_2, pos_5).hash_with_index(6); - assert_eq!(pmmr.peaks(), vec![pos_6]); + assert_eq!(pmmr.peaks().collect::>(), vec![pos_6]); assert_eq!(pmmr.root().unwrap(), pos_6); assert_eq!(pmmr.unpruned_size(), 7); @@ -314,7 +313,7 @@ fn pmmr_push_root() { pmmr.push(&elems[4]).unwrap(); pmmr.dump(false); let pos_7 = elems[4].hash_with_index(7); - assert_eq!(pmmr.peaks(), vec![pos_6, pos_7]); + assert_eq!(pmmr.peaks().collect::>(), vec![pos_6, pos_7]); assert_eq!(pmmr.root().unwrap(), (pos_6, pos_7).hash_with_index(8)); assert_eq!(pmmr.unpruned_size(), 8); @@ -322,14 +321,14 @@ fn pmmr_push_root() { pmmr.push(&elems[5]).unwrap(); let pos_8 = elems[5].hash_with_index(8); let pos_9 = (pos_7, pos_8).hash_with_index(9); - assert_eq!(pmmr.peaks(), vec![pos_6, pos_9]); + assert_eq!(pmmr.peaks().collect::>(), vec![pos_6, pos_9]); assert_eq!(pmmr.root().unwrap(), (pos_6, pos_9).hash_with_index(10)); assert_eq!(pmmr.unpruned_size(), 10); // seven elements pmmr.push(&elems[6]).unwrap(); let pos_10 = elems[6].hash_with_index(10); - assert_eq!(pmmr.peaks(), vec![pos_6, pos_9, pos_10]); + assert_eq!(pmmr.peaks().collect::>(), vec![pos_6, pos_9, pos_10]); assert_eq!( pmmr.root().unwrap(), (pos_6, (pos_9, pos_10).hash_with_index(11)).hash_with_index(11) @@ -343,14 +342,14 @@ fn pmmr_push_root() { let pos_12 = (pos_10, pos_11).hash_with_index(12); let pos_13 = (pos_9, pos_12).hash_with_index(13); let pos_14 = (pos_6, pos_13).hash_with_index(14); - assert_eq!(pmmr.peaks(), vec![pos_14]); + assert_eq!(pmmr.peaks().collect::>(), vec![pos_14]); assert_eq!(pmmr.root().unwrap(), pos_14); assert_eq!(pmmr.unpruned_size(), 15); // nine elements pmmr.push(&elems[8]).unwrap(); let pos_15 = elems[8].hash_with_index(15); - assert_eq!(pmmr.peaks(), vec![pos_14, pos_15]); + assert_eq!(pmmr.peaks().collect::>(), vec![pos_14, pos_15]); assert_eq!(pmmr.root().unwrap(), (pos_14, pos_15).hash_with_index(16)); assert_eq!(pmmr.unpruned_size(), 16); } diff --git a/store/src/prune_list.rs b/store/src/prune_list.rs index 14cd92fc0..64e27ce63 100644 --- a/store/src/prune_list.rs +++ b/store/src/prune_list.rs @@ -279,8 +279,7 @@ impl PruneList { let maximum = self.bitmap.maximum().unwrap_or(0); self.pruned_cache = Bitmap::create_with_capacity(maximum); for pos in 1..(maximum + 1) { - let path = path(pos as u64, maximum as u64); - let pruned = path.into_iter().any(|x| self.bitmap.contains(x as u32)); + let pruned = path(pos as u64, maximum as u64).any(|x| self.bitmap.contains(x as u32)); if pruned { self.pruned_cache.add(pos as u32) }