Skip to content

Commit

Permalink
Rollup merge of rust-lang#106918 - dtolnay:heapretain, r=the8472
Browse files Browse the repository at this point in the history
Rebuild BinaryHeap on unwind from retain

This closes the hole identified in rust-lang#71503 (comment) which had made it possible for the caller to end up with a heap in invalid state. As of rust-lang#105851, heaps in invalid state are not supposed to exist.
  • Loading branch information
Dylan-DPC authored Feb 24, 2023
2 parents 8c135ee + fa2ff4d commit b3657f9
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 6 deletions.
24 changes: 18 additions & 6 deletions library/alloc/src/collections/binary_heap/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
}
}

Expand Down
19 changes: 19 additions & 0 deletions library/alloc/src/collections/binary_heap/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit b3657f9

Please sign in to comment.