diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs index bcbdffabc7fbe..324e894bafd23 100644 --- a/library/alloc/src/vec/into_iter.rs +++ b/library/alloc/src/vec/into_iter.rs @@ -85,20 +85,29 @@ impl IntoIter { ptr::slice_from_raw_parts_mut(self.ptr as *mut T, self.len()) } - pub(super) fn drop_remaining(&mut self) { - unsafe { - ptr::drop_in_place(self.as_mut_slice()); - } - self.ptr = self.end; - } + /// Drops remaining elements and relinquishes the backing allocation. + /// + /// This is roughly equivalent to the following, but more efficient + /// + /// ``` + /// # let mut into_iter = Vec::::with_capacity(10).into_iter(); + /// (&mut into_iter).for_each(core::mem::drop); + /// unsafe { core::ptr::write(&mut into_iter, Vec::new().into_iter()); } + /// ``` + pub(super) fn forget_allocation_drop_remaining(&mut self) { + let remaining = self.as_raw_mut_slice(); - /// Relinquishes the backing allocation, equivalent to - /// `ptr::write(&mut self, Vec::new().into_iter())` - pub(super) fn forget_allocation(&mut self) { + // overwrite the individual fields instead of creating a new + // struct and then overwriting &mut self. + // this creates less assembly self.cap = 0; self.buf = unsafe { NonNull::new_unchecked(RawVec::NEW.ptr()) }; self.ptr = self.buf.as_ptr(); self.end = self.buf.as_ptr(); + + unsafe { + ptr::drop_in_place(remaining); + } } } diff --git a/library/alloc/src/vec/source_iter_marker.rs b/library/alloc/src/vec/source_iter_marker.rs index 50882fc17673e..e857d284d3ab6 100644 --- a/library/alloc/src/vec/source_iter_marker.rs +++ b/library/alloc/src/vec/source_iter_marker.rs @@ -69,9 +69,9 @@ where } // drop any remaining values at the tail of the source - src.drop_remaining(); // but prevent drop of the allocation itself once IntoIter goes out of scope - src.forget_allocation(); + // if the drop panics then we also leak any elements collected into dst_buf + src.forget_allocation_drop_remaining(); let vec = unsafe { Vec::from_raw_parts(dst_buf, len, cap) }; diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index c142536cd2dfb..b9fe07c73e55e 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -1,3 +1,4 @@ +use alloc::boxed::Box; use std::borrow::Cow; use std::cell::Cell; use std::collections::TryReserveError::*; @@ -1027,7 +1028,7 @@ fn test_from_iter_specialization_head_tail_drop() { } #[test] -fn test_from_iter_specialization_panic_drop() { +fn test_from_iter_specialization_panic_during_iteration_drops() { let drop_count: Vec<_> = (0..=2).map(|_| Rc::new(())).collect(); let src: Vec<_> = drop_count.iter().cloned().collect(); let iter = src.into_iter(); @@ -1050,6 +1051,51 @@ fn test_from_iter_specialization_panic_drop() { ); } +#[test] +fn test_from_iter_specialization_panic_during_drop_leaks() { + static mut DROP_COUNTER: usize = 0; + + #[derive(Debug)] + enum Droppable { + DroppedTwice, + PanicOnDrop, + } + + impl Drop for Droppable { + fn drop(&mut self) { + match self { + Droppable::DroppedTwice => { + unsafe { + DROP_COUNTER += 1; + } + println!("Dropping!") + } + Droppable::PanicOnDrop => { + if !std::thread::panicking() { + panic!(); + } + } + } + } + } + + let mut to_free: *mut Droppable = core::ptr::null_mut(); + let mut cap = 0; + + let _ = std::panic::catch_unwind(AssertUnwindSafe(|| { + let mut v = vec![Droppable::DroppedTwice, Droppable::PanicOnDrop]; + to_free = v.as_mut_ptr(); + cap = v.capacity(); + let _ = v.into_iter().take(0).collect::>(); + })); + + assert_eq!(unsafe { DROP_COUNTER }, 1); + // clean up the leak to keep miri happy + unsafe { + Vec::from_raw_parts(to_free, 0, cap); + } +} + #[test] fn test_cow_from() { let borrowed: &[_] = &["borrowed", "(slice)"];