diff --git a/library/alloc/src/vec/drain.rs b/library/alloc/src/vec/drain.rs index f0b63759ac70f..8630eef690c1d 100644 --- a/library/alloc/src/vec/drain.rs +++ b/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, pub(super) vec: NonNull>, + pub(super) phantom: PhantomData<&'a [T]>, } #[stable(feature = "collection_debug", since = "1.17.0")] impl 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 Iterator for Drain<'_, T, A> { #[inline] fn next(&mut self) -> Option { - self.iter.next().map(|elt| unsafe { ptr::read(elt as *const _) }) + self.iter.next() } fn size_hint(&self) -> (usize, Option) { @@ -166,7 +171,7 @@ impl Iterator for Drain<'_, T, A> { impl DoubleEndedIterator for Drain<'_, T, A> { #[inline] fn next_back(&mut self) -> Option { - self.iter.next_back().map(|elt| unsafe { ptr::read(elt as *const _) }) + self.iter.next_back() } } @@ -195,16 +200,13 @@ impl 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 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(); } } diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index c0b73ed5141bb..74ef3b9e00598 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -1389,7 +1389,11 @@ impl Vec { 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 { + self.buf.non_null() } /// Returns a reference to the underlying allocator. @@ -2199,12 +2203,13 @@ impl Vec { 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, } } } diff --git a/library/alloc/src/vec/splice.rs b/library/alloc/src/vec/splice.rs index 852fdcc3f5ce7..c1686cf1807aa 100644 --- a/library/alloc/src/vec/splice.rs +++ b/library/alloc/src/vec/splice.rs @@ -59,7 +59,7 @@ impl 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 { diff --git a/library/core/src/slice/drain.rs b/library/core/src/slice/drain.rs index d8e40f3003da0..6ed78d9a50f38 100644 --- a/library/core/src/slice/drain.rs +++ b/library/core/src/slice/drain.rs @@ -48,6 +48,21 @@ impl fmt::Debug for DrainRaw { } } +#[unstable(feature = "slice_drain_raw_iter", issue = "none")] +impl Default for DrainRaw { + /// 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 DrainRaw { /// Creates a new iterator which moves the `len` items starting at `ptr` /// while it's iterated, or drops them when the iterator is dropped.