Stop re-implementing slice iterators in vec::IntoIter

library/alloc/src/lib.rs

@@ -146,6 +146,7 @@
 #![feature(receiver_trait)]
 #![feature(set_ptr_value)]
 #![feature(sized_type_properties)]
+#![feature(slice_drain_raw_iter)]
 #![feature(slice_from_ptr_range)]
 #![feature(slice_index_methods)]
 #![feature(slice_ptr_get)]

library/alloc/src/vec/drain.rs

@@ -1,7 +1,8 @@
 use crate::alloc::{Allocator, Global};
 use core::fmt;
 use core::iter::{FusedIterator, TrustedLen};
-use core::mem::{self, ManuallyDrop, SizedTypeProperties};
+use core::marker::PhantomData;
+use core::mem::{ManuallyDrop, SizedTypeProperties};
 use core::ptr::{self, NonNull};
 use core::slice::{self};
 
@@ -29,14 +30,15 @@ pub struct Drain<
     /// Length of tail
     pub(super) tail_len: usize,
     /// Current remaining range to remove
-    pub(super) iter: slice::Iter<'a, T>,
+    pub(super) iter: slice::DrainRaw<T>,
     pub(super) vec: NonNull<Vec<T, A>>,
+    pub(super) phantom: PhantomData<&'a [T]>,
 }
 
 #[stable(feature = "collection_debug", since = "1.17.0")]
 impl<T: fmt::Debug, A: Allocator> fmt::Debug for Drain<'_, T, A> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_tuple("Drain").field(&self.iter.as_slice()).finish()
+        f.debug_tuple("Drain").field(&self.as_slice()).finish()
     }
 }
 
@@ -55,7 +57,9 @@ impl<'a, T, A: Allocator> Drain<'a, T, A> {
     #[must_use]
     #[stable(feature = "vec_drain_as_slice", since = "1.46.0")]
     pub fn as_slice(&self) -> &[T] {
-        self.iter.as_slice()
+        // SAFETY: Restricting the lifetime of the returned slice to the self
+        // borrow keeps anything else from dropping them.
+        unsafe { self.iter.as_nonnull_slice().as_ref() }
     }
 
     /// Returns a reference to the underlying allocator.
@@ -108,8 +112,9 @@ impl<'a, T, A: Allocator> Drain<'a, T, A> {
             let start = source_vec.len();
             let tail = this.tail_start;
 
-            let unyielded_len = this.iter.len();
-            let unyielded_ptr = this.iter.as_slice().as_ptr();
+            let keep_items = this.iter.forget_remaining();
+            let unyielded_len = keep_items.len();
+            let unyielded_ptr = keep_items.as_mut_ptr();
 
             // ZSTs have no identity, so we don't need to move them around.
             if !T::IS_ZST {
@@ -154,7 +159,7 @@ impl<T, A: Allocator> Iterator for Drain<'_, T, A> {
 
     #[inline]
     fn next(&mut self) -> Option<T> {
-        self.iter.next().map(|elt| unsafe { ptr::read(elt as *const _) })
+        self.iter.next()
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
@@ -166,7 +171,7 @@ impl<T, A: Allocator> Iterator for Drain<'_, T, A> {
 impl<T, A: Allocator> DoubleEndedIterator for Drain<'_, T, A> {
     #[inline]
     fn next_back(&mut self) -> Option<T> {
-        self.iter.next_back().map(|elt| unsafe { ptr::read(elt as *const _) })
+        self.iter.next_back()
     }
 }
 
@@ -195,16 +200,13 @@ impl<T, A: Allocator> Drop for Drain<'_, T, A> {
             }
         }
 
-        let iter = mem::take(&mut self.iter);
-        let drop_len = iter.len();
-
-        let mut vec = self.vec;
-
         if T::IS_ZST {
+            let drop_len = self.iter.forget_remaining().len();
+
             // ZSTs have no identity, so we don't need to move them around, we only need to drop the correct amount.
             // this can be achieved by manipulating the Vec length instead of moving values out from `iter`.
             unsafe {
-                let vec = vec.as_mut();
+                let vec = self.vec.as_mut();
                 let old_len = vec.len();
                 vec.set_len(old_len + drop_len + self.tail_len);
                 vec.truncate(old_len + self.tail_len);
@@ -214,28 +216,9 @@ impl<T, A: Allocator> Drop for Drain<'_, T, A> {
         }
 
         // ensure elements are moved back into their appropriate places, even when drop_in_place panics
-        let _guard = DropGuard(self);
+        let guard = DropGuard(self);
 
-        if drop_len == 0 {
-            return;
-        }
-
-        // as_slice() must only be called when iter.len() is > 0 because
-        // it also gets touched by vec::Splice which may turn it into a dangling pointer
-        // which would make it and the vec pointer point to different allocations which would
-        // lead to invalid pointer arithmetic below.
-        let drop_ptr = iter.as_slice().as_ptr();
-
-        unsafe {
-            // drop_ptr comes from a slice::Iter which only gives us a &[T] but for drop_in_place
-            // a pointer with mutable provenance is necessary. Therefore we must reconstruct
-            // it from the original vec but also avoid creating a &mut to the front since that could
-            // invalidate raw pointers to it which some unsafe code might rely on.
-            let vec_ptr = vec.as_mut().as_mut_ptr();
-            let drop_offset = drop_ptr.sub_ptr(vec_ptr);
-            let to_drop = ptr::slice_from_raw_parts_mut(vec_ptr.add(drop_offset), drop_len);
-            ptr::drop_in_place(to_drop);
-        }
+        guard.0.iter.drop_remaining();
     }
 }
 

library/alloc/src/vec/in_place_collect.rs

@@ -253,12 +253,14 @@ where
 {
     let (src_buf, src_ptr, src_cap, mut dst_buf, dst_end, dst_cap) = unsafe {
         let inner = iterator.as_inner().as_into_iter();
+        let inner_ptr = inner.ptr();
+        let inner_end = inner_ptr.add(inner.len());
         (
             inner.buf,
-            inner.ptr,
+            inner_ptr,
             inner.cap,
             inner.buf.cast::<T>(),
-            inner.end as *const T,
+            inner_end.as_ptr() as *const T,
             inner.cap * mem::size_of::<I::Src>() / mem::size_of::<T>(),
         )
     };
@@ -275,9 +277,9 @@ where
     // check InPlaceIterable contract. This is only possible if the iterator advanced the
     // source pointer at all. If it uses unchecked access via TrustedRandomAccess
     // then the source pointer will stay in its initial position and we can't use it as reference
-    if src.ptr != src_ptr {
+    if src.ptr() != src_ptr {
         debug_assert!(
-            unsafe { dst_buf.add(len).cast() } <= src.ptr,
+            unsafe { dst_buf.add(len).cast() } <= src.ptr(),
             "InPlaceIterable contract violation, write pointer advanced beyond read pointer"
         );
     }