Shrink from_raw_parts's MIR so that Vec::deref MIR-inlines again

scottmcm · scottmcm · commit fcea56fbe083 · 2024-03-29T02:04:19.000-07:00
Fixes 123174
cc CAD97
diff --git a/library/core/src/ptr/metadata.rs b/library/core/src/ptr/metadata.rs
@@ -92,10 +92,14 @@ pub trait Thin = Pointee<Metadata = ()>;
 #[rustc_const_unstable(feature = "ptr_metadata", issue = "81513")]
 #[inline]
 pub const fn metadata<T: ?Sized>(ptr: *const T) -> <T as Pointee>::Metadata {
-    // SAFETY: Accessing the value from the `PtrRepr` union is safe since *const T
-    // and PtrComponents<T> have the same memory layouts. Only std can make this
-    // guarantee.
-    unsafe { PtrRepr { const_ptr: ptr }.components.metadata }
+    metadata_flexible(ptr)
+}
+
+#[inline]
+const fn metadata_flexible<P: RawPointer>(ptr: P) -> <P::Pointee as Pointee>::Metadata {
+    // SAFETY: Transmuting like this is safe since `P` and `PtrComponents<P, _>`
+    // have the same memory layouts. Only std can make this guarantee.
+    unsafe { crate::intrinsics::transmute_unchecked::<P, PtrComponents<P::Pointee>>(ptr).metadata }
 }
 
 /// Forms a (possibly-wide) raw pointer from a data pointer and metadata.
@@ -112,10 +116,7 @@ pub const fn from_raw_parts<T: ?Sized>(
     data_pointer: *const (),
     metadata: <T as Pointee>::Metadata,
 ) -> *const T {
-    // SAFETY: Accessing the value from the `PtrRepr` union is safe since *const T
-    // and PtrComponents<T> have the same memory layouts. Only std can make this
-    // guarantee.
-    unsafe { PtrRepr { components: PtrComponents { data_pointer, metadata } }.const_ptr }
+    from_raw_parts_flexible(data_pointer, metadata)
 }
 
 /// Performs the same functionality as [`from_raw_parts`], except that a
@@ -129,30 +130,64 @@ pub const fn from_raw_parts_mut<T: ?Sized>(
     data_pointer: *mut (),
     metadata: <T as Pointee>::Metadata,
 ) -> *mut T {
-    // SAFETY: Accessing the value from the `PtrRepr` union is safe since *const T
-    // and PtrComponents<T> have the same memory layouts. Only std can make this
-    // guarantee.
-    unsafe { PtrRepr { components: PtrComponents { data_pointer, metadata } }.mut_ptr }
+    from_raw_parts_flexible(data_pointer, metadata)
 }
 
-#[repr(C)]
-union PtrRepr<T: ?Sized> {
-    const_ptr: *const T,
-    mut_ptr: *mut T,
-    components: PtrComponents<T>,
+/// Just like [`from_raw_parts`] and [`from_raw_parts_mut`], but more flexible
+/// in terms of which types it can take, allowing smaller MIR.
+// See <https://github.com/rust-lang/rust/issues/123174>
+#[rustc_const_unstable(feature = "ptr_metadata", issue = "81513")]
+#[inline]
+pub(crate) const fn from_raw_parts_flexible<P: RawPointer>(
+    data_pointer: impl RawPointer<Pointee: Thin>,
+    metadata: <P::Pointee as Pointee>::Metadata,
+) -> P {
+    // SAFETY: Transmuting like this is safe since `P` and `PtrComponents<P, _>`
+    // have the same memory layouts. Only std can make this guarantee.
+    unsafe {
+        crate::intrinsics::transmute_unchecked::<PtrComponents<P::Pointee, _>, P>(PtrComponents {
+            data_pointer,
+            metadata,
+        })
+    }
+}
+
+use private_bounds::RawPointer;
+mod private_bounds {
+    /// Internal trait to avoid bad instantiations of [`PtrComponents`]
+    ///
+    /// # Safety
+    ///
+    /// Must have the same layout as `*const Self::Pointee` and be able to hold provenance.
+    pub unsafe trait RawPointer {
+        type Pointee: ?Sized + super::Pointee;
+    }
+}
+
+// SAFETY: `*const T` is obviously a raw pointer
+unsafe impl<T: ?Sized> RawPointer for *const T {
+    type Pointee = T;
+}
+// SAFETY: `*mut T` is obviously a raw pointer
+unsafe impl<T: ?Sized> RawPointer for *mut T {
+    type Pointee = T;
+}
+// SAFETY: `NonNull<T>` is a transparent newtype around a `*const T`.
+unsafe impl<T: ?Sized> RawPointer for super::NonNull<T> {
+    type Pointee = T;
 }
 
 #[repr(C)]
-struct PtrComponents<T: ?Sized> {
-    data_pointer: *const (),
+struct PtrComponents<T: ?Sized, P = *const ()> {
+    data_pointer: P,
     metadata: <T as Pointee>::Metadata,
 }
 
 // Manual impl needed to avoid `T: Copy` bound.
-impl<T: ?Sized> Copy for PtrComponents<T> {}
+impl<T: ?Sized, P: Copy> Copy for PtrComponents<T, P> {}
 
 // Manual impl needed to avoid `T: Clone` bound.
-impl<T: ?Sized> Clone for PtrComponents<T> {
+impl<T: ?Sized, P: Copy> Clone for PtrComponents<T, P> {
     fn clone(&self) -> Self {
         *self
     }
diff --git a/library/core/src/ptr/mod.rs b/library/core/src/ptr/mod.rs
@@ -411,6 +411,7 @@ pub use crate::intrinsics::copy;
 pub use crate::intrinsics::write_bytes;
 
 mod metadata;
+pub(crate) use metadata::from_raw_parts_flexible;
 #[unstable(feature = "ptr_metadata", issue = "81513")]
 pub use metadata::{from_raw_parts, from_raw_parts_mut, metadata, DynMetadata, Pointee, Thin};
 
@@ -812,7 +813,7 @@ pub const fn from_mut<T: ?Sized>(r: &mut T) -> *mut T {
 #[rustc_allow_const_fn_unstable(ptr_metadata)]
 #[rustc_diagnostic_item = "ptr_slice_from_raw_parts"]
 pub const fn slice_from_raw_parts<T>(data: *const T, len: usize) -> *const [T] {
-    from_raw_parts(data.cast(), len)
+    from_raw_parts_flexible(data, len)
 }
 
 /// Forms a raw mutable slice from a pointer and a length.
@@ -858,7 +859,7 @@ pub const fn slice_from_raw_parts<T>(data: *const T, len: usize) -> *const [T] {
 #[rustc_const_unstable(feature = "const_slice_from_raw_parts_mut", issue = "67456")]
 #[rustc_diagnostic_item = "ptr_slice_from_raw_parts_mut"]
 pub const fn slice_from_raw_parts_mut<T>(data: *mut T, len: usize) -> *mut [T] {
-    from_raw_parts_mut(data.cast(), len)
+    from_raw_parts_flexible(data, len)
 }
 
 /// Swaps the values at two mutable locations of the same type, without
diff --git a/library/core/src/ptr/non_null.rs b/library/core/src/ptr/non_null.rs
@@ -265,10 +265,7 @@ impl<T: ?Sized> NonNull<T> {
         data_pointer: NonNull<()>,
         metadata: <T as super::Pointee>::Metadata,
     ) -> NonNull<T> {
-        // SAFETY: The result of `ptr::from::raw_parts_mut` is non-null because `data_pointer` is.
-        unsafe {
-            NonNull::new_unchecked(super::from_raw_parts_mut(data_pointer.as_ptr(), metadata))
-        }
+        super::from_raw_parts_flexible(data_pointer, metadata)
     }
 
     /// Decompose a (possibly wide) pointer into its data pointer and metadata components.
diff --git a/tests/mir-opt/pre-codegen/vec_deref.vec_deref_to_slice.PreCodegen.after.panic-abort.mir b/tests/mir-opt/pre-codegen/vec_deref.vec_deref_to_slice.PreCodegen.after.panic-abort.mir
@@ -3,12 +3,94 @@
 fn vec_deref_to_slice(_1: &Vec<u8>) -> &[u8] {
     debug v => _1;
     let mut _0: &[u8];
+    scope 1 (inlined <Vec<u8> as Deref>::deref) {
+        debug self => _1;
+        let mut _4: *const u8;
+        let mut _5: usize;
+        scope 2 {
+            scope 3 (inlined Vec::<u8>::as_ptr) {
+                debug self => _1;
+                let mut _2: &alloc::raw_vec::RawVec<u8>;
+                scope 4 (inlined alloc::raw_vec::RawVec::<u8>::ptr) {
+                    debug self => _2;
+                    let mut _3: std::ptr::NonNull<u8>;
+                    scope 5 (inlined Unique::<u8>::as_ptr) {
+                        debug ((self: Unique<u8>).0: std::ptr::NonNull<u8>) => _3;
+                        debug ((self: Unique<u8>).1: std::marker::PhantomData<u8>) => const PhantomData::<u8>;
+                        scope 6 (inlined NonNull::<u8>::as_ptr) {
+                            debug self => _3;
+                        }
+                    }
+                }
+            }
+            scope 7 (inlined std::slice::from_raw_parts::<'_, u8>) {
+                debug data => _4;
+                debug len => _5;
+                let mut _6: bool;
+                let mut _7: *mut ();
+                let _8: ();
+                let _10: *const [u8];
+                scope 8 {
+                    scope 9 (inlined core::ub_checks::check_language_ub) {
+                        scope 10 (inlined core::ub_checks::check_language_ub::runtime) {
+                        }
+                    }
+                    scope 11 (inlined std::mem::size_of::<u8>) {
+                    }
+                    scope 12 (inlined align_of::<u8>) {
+                    }
+                    scope 13 (inlined slice_from_raw_parts::<u8>) {
+                        debug data => _4;
+                        debug len => _5;
+                        scope 14 (inlined std::ptr::metadata::from_raw_parts_flexible::<*const [u8], *const u8>) {
+                            debug data_pointer => _4;
+                            debug metadata => _5;
+                            let mut _9: std::ptr::metadata::PtrComponents<[u8], *const u8>;
+                            scope 15 {
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
 
     bb0: {
-        _0 = <Vec<u8> as Deref>::deref(move _1) -> [return: bb1, unwind unreachable];
+        StorageLive(_4);
+        StorageLive(_2);
+        _2 = &((*_1).0: alloc::raw_vec::RawVec<u8>);
+        StorageLive(_3);
+        _3 = ((((*_1).0: alloc::raw_vec::RawVec<u8>).0: std::ptr::Unique<u8>).0: std::ptr::NonNull<u8>);
+        _4 = (_3.0: *const u8);
+        StorageDead(_3);
+        StorageDead(_2);
+        StorageLive(_5);
+        _5 = ((*_1).1: usize);
+        StorageLive(_6);
+        _6 = UbChecks();
+        switchInt(move _6) -> [0: bb3, otherwise: bb1];
     }
 
     bb1: {
+        StorageLive(_7);
+        _7 = _4 as *mut () (PtrToPtr);
+        _8 = std::slice::from_raw_parts::precondition_check(move _7, const 1_usize, const 1_usize, _5) -> [return: bb2, unwind unreachable];
+    }
+
+    bb2: {
+        StorageDead(_7);
+        goto -> bb3;
+    }
+
+    bb3: {
+        StorageDead(_6);
+        StorageLive(_9);
+        _9 = std::ptr::metadata::PtrComponents::<[u8], *const u8> { data_pointer: _4, metadata: _5 };
+        _10 = move _9 as *const [u8] (Transmute);
+        StorageDead(_9);
+        StorageDead(_5);
+        StorageDead(_4);
+        _0 = &(*_10);
         return;
     }
 }
diff --git a/tests/mir-opt/pre-codegen/vec_deref.vec_deref_to_slice.PreCodegen.after.panic-unwind.mir b/tests/mir-opt/pre-codegen/vec_deref.vec_deref_to_slice.PreCodegen.after.panic-unwind.mir
@@ -3,12 +3,94 @@
 fn vec_deref_to_slice(_1: &Vec<u8>) -> &[u8] {
     debug v => _1;
     let mut _0: &[u8];
+    scope 1 (inlined <Vec<u8> as Deref>::deref) {
+        debug self => _1;
+        let mut _4: *const u8;
+        let mut _5: usize;
+        scope 2 {
+            scope 3 (inlined Vec::<u8>::as_ptr) {
+                debug self => _1;
+                let mut _2: &alloc::raw_vec::RawVec<u8>;
+                scope 4 (inlined alloc::raw_vec::RawVec::<u8>::ptr) {
+                    debug self => _2;
+                    let mut _3: std::ptr::NonNull<u8>;
+                    scope 5 (inlined Unique::<u8>::as_ptr) {
+                        debug ((self: Unique<u8>).0: std::ptr::NonNull<u8>) => _3;
+                        debug ((self: Unique<u8>).1: std::marker::PhantomData<u8>) => const PhantomData::<u8>;
+                        scope 6 (inlined NonNull::<u8>::as_ptr) {
+                            debug self => _3;
+                        }
+                    }
+                }
+            }
+            scope 7 (inlined std::slice::from_raw_parts::<'_, u8>) {
+                debug data => _4;
+                debug len => _5;
+                let mut _6: bool;
+                let mut _7: *mut ();
+                let _8: ();
+                let _10: *const [u8];
+                scope 8 {
+                    scope 9 (inlined core::ub_checks::check_language_ub) {
+                        scope 10 (inlined core::ub_checks::check_language_ub::runtime) {
+                        }
+                    }
+                    scope 11 (inlined std::mem::size_of::<u8>) {
+                    }
+                    scope 12 (inlined align_of::<u8>) {
+                    }
+                    scope 13 (inlined slice_from_raw_parts::<u8>) {
+                        debug data => _4;
+                        debug len => _5;
+                        scope 14 (inlined std::ptr::metadata::from_raw_parts_flexible::<*const [u8], *const u8>) {
+                            debug data_pointer => _4;
+                            debug metadata => _5;
+                            let mut _9: std::ptr::metadata::PtrComponents<[u8], *const u8>;
+                            scope 15 {
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
 
     bb0: {
-        _0 = <Vec<u8> as Deref>::deref(move _1) -> [return: bb1, unwind continue];
+        StorageLive(_4);
+        StorageLive(_2);
+        _2 = &((*_1).0: alloc::raw_vec::RawVec<u8>);
+        StorageLive(_3);
+        _3 = ((((*_1).0: alloc::raw_vec::RawVec<u8>).0: std::ptr::Unique<u8>).0: std::ptr::NonNull<u8>);
+        _4 = (_3.0: *const u8);
+        StorageDead(_3);
+        StorageDead(_2);
+        StorageLive(_5);
+        _5 = ((*_1).1: usize);
+        StorageLive(_6);
+        _6 = UbChecks();
+        switchInt(move _6) -> [0: bb3, otherwise: bb1];
     }
 
     bb1: {
+        StorageLive(_7);
+        _7 = _4 as *mut () (PtrToPtr);
+        _8 = std::slice::from_raw_parts::precondition_check(move _7, const 1_usize, const 1_usize, _5) -> [return: bb2, unwind unreachable];
+    }
+
+    bb2: {
+        StorageDead(_7);
+        goto -> bb3;
+    }
+
+    bb3: {
+        StorageDead(_6);
+        StorageLive(_9);
+        _9 = std::ptr::metadata::PtrComponents::<[u8], *const u8> { data_pointer: _4, metadata: _5 };
+        _10 = move _9 as *const [u8] (Transmute);
+        StorageDead(_9);
+        StorageDead(_5);
+        StorageDead(_4);
+        _0 = &(*_10);
         return;
     }
 }
