diff --git a/store/src/prune_list.rs b/store/src/prune_list.rs index f1cc673fe4..5baa6fd86a 100644 --- a/store/src/prune_list.rs +++ b/store/src/prune_list.rs @@ -294,15 +294,16 @@ impl PruneList { } /// Iterator over all pos that are *not* pruned based on current prune_list. - pub fn unpruned_iter(&self) -> impl Iterator + '_ { + pub fn unpruned_iter(&self, cutoff_pos: u64) -> impl Iterator + '_ { UnprunedIterator::new(self.pruned_bintree_range_iter()) + .take_while(move |x| *x <= cutoff_pos) } /// Iterator over all leaf pos that are *not* pruned based on current prune_list. /// Note this is not necessarily the same as the "leaf_set" as an output /// can be spent but not yet pruned. - pub fn unpruned_leaf_iter(&self) -> impl Iterator + '_ { - self.unpruned_iter().filter(|x| pmmr::is_leaf(*x)) + pub fn unpruned_leaf_iter(&self, cutoff_pos: u64) -> impl Iterator + '_ { + self.unpruned_iter(cutoff_pos).filter(|x| pmmr::is_leaf(*x)) } } diff --git a/store/tests/prune_list.rs b/store/tests/prune_list.rs index cb44a5a664..b56bf349fb 100644 --- a/store/tests/prune_list.rs +++ b/store/tests/prune_list.rs @@ -313,16 +313,13 @@ pub fn test_pruned_bintree_range_iter() { #[test] pub fn test_unpruned_iter() { let pl = PruneList::empty(); - assert_eq!( - pl.unpruned_iter().take(5).collect::>(), - [1, 2, 3, 4, 5] - ); + assert_eq!(pl.unpruned_iter(5).collect::>(), [1, 2, 3, 4, 5]); let mut pl = PruneList::empty(); pl.add(2); assert_eq!(pl.iter().collect::>(), [2]); assert_eq!(pl.pruned_bintree_range_iter().collect::>(), [2..3]); - assert_eq!(pl.unpruned_iter().take(3).collect::>(), [1, 3, 4]); + assert_eq!(pl.unpruned_iter(4).collect::>(), [1, 3, 4]); let mut pl = PruneList::empty(); pl.add(2); @@ -333,17 +330,14 @@ pub fn test_unpruned_iter() { pl.pruned_bintree_range_iter().collect::>(), [2..3, 4..7] ); - assert_eq!( - pl.unpruned_iter().take(5).collect::>(), - [1, 3, 7, 8, 9] - ); + assert_eq!(pl.unpruned_iter(9).collect::>(), [1, 3, 7, 8, 9]); } #[test] fn test_unpruned_leaf_iter() { let pl = PruneList::empty(); assert_eq!( - pl.unpruned_leaf_iter().take(5).collect::>(), + pl.unpruned_leaf_iter(8).collect::>(), [1, 2, 4, 5, 8] ); @@ -351,10 +345,7 @@ fn test_unpruned_leaf_iter() { pl.add(2); assert_eq!(pl.iter().collect::>(), [2]); assert_eq!(pl.pruned_bintree_range_iter().collect::>(), [2..3]); - assert_eq!( - pl.unpruned_leaf_iter().take(3).collect::>(), - [1, 4, 5] - ); + assert_eq!(pl.unpruned_leaf_iter(5).collect::>(), [1, 4, 5]); let mut pl = PruneList::empty(); pl.add(2); @@ -365,8 +356,5 @@ fn test_unpruned_leaf_iter() { pl.pruned_bintree_range_iter().collect::>(), [2..3, 4..7] ); - assert_eq!( - pl.unpruned_leaf_iter().take(3).collect::>(), - [1, 8, 9] - ); + assert_eq!(pl.unpruned_leaf_iter(9).collect::>(), [1, 8, 9]); }