diff --git a/store/src/prune_list.rs b/store/src/prune_list.rs index e00b4f00d4..df3911156f 100644 --- a/store/src/prune_list.rs +++ b/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, 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>(path: P) -> io::Result { 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 { - 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() - } -} diff --git a/store/tests/prune_list.rs b/store/tests/prune_list.rs index ab6292d5bb..5be78f8132 100644 --- a/store/tests/prune_list.rs +++ b/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);