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 BinaryHeap { where F: FnMut(&T) -> bool, { - let mut first_removed = self.len(); + struct RebuildOnDrop<'a, T: Ord> { + heap: &'a mut BinaryHeap, + 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