Skip to content

Commit

Permalink
Auto merge of rust-lang#129845 - scottmcm:redo-layout, r=<try>
Browse files Browse the repository at this point in the history
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 rust-lang#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.
  • Loading branch information
bors committed Sep 1, 2024
2 parents a48861a + 972ff01 commit 9c711ef
Show file tree
Hide file tree
Showing 8 changed files with 223 additions and 84 deletions.
195 changes: 127 additions & 68 deletions library/core/src/alloc/layout.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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.
Expand All @@ -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.
Expand All @@ -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.
Expand Down Expand Up @@ -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`
Expand All @@ -279,29 +297,43 @@ 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()
}

#[inline]
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
Expand All @@ -315,12 +347,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()) }
Expand All @@ -333,20 +365,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
Expand Down Expand Up @@ -395,17 +443,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
Expand All @@ -421,11 +472,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
Expand All @@ -435,11 +490,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]`.
Expand All @@ -451,21 +510,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);
}

Expand Down
1 change: 1 addition & 0 deletions library/core/src/alloc/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading

0 comments on commit 9c711ef

Please sign in to comment.