PMMR Backend Support for append_pruned_root (Continued) (#3659)

* refactor prune_list with aim of allowing pruned subtree appending

* add test coverage around pmmr::is_leaf() and pmmr::bintree_leaf_pos_iter()

* comments

* cleanup

* implement append pruned subtree for prune_list

* commit

* we can now append to prune_list

* fix our prune_list corruption...

* rework how we rewrite the prune list during compaction

* test coverage for improved prune list api

* continuing to merge

* finish merge, tests passing again

* add function pmmr_leaf_to_insertion_index, and modify bintree_lef_pos_iter to use it. Note there's still an unwrap that needs to be dealt with sanely

* change pmmr_leaf_to_insertion_index to simpler version + handle conversion between 1 and 0 based in bintree_leaf_pos_iter

Co-authored-by: antiochp <[email protected]>

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, iter, marker, ops::Range, u64};
 
 use croaring::Bitmap;
 
@@ -495,6 +495,19 @@ pub fn insertion_to_pmmr_index(mut sz: u64) -> u64 {
 	2 * sz - sz.count_ones() as u64 + 1
 }
 
+/// Returns the insertion index of the given leaf index
+pub fn pmmr_leaf_to_insertion_index(pos1: u64) -> Option<u64> {
+	if pos1 == 0 {
+		return None;
+	}
+	let (insert_idx, height) = peak_map_height(pos1 - 1);
+	if height == 0 {
+		Some(insert_idx)
+	} else {
+		None
+	}
+}
+
 /// sizes of peaks and height of next node in mmr of given size
 /// Example: on input 5 returns ([3,1], 1) as mmr state before adding 5 was
 ///    2
@@ -568,6 +581,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 +681,35 @@ 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(pos1: u64) -> u64 {
+	pos1 - bintree_postorder_height(pos1)
 }
 
-/// 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)
+/// Gets the position of the leftmost node (i.e. leaf) beneath the provided subtree root.
+pub fn bintree_leftmost(pos1: u64) -> u64 {
+	let height = bintree_postorder_height(pos1);
+	pos1 + 2 - (2 << height)
+}
+
+/// Iterator over all leaf pos beneath the provided subtree root (including the root itself).
+pub fn bintree_leaf_pos_iter(pos1: u64) -> Box<dyn Iterator<Item = u64>> {
+	let leaf_start = pmmr_leaf_to_insertion_index(bintree_leftmost(pos1));
+	let leaf_end = pmmr_leaf_to_insertion_index(bintree_rightmost(pos1));
+	let leaf_start = match leaf_start {
+		Some(l) => l,
+		None => return Box::new(iter::empty::<u64>()),
+	};
+	let leaf_end = match leaf_end {
+		Some(l) => l,
+		None => return Box::new(iter::empty::<u64>()),
+	};
+	Box::new((leaf_start..=leaf_end).map(|n| insertion_to_pmmr_index(n + 1)))
+}
+
+/// Iterator over all pos beneath the provided subtree root (including the root itself).
+pub fn bintree_pos_iter(pos1: u64) -> impl Iterator<Item = u64> {
+	let leaf_start = max(1, bintree_leftmost(pos1 as u64));
+	(leaf_start..=pos1).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,80 @@ 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_pmmr_leaf_to_insertion_index() {
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(1), Some(0));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(2), Some(1));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(4), Some(2));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(5), Some(3));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(8), Some(4));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(9), Some(5));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(11), Some(6));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(12), Some(7));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(16), Some(8));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(17), Some(9));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(19), Some(10));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(20), Some(11));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(23), Some(12));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(24), Some(13));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(26), Some(14));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(27), Some(15));
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(32), Some(16));
+
+	// Not a leaf node
+	assert_eq!(pmmr::pmmr_leaf_to_insertion_index(31), None);
+
+	// Sanity check to make sure we don't get an explosion around the u64 max
+	// number of leaves
+	let n_leaves_max_u64 = pmmr::n_leaves(u64::MAX - 256);
+	assert_eq!(
+		pmmr::pmmr_leaf_to_insertion_index(n_leaves_max_u64),
+		Some(4611686018427387884)
+	);
+}
+
 #[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()?;
 		}