Use DrainRaw in vec::Drain too.

Author: scottmcm

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/mod.rs

@@ -1389,7 +1389,11 @@ impl<T, A: Allocator> Vec<T, A> {
     pub fn as_mut_ptr(&mut self) -> *mut T {
         // We shadow the slice method of the same name to avoid going through
         // `deref_mut`, which creates an intermediate reference.
-        self.buf.ptr()
+        self.as_nonnull_ptr().as_ptr()
+    }
+
+    fn as_nonnull_ptr(&mut self) -> NonNull<T> {
+        self.buf.non_null()
     }
 
     /// Returns a reference to the underlying allocator.
@@ -2199,12 +2203,13 @@ impl<T, A: Allocator> Vec<T, A> {
         unsafe {
             // set self.vec length's to start, to be safe in case Drain is leaked
             self.set_len(start);
-            let range_slice = slice::from_raw_parts(self.as_ptr().add(start), end - start);
+            let drain = DrainRaw::from_parts(self.as_nonnull_ptr().add(start), end - start);
             Drain {
                 tail_start: end,
                 tail_len: len - end,
-                iter: range_slice.iter(),
+                iter: drain,
                 vec: NonNull::from(self),
+                phantom: PhantomData,
             }
         }
     }

library/alloc/src/vec/splice.rs

@@ -59,7 +59,7 @@ impl<I: Iterator, A: Allocator> Drop for Splice<'_, I, A> {
         // Which means we can replace the slice::Iter with pointers that won't point to deallocated
         // memory, so that Drain::drop is still allowed to call iter.len(), otherwise it would break
         // the ptr.sub_ptr contract.
-        self.drain.iter = (&[]).iter();
+        self.drain.iter = Default::default();
 
         unsafe {
             if self.drain.tail_len == 0 {

library/core/src/slice/drain.rs

@@ -48,6 +48,21 @@ impl<T: fmt::Debug> fmt::Debug for DrainRaw<T> {
     }
 }
 
+#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
+impl<T> Default for DrainRaw<T> {
+    /// Creates an empty slice iterator.
+    ///
+    /// ```
+    /// # use core::slice::IterMut;
+    /// let iter: IterMut<'_, u8> = Default::default();
+    /// assert_eq!(iter.len(), 0);
+    /// ```
+    fn default() -> Self {
+        // SAFETY: dangling is sufficiently-aligned so zero-length is always fine
+        unsafe { DrainRaw::from_parts(NonNull::dangling(), 0) }
+    }
+}
+
 impl<T> DrainRaw<T> {
     /// Creates a new iterator which moves the `len` items starting at `ptr`
     /// while it's iterated, or drops them when the iterator is dropped.