Add a Layout::array_unchecked for RawVec to use

Instead of it writing the code for it itself.

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,9 +4,11 @@
 // 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::intrinsics;
 use crate::mem;
 use crate::ptr::{Alignment, NonNull};
 
@@ -425,43 +427,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(Layout::new::<T>(), 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(Layout::new::<T>(), 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 }
         }
     }
 }