Take more advantage of the isize::MAX limit in Layout

Things like `padding_needed_for` are current implemented being super careful to handle things like `Layout::size` potentially being `usize::MAX`.

But now that 95295 has happened, that's no longer a concern.  It's possible to add two `Layout::size`s together without risking overflow now.

So this PR adds a wrapper type to allow doing that kind of thing in safe code while still telling LLVM it can't overflow.

library/core/src/alloc/layout.rs

@@ -4,9 +4,12 @@
 // collections, resulting in having to optimize down excess IR multiple times.
 // Your performance intuition is useless. Run perf.
 
+use super::size_in_bytes::SizeInBytes;
 use crate::error::Error;
+use crate::intrinsics::unchecked_sub;
+use crate::mem::SizedTypeProperties;
 use crate::ptr::{Alignment, NonNull};
-use crate::{assert_unsafe_precondition, cmp, fmt, mem};
+use crate::{assert_unsafe_precondition, fmt, mem};
 
 // While this function is used in one place and its implementation
 // could be inlined, the previous attempts to do so made rustc
@@ -37,7 +40,7 @@ const fn size_align<T>() -> (usize, usize) {
 #[lang = "alloc_layout"]
 pub struct Layout {
     // size of the requested block of memory, measured in bytes.
-    size: usize,
+    size: SizeInBytes,
 
     // alignment of the requested block of memory, measured in bytes.
     // we ensure that this is always a power-of-two, because API's
@@ -68,22 +71,22 @@ impl Layout {
     pub const fn from_size_align(size: usize, align: usize) -> Result<Self, LayoutError> {
         if Layout::is_size_align_valid(size, align) {
             // SAFETY: Layout::is_size_align_valid checks the preconditions for this call.
-            unsafe { Ok(Layout { size, align: mem::transmute(align) }) }
+            unsafe { Ok(Layout { size: mem::transmute(size), align: mem::transmute(align) }) }
         } else {
             Err(LayoutError)
         }
     }
 
     const fn is_size_align_valid(size: usize, align: usize) -> bool {
         let Some(align) = Alignment::new(align) else { return false };
-        if size > Self::max_size_for_align(align) {
+        if size > Self::max_size_for_align(align).as_usize() {
             return false;
         }
         true
     }
 
     #[inline(always)]
-    const fn max_size_for_align(align: Alignment) -> usize {
+    const fn max_size_for_align(align: Alignment) -> SizeInBytes {
         // (power-of-two implies align != 0.)
 
         // Rounded up size is:
@@ -98,18 +101,28 @@ impl Layout {
         //
         // Above implies that checking for summation overflow is both
         // necessary and sufficient.
-        isize::MAX as usize - (align.as_usize() - 1)
+
+        // SAFETY: the maximum possible alignment is `isize::MAX + 1`,
+        // so the first subtraction cannot overflow.  The minimum possible
+        // alignment is `1`, so the subtraction returns as most `isize::MAX`,
+        // and thus the calculated `max_size` is guaranteed in-range.
+        unsafe {
+            let max_size = unchecked_sub(isize::MAX as usize + 1, align.as_usize());
+            SizeInBytes::new_unchecked(max_size)
+        }
     }
 
-    /// Internal helper constructor to skip revalidating alignment validity.
+    /// Internal helper constructor to check only the inter-field invariant,
+    /// trusting the types to enforce the per-field invariants.
     #[inline]
-    const fn from_size_alignment(size: usize, align: Alignment) -> Result<Self, LayoutError> {
-        if size > Self::max_size_for_align(align) {
-            return Err(LayoutError);
+    const fn from_size_alignment(size: SizeInBytes, align: Alignment) -> Result<Self, LayoutError> {
+        // FIXME: remove the `as_usize`s once we can use `const PartialOrd`
+        if size.as_usize() <= Self::max_size_for_align(align).as_usize() {
+            // SAFETY: Layout::size invariants checked above.
+            Ok(Layout { size, align })
+        } else {
+            Err(LayoutError)
         }
-
-        // SAFETY: Layout::size invariants checked above.
-        Ok(Layout { size, align })
     }
 
     /// Creates a layout, bypassing all checks.
@@ -134,7 +147,7 @@ impl Layout {
             ) => Layout::is_size_align_valid(size, align)
         );
         // SAFETY: the caller is required to uphold the preconditions.
-        unsafe { Layout { size, align: mem::transmute(align) } }
+        unsafe { Layout { size: mem::transmute(size), align: mem::transmute(align) } }
     }
 
     /// The minimum size in bytes for a memory block of this layout.
@@ -143,7 +156,7 @@ impl Layout {
     #[must_use]
     #[inline]
     pub const fn size(&self) -> usize {
-        self.size
+        self.size.as_usize()
     }
 
     /// The minimum byte alignment for a memory block of this layout.
@@ -252,9 +265,14 @@ impl Layout {
     /// Returns an error if the combination of `self.size()` and the given
     /// `align` violates the conditions listed in [`Layout::from_size_align`].
     #[stable(feature = "alloc_layout_manipulation", since = "1.44.0")]
+    #[rustc_const_unstable(feature = "const_alloc_layout", issue = "67521")]
     #[inline]
-    pub fn align_to(&self, align: usize) -> Result<Self, LayoutError> {
-        Layout::from_size_align(self.size(), cmp::max(self.align(), align))
+    pub const fn align_to(&self, align: usize) -> Result<Self, LayoutError> {
+        if let Some(align) = Alignment::new(align) {
+            Layout::from_size_alignment(self.size, Alignment::max(self.align, align))
+        } else {
+            Err(LayoutError)
+        }
     }
 
     /// Returns the amount of padding we must insert after `self`
@@ -279,29 +297,42 @@ impl Layout {
                   without modifying the `Layout`"]
     #[inline]
     pub const fn padding_needed_for(&self, align: usize) -> usize {
-        let len = self.size();
+        // FIXME: Can we just change the type on this to `Alignment`?
+        let Some(align) = Alignment::new(align) else { return usize::MAX };
+        self.padding_bytes_needed_for(align).as_usize()
+    }
+
+    const fn padding_bytes_needed_for(&self, align: Alignment) -> SizeInBytes {
+        let len = self.size;
+        let align_m1 = SizeInBytes::alignment_minus_one(align);
+        let len_rounded_up = len.add_wide(align_m1) & !align_m1.as_usize();
 
+        // SAFETY:
         // Rounded up value is:
         //   len_rounded_up = (len + align - 1) & !(align - 1);
         // and then we return the padding difference: `len_rounded_up - len`.
         //
-        // We use modular arithmetic throughout:
+        // The arithmetic we do here can never overflow:
         //
         // 1. align is guaranteed to be > 0, so align - 1 is always
         //    valid.
         //
-        // 2. `len + align - 1` can overflow by at most `align - 1`,
-        //    so the &-mask with `!(align - 1)` will ensure that in the
-        //    case of overflow, `len_rounded_up` will itself be 0.
-        //    Thus the returned padding, when added to `len`, yields 0,
-        //    which trivially satisfies the alignment `align`.
+        // 2. len is at most `isize::MAX`, so adding `align - 1` can never
+        //    overflow a `usize`.
         //
-        // (Of course, attempts to allocate blocks of memory whose
-        // size and padding overflow in the above manner should cause
-        // the allocator to yield an error anyway.)
-
-        let len_rounded_up = len.wrapping_add(align).wrapping_sub(1) & !align.wrapping_sub(1);
-        len_rounded_up.wrapping_sub(len)
+        // 3. masking by the alignment can remove at most `align - 1`,
+        //    which is what we just added, so the subtraction cannot overflow.
+        //
+        // 4. the resulting padding is thus at most `align - 1`, but the largest
+        //    possible alignment is `isize::MAX + 1`, and thus the padding
+        //    will always fit in `SizeInBytes`.
+        //
+        // (Size 0 Align MAX is already aligned, so doesn't need any padding,
+        // but Size 1 Align MAX has the largest padding requirement: `isize::MAX`.)
+        unsafe {
+            let padding = unchecked_sub(len_rounded_up, len.as_usize());
+            SizeInBytes::new_unchecked(padding)
+        }
     }
 
     /// Creates a layout by rounding the size of this layout up to a multiple
@@ -315,12 +346,12 @@ impl Layout {
                   without modifying the original"]
     #[inline]
     pub const fn pad_to_align(&self) -> Layout {
-        let pad = self.padding_needed_for(self.align());
+        let pad = self.padding_bytes_needed_for(self.align);
         // This cannot overflow. Quoting from the invariant of Layout:
         // > `size`, when rounded up to the nearest multiple of `align`,
         // > must not overflow isize (i.e., the rounded value must be
         // > less than or equal to `isize::MAX`)
-        let new_size = self.size() + pad;
+        let new_size = self.size.add_wide(pad);
 
         // SAFETY: padded size is guaranteed to not exceed `isize::MAX`.
         unsafe { Layout::from_size_align_unchecked(new_size, self.align()) }
@@ -333,20 +364,36 @@ impl Layout {
     /// layout of the array and `offs` is the distance between the start
     /// of each element in the array.
     ///
+    /// (That distance between elements is sometimes known as "stride".)
+    ///
     /// On arithmetic overflow, returns `LayoutError`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(alloc_layout_extra)]
+    /// use std::alloc::Layout;
+    ///
+    /// // All rust types have a size that's a multiple of their alignment.
+    /// let normal = Layout::from_size_align(12, 4).unwrap();
+    /// let repeated = normal.repeat(3).unwrap();
+    /// assert_eq!(repeated, (Layout::from_size_align(36, 4).unwrap(), 12));
+    ///
+    /// // But you can manually make layouts which don't meet that rule.
+    /// let padding_needed = Layout::from_size_align(6, 4).unwrap();
+    /// let repeated = padding_needed.repeat(3).unwrap();
+    /// assert_eq!(repeated, (Layout::from_size_align(24, 4).unwrap(), 8));
+    /// ```
     #[unstable(feature = "alloc_layout_extra", issue = "55724")]
+    #[rustc_const_unstable(feature = "const_alloc_layout", issue = "67521")]
     #[inline]
-    pub fn repeat(&self, n: usize) -> Result<(Self, usize), LayoutError> {
-        // This cannot overflow. Quoting from the invariant of Layout:
-        // > `size`, when rounded up to the nearest multiple of `align`,
-        // > must not overflow isize (i.e., the rounded value must be
-        // > less than or equal to `isize::MAX`)
-        let padded_size = self.size() + self.padding_needed_for(self.align());
-        let alloc_size = padded_size.checked_mul(n).ok_or(LayoutError)?;
-
-        // The safe constructor is called here to enforce the isize size limit.
-        let layout = Layout::from_size_alignment(alloc_size, self.align)?;
-        Ok((layout, padded_size))
+    pub const fn repeat(&self, n: usize) -> Result<(Self, usize), LayoutError> {
+        let padded = self.pad_to_align();
+        if let Ok(repeated) = padded.repeat_packed(n) {
+            Ok((repeated, padded.size()))
+        } else {
+            Err(LayoutError)
+        }
     }
 
     /// Creates a layout describing the record for `self` followed by
@@ -395,17 +442,20 @@ impl Layout {
     /// # assert_eq!(repr_c(&[u64, u32, u16, u32]), Ok((s, vec![0, 8, 12, 16])));
     /// ```
     #[stable(feature = "alloc_layout_manipulation", since = "1.44.0")]
+    #[rustc_const_unstable(feature = "const_alloc_layout", issue = "67521")]
     #[inline]
-    pub fn extend(&self, next: Self) -> Result<(Self, usize), LayoutError> {
-        let new_align = cmp::max(self.align, next.align);
-        let pad = self.padding_needed_for(next.align());
-
-        let offset = self.size().checked_add(pad).ok_or(LayoutError)?;
-        let new_size = offset.checked_add(next.size()).ok_or(LayoutError)?;
-
-        // The safe constructor is called here to enforce the isize size limit.
-        let layout = Layout::from_size_alignment(new_size, new_align)?;
-        Ok((layout, offset))
+    pub const fn extend(&self, next: Self) -> Result<(Self, usize), LayoutError> {
+        let new_align = Alignment::max(self.align, next.align);
+        let pad = self.padding_bytes_needed_for(next.align);
+
+        if let Some(offset) = self.size.checked_add(pad)
+            && let Some(new_size) = offset.checked_add(next.size)
+            && let Ok(layout) = Layout::from_size_alignment(new_size, new_align)
+        {
+            Ok((layout, offset.as_usize()))
+        } else {
+            Err(LayoutError)
+        }
     }
 
     /// Creates a layout describing the record for `n` instances of
@@ -421,11 +471,15 @@ impl Layout {
     ///
     /// On arithmetic overflow, returns `LayoutError`.
     #[unstable(feature = "alloc_layout_extra", issue = "55724")]
+    #[rustc_const_unstable(feature = "const_alloc_layout", issue = "67521")]
     #[inline]
-    pub fn repeat_packed(&self, n: usize) -> Result<Self, LayoutError> {
-        let size = self.size().checked_mul(n).ok_or(LayoutError)?;
-        // The safe constructor is called here to enforce the isize size limit.
-        Layout::from_size_alignment(size, self.align)
+    pub const fn repeat_packed(&self, n: usize) -> Result<Self, LayoutError> {
+        if let Some(size) = self.size.checked_mul(n) {
+            // The safe constructor is called here to enforce the isize size limit.
+            Layout::from_size_alignment(size, self.align)
+        } else {
+            Err(LayoutError)
+        }
     }
 
     /// Creates a layout describing the record for `self` followed by
@@ -435,11 +489,15 @@ impl Layout {
     ///
     /// On arithmetic overflow, returns `LayoutError`.
     #[unstable(feature = "alloc_layout_extra", issue = "55724")]
+    #[rustc_const_unstable(feature = "const_alloc_layout", issue = "67521")]
     #[inline]
-    pub fn extend_packed(&self, next: Self) -> Result<Self, LayoutError> {
-        let new_size = self.size().checked_add(next.size()).ok_or(LayoutError)?;
-        // The safe constructor is called here to enforce the isize size limit.
-        Layout::from_size_alignment(new_size, self.align)
+    pub const fn extend_packed(&self, next: Self) -> Result<Self, LayoutError> {
+        if let Some(new_size) = self.size.checked_add(next.size) {
+            // The safe constructor is called here to enforce the isize size limit.
+            Layout::from_size_alignment(new_size, self.align)
+        } else {
+            Err(LayoutError)
+        }
     }
 
     /// Creates a layout describing the record for a `[T; n]`.
@@ -451,21 +509,21 @@ impl Layout {
     #[inline]
     pub const fn array<T>(n: 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, n);
 
         #[inline]
-        const fn inner(
-            element_size: usize,
-            align: Alignment,
-            n: usize,
-        ) -> Result<Layout, LayoutError> {
+        const fn inner(element_layout: Layout, n: usize) -> Result<Layout, LayoutError> {
+            let Layout { size, align } = element_layout;
+            let element_size = size.as_usize();
+
             // 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 && n > Layout::max_size_for_align(align).as_usize() / element_size
+            {
                 return Err(LayoutError);
             }
 

library/core/src/alloc/mod.rs

@@ -4,6 +4,7 @@
 
 mod global;
 mod layout;
+mod size_in_bytes;
 
 #[stable(feature = "global_alloc", since = "1.28.0")]
 pub use self::global::GlobalAlloc;