From 0d3eaa848ce6b44a83502174eae16524c893dd28 Mon Sep 17 00:00:00 2001
From: David Tolnay <dtolnay@gmail.com>
Date: Sun, 15 Jan 2023 13:08:20 -0800
Subject: [PATCH 1/2] Add test showing broken behavior of BinaryHeap::retain

---
 .../alloc/src/collections/binary_heap/tests.rs  | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/library/alloc/src/collections/binary_heap/tests.rs b/library/alloc/src/collections/binary_heap/tests.rs
index ffbb6c80ac018..835b2e5b7e917 100644
--- a/library/alloc/src/collections/binary_heap/tests.rs
+++ b/library/alloc/src/collections/binary_heap/tests.rs
@@ -474,6 +474,23 @@ 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
+        });
+    }));
+
+    assert_eq!(heap.into_vec(), [1, 2]); // BAD!!
+}
+
 // old binaryheap failed this test
 //
 // Integrity means that all elements are present after a comparison panics,

From fa2ff4d7e5a88bfda186f0f3f621402470745788 Mon Sep 17 00:00:00 2001
From: David Tolnay <dtolnay@gmail.com>
Date: Sun, 15 Jan 2023 13:14:33 -0800
Subject: [PATCH 2/2] Rebuild BinaryHeap on unwind from retain

---
 .../alloc/src/collections/binary_heap/mod.rs  | 24 ++++++++++++++-----
 .../src/collections/binary_heap/tests.rs      |  4 +++-
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/library/alloc/src/collections/binary_heap/mod.rs b/library/alloc/src/collections/binary_heap/mod.rs
index 0b73b1af4eb35..f1d0a305d999f 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 835b2e5b7e917..500caa35678ab 100644
--- a/library/alloc/src/collections/binary_heap/tests.rs
+++ b/library/alloc/src/collections/binary_heap/tests.rs
@@ -488,7 +488,9 @@ fn test_retain_catch_unwind() {
         });
     }));
 
-    assert_eq!(heap.into_vec(), [1, 2]); // BAD!!
+    // 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