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

@@ -90,16 +90,13 @@ where
 	}
 
 	/// Returns a vec of the peaks of this MMR.
-	pub fn peaks(&self) -> Vec<Hash> {
+	pub fn peaks(&self) -> impl DoubleEndedIterator<Item = Hash> + '_ {
 		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<Hash> {
@@ -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::<Vec<_>>();
+			.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<u64> {
-	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<Item = u64> {
+	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<Self::Item> {
+		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

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::<Vec<_>>(), [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::<Vec<_>>(), [pos_2, pos_3]);
 
 	let proof = pmmr.merkle_proof(1).unwrap();
 	assert_eq!(proof.path, vec![pos_1, pos_3]);

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::<Vec<_>>(), [1, 3]);
+	assert_eq!(pmmr::path(2, 3).collect::<Vec<_>>(), [2, 3]);
+	assert_eq!(pmmr::path(4, 16).collect::<Vec<_>>(), [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<_>>(), 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<_>>(), 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<_>>(), 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<_>>(), 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<_>>(), 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<_>>(), 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<_>>(), 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<_>>(), 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<_>>(), vec![pos_14, pos_15]);
 	assert_eq!(pmmr.root().unwrap(), (pos_14, pos_15).hash_with_index(16));
 	assert_eq!(pmmr.unpruned_size(), 16);
 }

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)
 			}