Rollup merge of #106918 - dtolnay:heapretain, r=the8472

Rebuild BinaryHeap on unwind from retain

This closes the hole identified in https://github.com/rust-lang/rust/issues/71503#issuecomment-1383251315 which had made it possible for the caller to end up with a heap in invalid state. As of #105851, heaps in invalid state are not supposed to exist.

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 0b73b1af4eb..f1d0a305d99 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 ffbb6c80ac0..500caa35678 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,