Avoid more NonNull-raw-NonNull roundtrips in Vec

Author: saethlin

library/alloc/src/raw_vec.rs

@@ -259,6 +259,11 @@ impl<T, A: Allocator> RawVec<T, A> {
         Self { ptr: unsafe { Unique::new_unchecked(ptr) }, cap, alloc }
     }
 
+    pub(crate) unsafe fn from_nonnull_in(ptr: NonNull<T>, capacity: usize, alloc: A) -> Self {
+        let cap = if T::IS_ZST { Cap::ZERO } else { unsafe { Cap(capacity) } };
+        Self { ptr: Unique::from(ptr), cap, alloc }
+    }
+
     /// Gets a raw pointer to the start of the allocation. Note that this is
     /// `Unique::dangling()` if `capacity == 0` or `T` is zero-sized. In the former case, you must
     /// be careful.

library/alloc/src/vec/in_place_collect.rs

@@ -161,7 +161,7 @@ use core::iter::{InPlaceIterable, SourceIter, TrustedRandomAccessNoCoerce};
 use core::marker::PhantomData;
 use core::mem::{self, ManuallyDrop, SizedTypeProperties};
 use core::num::NonZero;
-use core::ptr::{self, NonNull};
+use core::ptr;
 
 use super::{InPlaceDrop, InPlaceDstDataSrcBufDrop, SpecFromIter, SpecFromIterNested, Vec};
 
@@ -254,28 +254,30 @@ where
     let (src_buf, src_ptr, src_cap, mut dst_buf, dst_end, dst_cap) = unsafe {
         let inner = iterator.as_inner().as_into_iter();
         (
-            inner.buf.as_ptr(),
+            inner.buf,
             inner.ptr,
             inner.cap,
-            inner.buf.as_ptr() as *mut T,
+            inner.buf.cast::<T>(),
             inner.end as *const T,
             inner.cap * mem::size_of::<I::Src>() / mem::size_of::<T>(),
         )
     };
 
     // SAFETY: `dst_buf` and `dst_end` are the start and end of the buffer.
-    let len = unsafe { SpecInPlaceCollect::collect_in_place(&mut iterator, dst_buf, dst_end) };
+    let len = unsafe {
+        SpecInPlaceCollect::collect_in_place(&mut iterator, dst_buf.as_ptr() as *mut T, dst_end)
+    };
 
     let src = unsafe { iterator.as_inner().as_into_iter() };
     // check if SourceIter contract was upheld
     // caveat: if they weren't we might not even make it to this point
-    debug_assert_eq!(src_buf, src.buf.as_ptr());
+    debug_assert_eq!(src_buf, src.buf);
     // 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 {
         debug_assert!(
-            unsafe { dst_buf.add(len) as *const _ } <= src.ptr.as_ptr(),
+            unsafe { dst_buf.add(len).cast() } <= src.ptr,
             "InPlaceIterable contract violation, write pointer advanced beyond read pointer"
         );
     }
@@ -315,18 +317,17 @@ where
             let dst_size = mem::size_of::<T>().unchecked_mul(dst_cap);
             let new_layout = Layout::from_size_align_unchecked(dst_size, dst_align);
 
-            let result =
-                alloc.shrink(NonNull::new_unchecked(dst_buf as *mut u8), old_layout, new_layout);
+            let result = alloc.shrink(dst_buf.cast(), old_layout, new_layout);
             let Ok(reallocated) = result else { handle_alloc_error(new_layout) };
-            dst_buf = reallocated.as_ptr() as *mut T;
+            dst_buf = reallocated.cast();
         }
     } else {
         debug_assert_eq!(src_cap * mem::size_of::<I::Src>(), dst_cap * mem::size_of::<T>());
     }
 
     mem::forget(dst_guard);
 
-    let vec = unsafe { Vec::from_raw_parts(dst_buf, len, dst_cap) };
+    let vec = unsafe { Vec::from_nonnull(dst_buf, len, dst_cap) };
 
     vec
 }

library/alloc/src/vec/in_place_drop.rs

@@ -1,4 +1,5 @@
 use core::marker::PhantomData;
+use core::ptr::NonNull;
 use core::ptr::{self, drop_in_place};
 use core::slice::{self};
 
@@ -31,7 +32,7 @@ impl<T> Drop for InPlaceDrop<T> {
 // the source allocation - i.e. before the reallocation happened - to avoid leaking them
 // if some other destructor panics.
 pub(super) struct InPlaceDstDataSrcBufDrop<Src, Dest> {
-    pub(super) ptr: *mut Dest,
+    pub(super) ptr: NonNull<Dest>,
     pub(super) len: usize,
     pub(super) src_cap: usize,
     pub(super) src: PhantomData<Src>,
@@ -42,8 +43,8 @@ impl<Src, Dest> Drop for InPlaceDstDataSrcBufDrop<Src, Dest> {
     fn drop(&mut self) {
         unsafe {
             let _drop_allocation =
-                RawVec::<Src>::from_raw_parts_in(self.ptr.cast::<Src>(), self.src_cap, Global);
-            drop_in_place(core::ptr::slice_from_raw_parts_mut::<Dest>(self.ptr, self.len));
+                RawVec::<Src>::from_nonnull_in(self.ptr.cast::<Src>(), self.src_cap, Global);
+            drop_in_place(core::ptr::slice_from_raw_parts_mut::<Dest>(self.ptr.as_ptr(), self.len));
         };
     }
 }

library/alloc/src/vec/into_iter.rs

@@ -433,7 +433,7 @@ unsafe impl<#[may_dangle] T, A: Allocator> Drop for IntoIter<T, A> {
                     // `IntoIter::alloc` is not used anymore after this and will be dropped by RawVec
                     let alloc = ManuallyDrop::take(&mut self.0.alloc);
                     // RawVec handles deallocation
-                    let _ = RawVec::from_raw_parts_in(self.0.buf.as_ptr(), self.0.cap, alloc);
+                    let _ = RawVec::from_nonnull_in(self.0.buf, self.0.cap, alloc);
                 }
             }
         }

library/alloc/src/vec/mod.rs

@@ -603,6 +603,11 @@ impl<T> Vec<T> {
     pub unsafe fn from_raw_parts(ptr: *mut T, length: usize, capacity: usize) -> Self {
         unsafe { Self::from_raw_parts_in(ptr, length, capacity, Global) }
     }
+
+    #[inline]
+    pub(crate) unsafe fn from_nonnull(ptr: NonNull<T>, length: usize, capacity: usize) -> Self {
+        unsafe { Self::from_nonnull_in(ptr, length, capacity, Global) }
+    }
 }
 
 impl<T, A: Allocator> Vec<T, A> {
@@ -820,6 +825,16 @@ impl<T, A: Allocator> Vec<T, A> {
         unsafe { Vec { buf: RawVec::from_raw_parts_in(ptr, capacity, alloc), len: length } }
     }
 
+    #[inline]
+    pub(crate) unsafe fn from_nonnull_in(
+        ptr: NonNull<T>,
+        length: usize,
+        capacity: usize,
+        alloc: A,
+    ) -> Self {
+        unsafe { Vec { buf: RawVec::from_nonnull_in(ptr, capacity, alloc), len: length } }
+    }
+
     /// Decomposes a `Vec<T>` into its raw components: `(pointer, length, capacity)`.
     ///
     /// Returns the raw pointer to the underlying data, the length of

library/alloc/src/vec/spec_from_iter.rs

@@ -51,7 +51,7 @@ impl<T> SpecFromIter<T, IntoIter<T>> for Vec<T> {
                 if has_advanced {
                     ptr::copy(it.ptr.as_ptr(), it.buf.as_ptr(), it.len());
                 }
-                return Vec::from_raw_parts(it.buf.as_ptr(), it.len(), it.cap);
+                return Vec::from_nonnull(it.buf, it.len(), it.cap);
             }
         }