PMMR Backend Support for append_pruned_root (Continued)

chain/tests/mine_simple_chain.rs

@@ -754,6 +754,7 @@ fn output_header_mappings() {
 	global::set_local_chain_type(ChainTypes::AutomatedTesting);
 	util::init_test_logger();
 	{
+		clean_output_dir(".grin_header_for_output");
 		let chain = init_chain(
 			".grin_header_for_output",
 			pow::mine_genesis_block().unwrap(),

core/src/core/pmmr/backend.rs

@@ -29,6 +29,10 @@ pub trait Backend<T: PMMRable> {
 	/// help the implementation.
 	fn append(&mut self, data: &T, hashes: &[Hash]) -> Result<(), String>;
 
+	/// Rebuilding a PMMR locally from PIBD segments requires pruned subtree support.
+	/// This allows us to append an existing pruned subtree directly without the underlying leaf nodes.
+	fn append_pruned_subtree(&mut self, hash: Hash, pos: u64) -> Result<(), String>;
+
 	/// Rewind the backend state to a previous position, as if all append
 	/// operations after that had been canceled. Expects a position in the PMMR
 	/// to rewind to as well as bitmaps representing the positions added and

core/src/core/pmmr/pmmr.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::{marker, ops::Range, u64};
+use std::{cmp::max, marker, ops::Range, u64};
 
 use croaring::Bitmap;
 
@@ -568,6 +568,9 @@ pub fn bintree_postorder_height(num: u64) -> u64 {
 /// of any size (somewhat unintuitively but this is how the PMMR is "append
 /// only").
 pub fn is_leaf(pos: u64) -> bool {
+	if pos == 0 {
+		return false;
+	}
 	bintree_postorder_height(pos) == 0
 }
 
@@ -665,14 +668,27 @@ pub fn family_branch(pos: u64, last_pos: u64) -> Vec<(u64, u64)> {
 }
 
 /// Gets the position of the rightmost node (i.e. leaf) beneath the provided subtree root.
-pub fn bintree_rightmost(num: u64) -> u64 {
-	num - bintree_postorder_height(num)
+pub fn bintree_rightmost(pos: u64) -> u64 {
+	pos - bintree_postorder_height(pos)
 }
 
 /// Gets the position of the rightmost node (i.e. leaf) beneath the provided subtree root.
-pub fn bintree_leftmost(num: u64) -> u64 {
-	let height = bintree_postorder_height(num);
-	num + 2 - (2 << height)
+pub fn bintree_leftmost(pos: u64) -> u64 {
+	let height = bintree_postorder_height(pos);
+	pos + 2 - (2 << height)
+}
+
+/// Iterator over all leaf pos beneath the provided subtree root (including the root itself).
+pub fn bintree_leaf_pos_iter(pos: u64) -> impl Iterator<Item = u64> {
+	let leaf_start = max(1, bintree_leftmost(pos as u64));
+	let leaf_end = bintree_rightmost(pos as u64);
+	(leaf_start..=leaf_end).filter(|x| is_leaf(*x))
+}
+
+/// Iterator over all pos beneath the provided subtree root (including the root itself).
+pub fn bintree_pos_iter(pos: u64) -> impl Iterator<Item = u64> {
+	let leaf_start = max(1, bintree_leftmost(pos as u64));
+	(leaf_start..=pos).into_iter()
 }
 
 /// All pos in the subtree beneath the provided root, including root itself.

core/src/core/pmmr/vec_backend.rs

@@ -43,6 +43,10 @@ impl<T: PMMRable> Backend<T> for VecBackend<T> {
 		Ok(())
 	}
 
+	fn append_pruned_subtree(&mut self, _hash: Hash, _pos: u64) -> Result<(), String> {
+		unimplemented!()
+	}
+
 	fn get_hash(&self, position: u64) -> Option<Hash> {
 		if self.removed.contains(&position) {
 			None

core/tests/pmmr.rs

@@ -20,7 +20,6 @@ use self::core::ser::PMMRIndexHashable;
 use crate::common::TestElem;
 use chrono::prelude::Utc;
 use grin_core as core;
-use std::u64;
 
 #[test]
 fn some_peak_map() {
@@ -128,6 +127,48 @@ fn test_bintree_leftmost() {
 	assert_eq!(pmmr::bintree_leftmost(7), 1);
 }
 
+#[test]
+fn test_bintree_leaf_pos_iter() {
+	assert_eq!(pmmr::bintree_leaf_pos_iter(0).count(), 0);
+	assert_eq!(pmmr::bintree_leaf_pos_iter(1).collect::<Vec<_>>(), [1]);
+	assert_eq!(pmmr::bintree_leaf_pos_iter(2).collect::<Vec<_>>(), [2]);
+	assert_eq!(pmmr::bintree_leaf_pos_iter(3).collect::<Vec<_>>(), [1, 2]);
+	assert_eq!(pmmr::bintree_leaf_pos_iter(4).collect::<Vec<_>>(), [4]);
+	assert_eq!(pmmr::bintree_leaf_pos_iter(5).collect::<Vec<_>>(), [5]);
+	assert_eq!(pmmr::bintree_leaf_pos_iter(6).collect::<Vec<_>>(), [4, 5]);
+	assert_eq!(
+		pmmr::bintree_leaf_pos_iter(7).collect::<Vec<_>>(),
+		[1, 2, 4, 5]
+	);
+}
+
+#[test]
+fn test_bintree_pos_iter() {
+	assert_eq!(pmmr::bintree_pos_iter(0).count(), 0);
+	assert_eq!(pmmr::bintree_pos_iter(1).collect::<Vec<_>>(), [1]);
+	assert_eq!(pmmr::bintree_pos_iter(2).collect::<Vec<_>>(), [2]);
+	assert_eq!(pmmr::bintree_pos_iter(3).collect::<Vec<_>>(), [1, 2, 3]);
+	assert_eq!(pmmr::bintree_pos_iter(4).collect::<Vec<_>>(), [4]);
+	assert_eq!(pmmr::bintree_pos_iter(5).collect::<Vec<_>>(), [5]);
+	assert_eq!(pmmr::bintree_pos_iter(6).collect::<Vec<_>>(), [4, 5, 6]);
+	assert_eq!(
+		pmmr::bintree_pos_iter(7).collect::<Vec<_>>(),
+		[1, 2, 3, 4, 5, 6, 7]
+	);
+}
+
+#[test]
+fn test_is_leaf() {
+	assert_eq!(pmmr::is_leaf(0), false);
+	assert_eq!(pmmr::is_leaf(1), true);
+	assert_eq!(pmmr::is_leaf(2), true);
+	assert_eq!(pmmr::is_leaf(3), false);
+	assert_eq!(pmmr::is_leaf(4), true);
+	assert_eq!(pmmr::is_leaf(5), true);
+	assert_eq!(pmmr::is_leaf(6), false);
+	assert_eq!(pmmr::is_leaf(7), false);
+}
+
 #[test]
 fn test_n_leaves() {
 	// make sure we handle an empty MMR correctly

store/src/pmmr.rs

@@ -65,7 +65,6 @@ pub struct PMMRBackend<T: PMMRable> {
 impl<T: PMMRable> Backend<T> for PMMRBackend<T> {
 	/// Append the provided data and hashes to the backend storage.
 	/// Add the new leaf pos to our leaf_set if this is a prunable MMR.
-	#[allow(unused_variables)]
 	fn append(&mut self, data: &T, hashes: &[Hash]) -> Result<(), String> {
 		let size = self
 			.data_file
@@ -86,6 +85,22 @@ impl<T: PMMRable> Backend<T> for PMMRBackend<T> {
 		Ok(())
 	}
 
+	// Supports appending a pruned subtree (single root hash) to an existing hash file.
+	// Update the prune_list "shift cache" to reflect the new pruned leaf pos in the subtree.
+	fn append_pruned_subtree(&mut self, hash: Hash, pos: u64) -> Result<(), String> {
+		if !self.prunable {
+			return Err("Not prunable, cannot append pruned subtree.".into());
+		}
+
+		self.hash_file
+			.append(&hash)
+			.map_err(|e| format!("Failed to append subtree hash to file. {}", e))?;
+
+		self.prune_list.append(pos);
+
+		Ok(())
+	}
+
 	fn get_from_file(&self, position: u64) -> Option<Hash> {
 		if self.is_compacted(position) {
 			return None;
@@ -402,9 +417,9 @@ impl<T: PMMRable> PMMRBackend<T> {
 
 		// Update the prune list and write to disk.
 		{
-			for pos in leaves_removed.iter() {
-				self.prune_list.add(pos.into());
-			}
+			let mut bitmap = self.prune_list.bitmap();
+			bitmap.or_inplace(&leaves_removed);
+			self.prune_list = PruneList::new(Some(self.data_dir.join(PMMR_PRUN_FILE)), bitmap);
 			self.prune_list.flush()?;
 		}