diff --git a/library/alloc/src/vec/in_place_collect.rs b/library/alloc/src/vec/in_place_collect.rs
index 55dcb84ad16f9..9f2555ee1611a 100644
--- a/library/alloc/src/vec/in_place_collect.rs
+++ b/library/alloc/src/vec/in_place_collect.rs
@@ -55,6 +55,9 @@
//! This is handled by the [`InPlaceDrop`] guard for sink items (`U`) and by
//! [`vec::IntoIter::forget_allocation_drop_remaining()`] for remaining source items (`T`).
//!
+//! If dropping any remaining source item (`T`) panics then [`InPlaceDstBufDrop`] will handle dropping
+//! the already collected sink items (`U`) and freeing the allocation.
+//!
//! [`vec::IntoIter::forget_allocation_drop_remaining()`]: super::IntoIter::forget_allocation_drop_remaining()
//!
//! # O(1) collect
@@ -138,7 +141,7 @@ use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccessNoCoerce};
use core::mem::{self, ManuallyDrop};
use core::ptr::{self};
-use super::{InPlaceDrop, SpecFromIter, SpecFromIterNested, Vec};
+use super::{InPlaceDrop, InPlaceDstBufDrop, SpecFromIter, SpecFromIterNested, Vec};
/// Specialization marker for collecting an iterator pipeline into a Vec while reusing the
/// source allocation, i.e. executing the pipeline in place.
@@ -191,14 +194,17 @@ where
);
}
- // Drop any remaining values at the tail of the source but prevent drop of the allocation
- // itself once IntoIter goes out of scope.
- // If the drop panics then we also leak any elements collected into dst_buf.
+ // The ownership of the allocation and the new `T` values is temporarily moved into `dst_guard`.
+ // This is safe because `forget_allocation_drop_remaining` immediately forgets the allocation
+ // before any panic can occur in order to avoid any double free, and then proceeds to drop
+ // any remaining values at the tail of the source.
//
// Note: This access to the source wouldn't be allowed by the TrustedRandomIteratorNoCoerce
// contract (used by SpecInPlaceCollect below). But see the "O(1) collect" section in the
// module documenttation why this is ok anyway.
+ let dst_guard = InPlaceDstBufDrop { ptr: dst_buf, len, cap };
src.forget_allocation_drop_remaining();
+ mem::forget(dst_guard);
let vec = unsafe { Vec::from_raw_parts(dst_buf, len, cap) };
diff --git a/library/alloc/src/vec/in_place_drop.rs b/library/alloc/src/vec/in_place_drop.rs
index 1b1ef9130face..25ca33c6a7bf0 100644
--- a/library/alloc/src/vec/in_place_drop.rs
+++ b/library/alloc/src/vec/in_place_drop.rs
@@ -22,3 +22,18 @@ impl<T> Drop for InPlaceDrop<T> {
}
}
}
+
+// A helper struct for in-place collection that drops the destination allocation and elements,
+// to avoid leaking them if some other destructor panics.
+pub(super) struct InPlaceDstBufDrop<T> {
+ pub(super) ptr: *mut T,
+ pub(super) len: usize,
+ pub(super) cap: usize,
+}
+
+impl<T> Drop for InPlaceDstBufDrop<T> {
+ #[inline]
+ fn drop(&mut self) {
+ unsafe { super::Vec::from_raw_parts(self.ptr, self.len, self.cap) };
+ }
+}
diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs
index 1b483e3fc7793..b5a03392a1797 100644
--- a/library/alloc/src/vec/into_iter.rs
+++ b/library/alloc/src/vec/into_iter.rs
@@ -96,13 +96,16 @@ impl<T, A: Allocator> IntoIter<T, A> {
}
/// Drops remaining elements and relinquishes the backing allocation.
+ /// This method guarantees it won't panic before relinquishing
+ /// the backing allocation.
///
/// This is roughly equivalent to the following, but more efficient
///
/// ```
/// # let mut into_iter = Vec::<u8>::with_capacity(10).into_iter();
+ /// let mut into_iter = std::mem::replace(&mut into_iter, Vec::new().into_iter());
/// (&mut into_iter).for_each(core::mem::drop);
- /// unsafe { core::ptr::write(&mut into_iter, Vec::new().into_iter()); }
+ /// std::mem::forget(into_iter);
/// ```
///
/// This method is used by in-place iteration, refer to the vec::in_place_collect
@@ -119,6 +122,8 @@ impl<T, A: Allocator> IntoIter<T, A> {
self.ptr = self.buf.as_ptr();
self.end = self.buf.as_ptr();
+ // Dropping the remaining elements can panic, so this needs to be
+ // done only after updating the other fields.
unsafe {
ptr::drop_in_place(remaining);
}
diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index fa9f2131c0c1d..acfbb98272dd5 100644
--- a/library/alloc/src/vec/mod.rs
+++ b/library/alloc/src/vec/mod.rs
@@ -125,7 +125,7 @@ use self::set_len_on_drop::SetLenOnDrop;
mod set_len_on_drop;
#[cfg(not(no_global_oom_handling))]
-use self::in_place_drop::InPlaceDrop;
+use self::in_place_drop::{InPlaceDrop, InPlaceDstBufDrop};
#[cfg(not(no_global_oom_handling))]
mod in_place_drop;
diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs
index d94da8f5f5a0e..990ce38e22edd 100644
--- a/library/alloc/tests/vec.rs
+++ b/library/alloc/tests/vec.rs
@@ -1110,48 +1110,53 @@ fn test_from_iter_specialization_panic_during_iteration_drops() {
}
#[test]
-fn test_from_iter_specialization_panic_during_drop_leaks() {
- static mut DROP_COUNTER: usize = 0;
+fn test_from_iter_specialization_panic_during_drop_doesnt_leak() {
+ static mut DROP_COUNTER_OLD: [usize; 5] = [0; 5];
+ static mut DROP_COUNTER_NEW: [usize; 2] = [0; 2];
#[derive(Debug)]
- enum Droppable {
- DroppedTwice(Box<i32>),
- PanicOnDrop,
- }
+ struct Old(usize);
- impl Drop for Droppable {
+ impl Drop for Old {
fn drop(&mut self) {
- match self {
- Droppable::DroppedTwice(_) => {
- unsafe {
- DROP_COUNTER += 1;
- }
- println!("Dropping!")
- }
- Droppable::PanicOnDrop => {
- if !std::thread::panicking() {
- panic!();
- }
- }
+ unsafe {
+ DROP_COUNTER_OLD[self.0] += 1;
+ }
+
+ if self.0 == 3 {
+ panic!();
}
+
+ println!("Dropped Old: {}", self.0);
}
}
- let mut to_free: *mut Droppable = core::ptr::null_mut();
- let mut cap = 0;
+ #[derive(Debug)]
+ struct New(usize);
+
+ impl Drop for New {
+ fn drop(&mut self) {
+ unsafe {
+ DROP_COUNTER_NEW[self.0] += 1;
+ }
+
+ println!("Dropped New: {}", self.0);
+ }
+ }
let _ = std::panic::catch_unwind(AssertUnwindSafe(|| {
- let mut v = vec![Droppable::DroppedTwice(Box::new(123)), Droppable::PanicOnDrop];
- to_free = v.as_mut_ptr();
- cap = v.capacity();
- let _ = v.into_iter().take(0).collect::<Vec<_>>();
+ let v = vec![Old(0), Old(1), Old(2), Old(3), Old(4)];
+ let _ = v.into_iter().map(|x| New(x.0)).take(2).collect::<Vec<_>>();
}));
- assert_eq!(unsafe { DROP_COUNTER }, 1);
- // clean up the leak to keep miri happy
- unsafe {
- drop(Vec::from_raw_parts(to_free, 0, cap));
- }
+ assert_eq!(unsafe { DROP_COUNTER_OLD[0] }, 1);
+ assert_eq!(unsafe { DROP_COUNTER_OLD[1] }, 1);
+ assert_eq!(unsafe { DROP_COUNTER_OLD[2] }, 1);
+ assert_eq!(unsafe { DROP_COUNTER_OLD[3] }, 1);
+ assert_eq!(unsafe { DROP_COUNTER_OLD[4] }, 1);
+
+ assert_eq!(unsafe { DROP_COUNTER_NEW[0] }, 1);
+ assert_eq!(unsafe { DROP_COUNTER_NEW[1] }, 1);
}
// regression test for issue #85322. Peekable previously implemented InPlaceIterable,