Skip to content

Use usize::checked_mul for Layout::array checks again #118239

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed

Use usize::checked_mul for Layout::array checks again #118239

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 9 additions & 17 deletions library/core/src/alloc/layout.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -432,30 +432,22 @@ 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);
// By multiplying by twice the size of the type, we can use `usize`
// overflow checking to enforce the `isize::MAX` limits, and since
// a single `T` can never be more than `isize::MAX` either, doubling
// it can never overflow.
return inner(mem::size_of::<T>() * 2, Alignment::of::<T>(), n);

#[inline]
const fn inner(
element_size: usize,
element_size_times_two: usize,
align: Alignment,
n: 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 {
let Some(total_size_times_two) = n.checked_mul(element_size_times_two) else {
return Err(LayoutError);
}

let array_size = element_size * 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())) }
};
Layout::from_size_alignment(total_size_times_two >> 1, align)
}
}
}
Expand Down
17 changes: 12 additions & 5 deletions tests/codegen/layout-size-checks.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,16 @@ type RGB48 = [u16; 3];
// CHECK-LABEL: @layout_array_rgb48
#[no_mangle]
pub fn layout_array_rgb48(n: usize) -> Layout {
// CHECK-NOT: llvm.umul.with.overflow.i64
// CHECK: icmp ugt i64 %n, 1537228672809129301
// CHECK-NOT: llvm.umul.with.overflow.i64
// CHECK: mul nuw nsw i64 %n, 6
// CHECK-NOT: llvm.umul.with.overflow.i64
// CHECK-NOT: icmp
// CHECK-NOT: mul
// CHECK: %[[TUP:.+]] = tail call { i64, i1 } @llvm.umul.with.overflow.i64(i64 %n, i64 12)
// CHECK-NOT: icmp
// CHECK-NOT: mul
// CHECK: %[[PROD:.+]] = extractvalue { i64, i1 } %[[TUP]], 0
// CHECK-NEXT: lshr exact i64 %[[PROD]], 1
// CHECK-NOT: icmp
// CHECK-NOT: mul

Layout::array::<RGB48>(n).unwrap()
}

Expand All @@ -29,3 +34,5 @@ pub fn layout_array_i32(n: usize) -> Layout {
// CHECK-NOT: llvm.umul.with.overflow.i64
Layout::array::<i32>(n).unwrap()
}

// CHECK: declare { i64, i1 } @llvm.umul.with.overflow.i64(i64, i64)