Prune list robustness

store/src/prune_list.rs

@@ -51,17 +51,25 @@ pub struct PruneList {
 }
 
 impl PruneList {
-	/// Instantiate a new empty prune list
-	pub fn new() -> PruneList {
+	/// Instantiate a new prune list from the provided path and bitmap.
+	pub fn new(path: Option<PathBuf>, mut bitmap: Bitmap) -> PruneList {
+		// Note: prune list is 1-indexed so remove any 0 value for safety.
+		bitmap.remove(0);
+
 		PruneList {
-			path: None,
-			bitmap: Bitmap::create(),
+			path,
+			bitmap,
 			pruned_cache: Bitmap::create(),
 			shift_cache: vec![],
 			leaf_shift_cache: vec![],
 		}
 	}
 
+	/// Instatiate a new empty prune list.
+	pub fn empty() -> PruneList {
+		PruneList::new(None, Bitmap::create())
+	}
+
 	/// Open an existing prune_list or create a new one.
 	pub fn open<P: AsRef<Path>>(path: P) -> io::Result<PruneList> {
 		let file_path = PathBuf::from(path.as_ref());
@@ -71,13 +79,7 @@ impl PruneList {
 			Bitmap::create()
 		};
 
-		let mut prune_list = PruneList {
-			path: Some(file_path),
-			bitmap,
-			pruned_cache: Bitmap::create(),
-			shift_cache: vec![],
-			leaf_shift_cache: vec![],
-		};
+		let mut prune_list = PruneList::new(Some(file_path), bitmap);
 
 		// Now built the shift and pruned caches from the bitmap we read from disk.
 		prune_list.init_caches();
@@ -96,7 +98,8 @@ impl PruneList {
 		Ok(prune_list)
 	}
 
-	fn init_caches(&mut self) {
+	/// Init our internal shift caches.
+	pub fn init_caches(&mut self) {
 		self.build_shift_cache();
 		self.build_leaf_shift_cache();
 		self.build_pruned_cache();
@@ -152,9 +155,9 @@ impl PruneList {
 		}
 
 		if idx > self.shift_cache.len() as u64 {
-			self.shift_cache[self.shift_cache.len() - 1]
+			self.shift_cache[self.shift_cache.len().saturating_sub(1)]
 		} else {
-			self.shift_cache[idx as usize - 1]
+			self.shift_cache[(idx as usize).saturating_sub(1)]
 		}
 	}
 
@@ -164,9 +167,9 @@ impl PruneList {
 		}
 
 		self.shift_cache.clear();
-		for pos in self.bitmap.iter() {
+		for pos in self.bitmap.iter().filter(|x| *x > 0) {
 			let pos = pos as u64;
-			let prev_shift = self.get_shift(pos - 1);
+			let prev_shift = self.get_shift(pos.saturating_sub(1));
 
 			let curr_shift = if self.is_pruned_root(pos) {
 				let height = bintree_postorder_height(pos);
@@ -193,9 +196,9 @@ impl PruneList {
 		}
 
 		if idx > self.leaf_shift_cache.len() as u64 {
-			self.leaf_shift_cache[self.leaf_shift_cache.len() - 1]
+			self.leaf_shift_cache[self.leaf_shift_cache.len().saturating_sub(1)]
 		} else {
-			self.leaf_shift_cache[idx as usize - 1]
+			self.leaf_shift_cache[(idx as usize).saturating_sub(1)]
 		}
 	}
 
@@ -206,9 +209,9 @@ impl PruneList {
 
 		self.leaf_shift_cache.clear();
 
-		for pos in self.bitmap.iter() {
+		for pos in self.bitmap.iter().filter(|x| *x > 0) {
 			let pos = pos as u64;
-			let prev_shift = self.get_leaf_shift(pos - 1);
+			let prev_shift = self.get_leaf_shift(pos.saturating_sub(1));
 
 			let curr_shift = if self.is_pruned_root(pos) {
 				let height = bintree_postorder_height(pos);
@@ -229,6 +232,8 @@ impl PruneList {
 	/// list if pruning the additional node means a parent can get pruned as
 	/// well.
 	pub fn add(&mut self, pos: u64) {
+		assert!(pos > 0, "prune list 1-indexed, 0 not valid pos");
+
 		let mut current = pos;
 		loop {
 			let (parent, sibling) = family(current);
@@ -257,12 +262,13 @@ impl PruneList {
 
 	/// Convert the prune_list to a vec of pos.
 	pub fn to_vec(&self) -> Vec<u64> {
-		self.bitmap.to_vec().into_iter().map(|x| x as u64).collect()
+		self.bitmap.iter().map(|x| x as u64).collect()
 	}
 
 	/// Is the pos pruned?
 	/// Assumes the pruned_cache is fully built and up to date.
 	pub fn is_pruned(&self, pos: u64) -> bool {
+		assert!(pos > 0, "prune list 1-indexed, 0 not valid pos");
 		self.pruned_cache.contains(pos as u32)
 	}
 
@@ -283,12 +289,7 @@ impl PruneList {
 
 	/// Is the specified position a root of a pruned subtree?
 	pub fn is_pruned_root(&self, pos: u64) -> bool {
+		assert!(pos > 0, "prune list 1-indexed, 0 not valid pos");
 		self.bitmap.contains(pos as u32)
 	}
 }
-
-impl Default for PruneList {
-	fn default() -> Self {
-		Self::new()
-	}
-}

store/tests/prune_list.rs

@@ -15,10 +15,26 @@
 use grin_store as store;
 
 use crate::store::prune_list::PruneList;
+use croaring::Bitmap;
+
+// Prune list is 1-indexed but we implement this internally with a bitmap that supports a 0 value.
+// We need to make sure we safely handle 0 safely.
+#[test]
+fn test_zero_value() {
+	// Create a bitmap with a 0 value in it.
+	let mut bitmap = Bitmap::create();
+	bitmap.add(0);
+
+	// Instantiate a prune list from our existing bitmap.
+	let pl = PruneList::new(None, bitmap);
+
+	// Our prune list should be empty (0 filtered out during creation).
+	assert!(pl.is_empty());
+}
 
 #[test]
 fn test_is_pruned() {
-	let mut pl = PruneList::new();
+	let mut pl = PruneList::empty();
 
 	assert_eq!(pl.len(), 0);
 	assert_eq!(pl.is_pruned(1), false);
@@ -62,7 +78,7 @@ fn test_is_pruned() {
 
 #[test]
 fn test_get_leaf_shift() {
-	let mut pl = PruneList::new();
+	let mut pl = PruneList::empty();
 
 	// start with an empty prune list (nothing shifted)
 	assert_eq!(pl.len(), 0);
@@ -135,7 +151,7 @@ fn test_get_leaf_shift() {
 
 	// now check we can prune some unconnected nodes in arbitrary order
 	// and that leaf_shift is correct for various pos
-	let mut pl = PruneList::new();
+	let mut pl = PruneList::empty();
 	pl.add(5);
 	pl.add(11);
 	pl.add(12);
@@ -154,7 +170,7 @@ fn test_get_leaf_shift() {
 
 #[test]
 fn test_get_shift() {
-	let mut pl = PruneList::new();
+	let mut pl = PruneList::empty();
 	assert!(pl.is_empty());
 	assert_eq!(pl.get_shift(1), 0);
 	assert_eq!(pl.get_shift(2), 0);
@@ -231,7 +247,7 @@ fn test_get_shift() {
 	// and check we shift by a large number (hopefully the correct number...)
 	assert_eq!(pl.get_shift(1010), 996);
 
-	let mut pl = PruneList::new();
+	let mut pl = PruneList::empty();
 	pl.add(9);
 	pl.add(8);
 	pl.add(5);