summaryrefslogtreecommitdiffstats
path: root/library/alloc/src/collections/binary_heap
diff options
context:
space:
mode:
Diffstat (limited to 'library/alloc/src/collections/binary_heap')
-rw-r--r--library/alloc/src/collections/binary_heap/mod.rs24
-rw-r--r--library/alloc/src/collections/binary_heap/tests.rs19
2 files changed, 37 insertions, 6 deletions
diff --git a/library/alloc/src/collections/binary_heap/mod.rs b/library/alloc/src/collections/binary_heap/mod.rs
index 0b73b1af4..f1d0a305d 100644
--- a/library/alloc/src/collections/binary_heap/mod.rs
+++ b/library/alloc/src/collections/binary_heap/mod.rs
@@ -851,18 +851,30 @@ impl<T: Ord> BinaryHeap<T> {
where
F: FnMut(&T) -> bool,
{
- let mut first_removed = self.len();
+ struct RebuildOnDrop<'a, T: Ord> {
+ heap: &'a mut BinaryHeap<T>,
+ first_removed: usize,
+ }
+
+ let mut guard = RebuildOnDrop { first_removed: self.len(), heap: self };
+
let mut i = 0;
- self.data.retain(|e| {
+ guard.heap.data.retain(|e| {
let keep = f(e);
- if !keep && i < first_removed {
- first_removed = i;
+ if !keep && i < guard.first_removed {
+ guard.first_removed = i;
}
i += 1;
keep
});
- // data[0..first_removed] is untouched, so we only need to rebuild the tail:
- self.rebuild_tail(first_removed);
+
+ impl<'a, T: Ord> Drop for RebuildOnDrop<'a, T> {
+ fn drop(&mut self) {
+ // data[..first_removed] is untouched, so we only need to
+ // rebuild the tail:
+ self.heap.rebuild_tail(self.first_removed);
+ }
+ }
}
}
diff --git a/library/alloc/src/collections/binary_heap/tests.rs b/library/alloc/src/collections/binary_heap/tests.rs
index ffbb6c80a..500caa356 100644
--- a/library/alloc/src/collections/binary_heap/tests.rs
+++ b/library/alloc/src/collections/binary_heap/tests.rs
@@ -474,6 +474,25 @@ fn test_retain() {
assert!(a.is_empty());
}
+#[test]
+fn test_retain_catch_unwind() {
+ let mut heap = BinaryHeap::from(vec![3, 1, 2]);
+
+ // Removes the 3, then unwinds out of retain.
+ let _ = catch_unwind(AssertUnwindSafe(|| {
+ heap.retain(|e| {
+ if *e == 1 {
+ panic!();
+ }
+ false
+ });
+ }));
+
+ // Naively this would be [1, 2] (an invalid heap) if BinaryHeap delegates to
+ // Vec's retain impl and then does not rebuild the heap after that unwinds.
+ assert_eq!(heap.into_vec(), [2, 1]);
+}
+
// old binaryheap failed this test
//
// Integrity means that all elements are present after a comparison panics,