From 5092652b0ca4c11b52d59dfc197b41087dc04c16 Mon Sep 17 00:00:00 2001
From: Antioch Peverell <apeverell@protonmail.com>
Date: Fri, 19 Feb 2021 11:59:40 +0000
Subject: [PATCH] rework prune rewrite with iterators (#3568)

* rework prune rewrite with iterators

* fix
---
 store/src/lib.rs   |  1 -
 store/src/pmmr.rs  | 28 ++++++++++++------------
 store/src/types.rs | 54 +++++++++++++++++++++++++++++++---------------
 3 files changed, 51 insertions(+), 32 deletions(-)

diff --git a/store/src/lib.rs b/store/src/lib.rs
index ac2e0800a..f518bda0d 100644
--- a/store/src/lib.rs
+++ b/store/src/lib.rs
@@ -25,7 +25,6 @@ extern crate log;
 use failure;
 #[macro_use]
 extern crate failure_derive;
-#[macro_use]
 extern crate grin_core as core;
 extern crate grin_util as util;
 
diff --git a/store/src/pmmr.rs b/store/src/pmmr.rs
index 269afbb11..19e348f5e 100644
--- a/store/src/pmmr.rs
+++ b/store/src/pmmr.rs
@@ -354,29 +354,29 @@ impl<T: PMMRable> PMMRBackend<T> {
 
 		// 1. Save compact copy of the hash file, skipping removed data.
 		{
-			let pos_to_rm = map_vec!(pos_to_rm, |pos| {
-				let shift = self.prune_list.get_shift(pos.into());
-				pos as u64 - shift
+			let pos_to_rm = pos_to_rm.iter().map(|x| x as u64).map(|pos| {
+				let shift = self.prune_list.get_shift(pos);
+				pos - shift
 			});
 
-			self.hash_file.save_prune(&pos_to_rm)?;
+			self.hash_file.write_tmp_pruned(pos_to_rm)?;
+			self.hash_file.replace_with_tmp()?;
 		}
 
 		// 2. Save compact copy of the data file, skipping removed leaves.
 		{
-			let leaf_pos_to_rm = pos_to_rm
+			let pos_to_rm = pos_to_rm
 				.iter()
-				.filter(|&x| pmmr::is_leaf(x.into()))
 				.map(|x| x as u64)
-				.collect::<Vec<_>>();
+				.filter(|&x| pmmr::is_leaf(x))
+				.map(|pos| {
+					let flat_pos = pmmr::n_leaves(pos);
+					let shift = self.prune_list.get_leaf_shift(pos);
+					flat_pos - shift
+				});
 
-			let pos_to_rm = map_vec!(leaf_pos_to_rm, |&pos| {
-				let flat_pos = pmmr::n_leaves(pos);
-				let shift = self.prune_list.get_leaf_shift(pos);
-				flat_pos - shift
-			});
-
-			self.data_file.save_prune(&pos_to_rm)?;
+			self.data_file.write_tmp_pruned(pos_to_rm)?;
+			self.data_file.replace_with_tmp()?;
 		}
 
 		// 3. Update the prune list and write to disk.
diff --git a/store/src/types.rs b/store/src/types.rs
index df21a1719..46f231961 100644
--- a/store/src/types.rs
+++ b/store/src/types.rs
@@ -154,10 +154,18 @@ where
 	}
 
 	/// Write the file out to disk, pruning removed elements.
-	pub fn save_prune(&mut self, prune_pos: &[u64]) -> io::Result<()> {
+	pub fn write_tmp_pruned<I>(&self, prune_pos: I) -> io::Result<()>
+	where
+		I: IntoIterator<Item = u64>,
+	{
 		// Need to convert from 1-index to 0-index (don't ask).
-		let prune_idx: Vec<_> = prune_pos.iter().map(|x| x - 1).collect();
-		self.file.save_prune(prune_idx.as_slice())
+		let prune_idx = prune_pos.into_iter().map(|x| x - 1);
+		self.file.write_tmp_pruned(prune_idx)
+	}
+
+	/// Replace underlying file with the file at our tmp path.
+	pub fn replace_with_tmp(&mut self) -> io::Result<()> {
+		self.file.replace_with_tmp()
 	}
 }
 
@@ -485,39 +493,52 @@ where
 		Ok(file)
 	}
 
+	fn tmp_file_path(&self) -> PathBuf {
+		self.path.with_extension("tmp")
+	}
+
 	/// Saves a copy of the current file content, skipping data at the provided
 	/// prune positions. prune_pos must be ordered.
-	pub fn save_prune(&mut self, prune_pos: &[u64]) -> io::Result<()> {
-		let tmp_path = self.path.with_extension("tmp");
+	fn write_tmp_pruned<I>(&self, prune_pos: I) -> io::Result<()>
+	where
+		I: IntoIterator<Item = u64>,
+	{
+		let mut prune_pos = prune_pos.into_iter().peekable();
 
 		// Scope the reader and writer to within the block so we can safely replace files later on.
 		{
 			let reader = File::open(&self.path)?;
 			let mut buf_reader = BufReader::new(reader);
 			let mut streaming_reader = StreamingReader::new(&mut buf_reader, self.version);
-
-			let mut buf_writer = BufWriter::new(File::create(&tmp_path)?);
+			let tmp_path = self.tmp_file_path();
+			let mut buf_writer = BufWriter::new(File::create(tmp_path)?);
 			let mut bin_writer = BinWriter::new(&mut buf_writer, self.version);
 
 			let mut current_pos = 0;
-			let mut prune_pos = prune_pos;
 			while let Ok(elmt) = T::read(&mut streaming_reader) {
-				if prune_pos.contains(&current_pos) {
-					// Pruned pos, moving on.
-					prune_pos = &prune_pos[1..];
-				} else {
-					// Not pruned, write to file.
-					elmt.write(&mut bin_writer)
-						.map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+				if let Some(pos) = prune_pos.peek() {
+					if *pos == current_pos {
+						// Pruned pos, moving on.
+						prune_pos.next();
+						current_pos += 1;
+						continue;
+					}
 				}
+				// Not pruned, write to file.
+				elmt.write(&mut bin_writer)
+					.map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
 				current_pos += 1;
 			}
 			buf_writer.flush()?;
 		}
 
+		Ok(())
+	}
+
+	fn replace_with_tmp(&mut self) -> io::Result<()> {
 		// Replace the underlying file -
 		// pmmr_data.tmp -> pmmr_data.bin
-		self.replace(&tmp_path)?;
+		self.replace(self.tmp_file_path())?;
 
 		// Now rebuild our size file to reflect the pruned data file.
 		// This will replace the underlying file internally.
@@ -575,7 +596,6 @@ where
 	}
 
 	/// Replace the underlying file with another file, deleting the original.
-	/// Takes an optional size_file path in addition to path.
 	fn replace<P>(&mut self, with: P) -> io::Result<()>
 	where
 		P: AsRef<Path> + Debug,