diff --git a/src/liballoc/collections/binary_heap.rs b/src/liballoc/collections/binary_heap.rs index c527b378f74..f38fe997b73 100644 --- a/src/liballoc/collections/binary_heap.rs +++ b/src/liballoc/collections/binary_heap.rs @@ -147,7 +147,7 @@ use core::fmt; use core::iter::{FromIterator, FusedIterator, TrustedLen}; -use core::mem::{size_of, swap, ManuallyDrop}; +use core::mem::{self, size_of, swap, ManuallyDrop}; use core::ops::{Deref, DerefMut}; use core::ptr; @@ -1239,7 +1239,19 @@ pub struct DrainSorted<'a, T: Ord> { impl<'a, T: Ord> Drop for DrainSorted<'a, T> { /// Removes heap elements in heap order. fn drop(&mut self) { - while let Some(_) = self.inner.pop() {} + struct DropGuard<'r, 'a, T: Ord>(&'r mut DrainSorted<'a, T>); + + impl<'r, 'a, T: Ord> Drop for DropGuard<'r, 'a, T> { + fn drop(&mut self) { + while let Some(_) = self.0.inner.pop() {} + } + } + + while let Some(item) = self.inner.pop() { + let guard = DropGuard(self); + drop(item); + mem::forget(guard); + } } } diff --git a/src/liballoc/tests/binary_heap.rs b/src/liballoc/tests/binary_heap.rs index f49ca713921..be5516f54f3 100644 --- a/src/liballoc/tests/binary_heap.rs +++ b/src/liballoc/tests/binary_heap.rs @@ -1,6 +1,8 @@ use std::collections::binary_heap::{Drain, PeekMut}; use std::collections::BinaryHeap; use std::iter::TrustedLen; +use std::panic::{catch_unwind, AssertUnwindSafe}; +use std::sync::atomic::{AtomicU32, Ordering}; #[test] fn test_iterator() { @@ -275,6 +277,37 @@ fn test_drain_sorted() { assert!(q.is_empty()); } +#[test] +fn test_drain_sorted_leak() { + static DROPS: AtomicU32 = AtomicU32::new(0); + + #[derive(Clone, PartialEq, Eq, PartialOrd, Ord)] + struct D(u32, bool); + + impl Drop for D { + fn drop(&mut self) { + DROPS.fetch_add(1, Ordering::SeqCst); + + if self.1 { + panic!("panic in `drop`"); + } + } + } + + let mut q = BinaryHeap::from(vec![ + D(0, false), + D(1, false), + D(2, false), + D(3, true), + D(4, false), + D(5, false), + ]); + + catch_unwind(AssertUnwindSafe(|| drop(q.drain_sorted()))).ok(); + + assert_eq!(DROPS.load(Ordering::SeqCst), 6); +} + #[test] fn test_extend_ref() { let mut a = BinaryHeap::new();