Auto merge of rust-lang#126556 - saethlin:layout-precondition, r=<try>

Add a precondition check for Layout::from_size_align_unchecked

Ran into this while looking into rust-lang/miri#3679. This is of course not the cause of the ICE, but the reproducer doesn't encounter a precondition check and it ought to.

Author: bors

library/alloc/src/lib.rs

@@ -138,6 +138,7 @@
 #![feature(maybe_uninit_uninit_array_transpose)]
 #![feature(panic_internals)]
 #![feature(pattern)]
+#![feature(ptr_alignment_type)]
 #![feature(ptr_internals)]
 #![feature(ptr_metadata)]
 #![feature(ptr_sub_ptr)]

library/alloc/src/raw_vec.rs

@@ -297,7 +297,20 @@ impl<T, A: Allocator> RawVec<T, A> {
     }
 
     fn current_memory(&self) -> Option<(NonNull<u8>, Layout)> {
-        if T::IS_ZST || self.cap.0 == 0 {
+        // Reduce the amount of code we need to monomorphize per `T`.
+        #[inline]
+        #[rustc_no_mir_inline]
+        unsafe fn inner(size: usize, align: usize, cap: usize) -> Layout {
+            // SAFETY: Precondition guaranteed by the caller
+            unsafe {
+                let size = size.unchecked_mul(cap);
+                Layout::from_size_align_unchecked(size, align)
+            }
+        }
+
+        let cap = self.cap.0;
+
+        if T::IS_ZST || cap == 0 {
             None
         } else {
             // We could use Layout::array here which ensures the absence of isize and usize overflows
@@ -306,10 +319,8 @@ impl<T, A: Allocator> RawVec<T, A> {
             // support such types. So we can do better by skipping some checks and avoid an unwrap.
             const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };
             unsafe {
-                let align = mem::align_of::<T>();
-                let size = mem::size_of::<T>().unchecked_mul(self.cap.0);
-                let layout = Layout::from_size_align_unchecked(size, align);
-                Some((self.ptr.cast().into(), layout))
+                let layout = inner(mem::size_of::<T>(), mem::align_of::<T>(), cap);
+                Some((self.non_null().cast(), layout))
             }
         }
     }

library/core/src/alloc/layout.rs

@@ -4,6 +4,7 @@
 // collections, resulting in having to optimize down excess IR multiple times.
 // Your performance intuition is useless. Run perf.
 
+use crate::assert_unsafe_precondition;
 use crate::cmp;
 use crate::error::Error;
 use crate::fmt;
@@ -118,6 +119,15 @@ impl Layout {
     #[inline]
     #[rustc_allow_const_fn_unstable(ptr_alignment_type)]
     pub const unsafe fn from_size_align_unchecked(size: usize, align: usize) -> Self {
+        assert_unsafe_precondition!(
+            check_library_ub,
+            "Layout::from_size_align_unchecked requires that align is a power of 2 \
+            and the rounded-up allocation size does not exceed isize::MAX",
+            (
+                size: usize = size,
+                align: usize = align,
+            ) => Layout::from_size_align(size, align).is_ok()
+        );
         // SAFETY: the caller is required to uphold the preconditions.
         unsafe { Layout { size, align: Alignment::new_unchecked(align) } }
     }