use take_while so unpruned iterators are not infinite

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<Item = u64> + '_ {
+	pub fn unpruned_iter(&self, cutoff_pos: u64) -> impl Iterator<Item = u64> + '_ {
 		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<Item = u64> + '_ {
-		self.unpruned_iter().filter(|x| pmmr::is_leaf(*x))
+	pub fn unpruned_leaf_iter(&self, cutoff_pos: u64) -> impl Iterator<Item = u64> + '_ {
+		self.unpruned_iter(cutoff_pos).filter(|x| pmmr::is_leaf(*x))
 	}
 }
 

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::<Vec<_>>(),
-		[1, 2, 3, 4, 5]
-	);
+	assert_eq!(pl.unpruned_iter(5).collect::<Vec<_>>(), [1, 2, 3, 4, 5]);
 
 	let mut pl = PruneList::empty();
 	pl.add(2);
 	assert_eq!(pl.iter().collect::<Vec<_>>(), [2]);
 	assert_eq!(pl.pruned_bintree_range_iter().collect::<Vec<_>>(), [2..3]);
-	assert_eq!(pl.unpruned_iter().take(3).collect::<Vec<_>>(), [1, 3, 4]);
+	assert_eq!(pl.unpruned_iter(4).collect::<Vec<_>>(), [1, 3, 4]);
 
 	let mut pl = PruneList::empty();
 	pl.add(2);
@@ -333,28 +330,22 @@ pub fn test_unpruned_iter() {
 		pl.pruned_bintree_range_iter().collect::<Vec<_>>(),
 		[2..3, 4..7]
 	);
-	assert_eq!(
-		pl.unpruned_iter().take(5).collect::<Vec<_>>(),
-		[1, 3, 7, 8, 9]
-	);
+	assert_eq!(pl.unpruned_iter(9).collect::<Vec<_>>(), [1, 3, 7, 8, 9]);
 }
 
 #[test]
 fn test_unpruned_leaf_iter() {
 	let pl = PruneList::empty();
 	assert_eq!(
-		pl.unpruned_leaf_iter().take(5).collect::<Vec<_>>(),
+		pl.unpruned_leaf_iter(8).collect::<Vec<_>>(),
 		[1, 2, 4, 5, 8]
 	);
 
 	let mut pl = PruneList::empty();
 	pl.add(2);
 	assert_eq!(pl.iter().collect::<Vec<_>>(), [2]);
 	assert_eq!(pl.pruned_bintree_range_iter().collect::<Vec<_>>(), [2..3]);
-	assert_eq!(
-		pl.unpruned_leaf_iter().take(3).collect::<Vec<_>>(),
-		[1, 4, 5]
-	);
+	assert_eq!(pl.unpruned_leaf_iter(5).collect::<Vec<_>>(), [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::<Vec<_>>(),
 		[2..3, 4..7]
 	);
-	assert_eq!(
-		pl.unpruned_leaf_iter().take(3).collect::<Vec<_>>(),
-		[1, 8, 9]
-	);
+	assert_eq!(pl.unpruned_leaf_iter(9).collect::<Vec<_>>(), [1, 8, 9]);
 }