Rebuild BinaryHeap on unwind from retain

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 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 835b2e5b7e9..500caa35678 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