fix double-drop in in-place collect specialization

2 files changed, 20 insertions, 11 deletions

diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs
index bcbdffabc7f..324e894bafd 100644
--- a/library/alloc/src/vec/into_iter.rs
+++ b/library/alloc/src/vec/into_iter.rs
@@ -85,20 +85,29 @@ impl<T, A: Allocator> IntoIter<T, A> {
         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::<u8>::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 50882fc1767..e857d284d3a 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) };