Make Layout::new::<T>() a constant instead of multiple NullOps

Author: scottmcm

library/alloc/src/raw_vec.rs

@@ -300,17 +300,9 @@ impl<T, A: Allocator> RawVec<T, A> {
         if T::IS_ZST || self.cap.0 == 0 {
             None
         } else {
-            // We could use Layout::array here which ensures the absence of isize and usize overflows
-            // and could hypothetically handle differences between stride and size, but this memory
-            // has already been allocated so we know it can't overflow and currently Rust does not
-            // 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))
-            }
+            // SAFETY: This memory has already been allocated so we know it can't overflow
+            let layout = unsafe { Layout::array_unchecked::<T>(self.cap.0) };
+            Some((self.ptr.cast().into(), layout))
         }
     }
 

library/core/src/alloc/layout.rs

@@ -4,22 +4,14 @@
 // 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;
-use crate::mem;
+use crate::intrinsics;
+use crate::mem::{self, SizedTypeProperties};
 use crate::ptr::{Alignment, NonNull};
 
-// While this function is used in one place and its implementation
-// could be inlined, the previous attempts to do so made rustc
-// slower:
-//
-// * https://github.com/rust-lang/rust/pull/72189
-// * https://github.com/rust-lang/rust/pull/79827
-const fn size_align<T>() -> (usize, usize) {
-    (mem::size_of::<T>(), mem::align_of::<T>())
-}
-
 /// Layout of a block of memory.
 ///
 /// An instance of `Layout` describes a particular layout of memory.
@@ -150,11 +142,7 @@ impl Layout {
     #[must_use]
     #[inline]
     pub const fn new<T>() -> Self {
-        let (size, align) = size_align::<T>();
-        // SAFETY: if the type is instantiated, rustc already ensures that its
-        // layout is valid. Use the unchecked constructor to avoid inserting a
-        // panicking codepath that needs to be optimized out.
-        unsafe { Layout::from_size_align_unchecked(size, align) }
+        T::LAYOUT
     }
 
     /// Produces layout describing a record that could be used to
@@ -425,43 +413,76 @@ impl Layout {
         Layout::from_size_alignment(new_size, self.align)
     }
 
-    /// Creates a layout describing the record for a `[T; n]`.
+    /// Creates a layout describing the record for a `[T; len]`.
     ///
     /// On arithmetic overflow or when the total size would exceed
     /// `isize::MAX`, returns `LayoutError`.
     #[stable(feature = "alloc_layout_manipulation", since = "1.44.0")]
     #[rustc_const_unstable(feature = "const_alloc_layout", issue = "67521")]
     #[inline]
-    pub const fn array<T>(n: usize) -> Result<Self, LayoutError> {
+    pub const fn array<T>(len: usize) -> Result<Self, LayoutError> {
         // Reduce the amount of code we need to monomorphize per `T`.
-        return inner(mem::size_of::<T>(), Alignment::of::<T>(), n);
+        return inner(T::LAYOUT, len);
 
         #[inline]
-        const fn inner(
-            element_size: usize,
-            align: Alignment,
-            n: usize,
-        ) -> Result<Layout, LayoutError> {
+        const fn inner(element: Layout, len: usize) -> Result<Layout, LayoutError> {
             // We need to check two things about the size:
             //  - That the total size won't overflow a `usize`, and
             //  - That the total size still fits in an `isize`.
             // By using division we can check them both with a single threshold.
             // That'd usually be a bad idea, but thankfully here the element size
             // and alignment are constants, so the compiler will fold all of it.
-            if element_size != 0 && n > Layout::max_size_for_align(align) / element_size {
+            if element.size != 0 && len > Layout::max_size_for_align(element.align) / element.size {
                 return Err(LayoutError);
             }
 
-            // SAFETY: We just checked that we won't overflow `usize` when we multiply.
+            // SAFETY: We just checked above that the `array_size` will not
+            // exceed `isize::MAX` even when rounded up to the alignment.
+            unsafe { Ok(Layout::array_of_unchecked(element, len)) }
+        }
+    }
+
+    /// Creates a layout describing the record for a `[T; n]`, without checking
+    /// whether the `len` is an allowed length.
+    ///
+    /// # Safety
+    ///
+    /// `len` must be small enough that the total size *in bytes* of the
+    /// resulting layout is `isize::MAX` or less.
+    ///
+    /// Equivalently, either `T` is a ZST or `len <= isize::MAX / size_of::<T>()`.
+    #[unstable(feature = "alloc_layout_extra", issue = "55724")]
+    #[rustc_const_unstable(feature = "const_alloc_layout", issue = "67521")]
+    #[inline]
+    pub const unsafe fn array_unchecked<T>(len: usize) -> Self {
+        // SAFETY: `T` is a Rust type, so its size is a multiple of align,
+        // and our safety precondition ensures that len is small enough.
+        unsafe { Self::array_of_unchecked(T::LAYOUT, len) }
+    }
+
+    /// # Safety
+    /// - The element layout must be that of a rust type, with size a multiple of align
+    /// - The total size of the resulting array must be at most `isize::MAX`
+    #[inline]
+    const unsafe fn array_of_unchecked(element: Layout, len: usize) -> Layout {
+        assert_unsafe_precondition!(
+            check_language_ub,
+            "array_of_unchecked element must have size an integer multiple of \
+             alignment and the total length must be be isize::MAX or less",
+            (
+                size: usize = element.size,
+                align: usize = element.align.as_usize(),
+                len: usize = len,
+            ) => size % align == 0 && (size as u128) * (len as u128) <= (isize::MAX as u128),
+        );
+
+        // SAFETY: Out preconditions are exactly what's needed here
+        unsafe {
             // This is a useless hint inside this function, but after inlining this helps
             // deduplicate checks for whether the overall capacity is zero (e.g., in RawVec's
             // allocation path) before/after this multiplication.
-            let array_size = unsafe { element_size.unchecked_mul(n) };
-
-            // SAFETY: We just checked above that the `array_size` will not
-            // exceed `isize::MAX` even when rounded up to the alignment.
-            // And `Alignment` guarantees it's a power of two.
-            unsafe { Ok(Layout::from_size_align_unchecked(array_size, align.as_usize())) }
+            let bytes = intrinsics::unchecked_mul(element.size, len);
+            Self { size: bytes, align: element.align }
         }
     }
 }

library/core/src/mem/mod.rs

@@ -5,6 +5,7 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
+use crate::alloc::Layout;
 use crate::clone;
 use crate::cmp;
 use crate::fmt;
@@ -1235,6 +1236,14 @@ pub trait SizedTypeProperties: Sized {
     #[doc(hidden)]
     #[unstable(feature = "sized_type_properties", issue = "none")]
     const IS_ZST: bool = size_of::<Self>() == 0;
+
+    #[doc(hidden)]
+    #[unstable(feature = "sized_type_properties", issue = "none")]
+    const LAYOUT: Layout =
+        // SAFETY: if the type is instantiated, rustc already ensures that its
+        // layout is valid. Use the unchecked constructor to avoid inserting a
+        // panicking codepath that needs to be optimized out.
+        unsafe { Layout::from_size_align_unchecked(size_of::<Self>(), align_of::<Self>()) };
 }
 #[doc(hidden)]
 #[unstable(feature = "sized_type_properties", issue = "none")]