Skip to content

Commit d787881

Browse files
Fix Itertools::k_smallest on unfused iterators
In the case of an unfused iterator yielding at most k-1 elements, then None, then other elements, the heap is not k-long and `for_each` is called on the other elements leading `debug_assert_eq` to rightfully panic. Our recent variants of `k_smallest` do not have this bug!
1 parent 8ed734b commit d787881

File tree

1 file changed

+11
-9
lines changed

1 file changed

+11
-9
lines changed

src/lib.rs

+11-9
Original file line numberDiff line numberDiff line change
@@ -2965,15 +2965,17 @@ pub trait Itertools: Iterator {
29652965

29662966
let mut heap = self.by_ref().take(k).collect::<BinaryHeap<_>>();
29672967

2968-
self.for_each(|i| {
2969-
debug_assert_eq!(heap.len(), k);
2970-
// Equivalent to heap.push(min(i, heap.pop())) but more efficient.
2971-
// This should be done with a single `.peek_mut().unwrap()` but
2972-
// `PeekMut` sifts-down unconditionally on Rust 1.46.0 and prior.
2973-
if *heap.peek().unwrap() > i {
2974-
*heap.peek_mut().unwrap() = i;
2975-
}
2976-
});
2968+
if heap.len() == k {
2969+
self.for_each(|i| {
2970+
debug_assert_eq!(heap.len(), k);
2971+
// Equivalent to heap.push(min(i, heap.pop())) but more efficient.
2972+
// This should be done with a single `.peek_mut().unwrap()` but
2973+
// `PeekMut` sifts-down unconditionally on Rust 1.46.0 and prior.
2974+
if *heap.peek().unwrap() > i {
2975+
*heap.peek_mut().unwrap() = i;
2976+
}
2977+
});
2978+
}
29772979

29782980
heap.into_sorted_vec().into_iter()
29792981
}

0 commit comments

Comments
 (0)