Skip to content

Commit 28c2d32

Browse files
committed
Use usize::checked_mul for Layout::array checks again
I found a little trick that allows us to check the `isize::MAX` limit using `usize::checked_mul`. No idea whether that's a good idea, but let's see what perf has to say...
1 parent 1e9dda7 commit 28c2d32

File tree

2 files changed

+21
-22
lines changed

2 files changed

+21
-22
lines changed

library/core/src/alloc/layout.rs

Lines changed: 9 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -432,30 +432,22 @@ impl Layout {
432432
#[inline]
433433
pub const fn array<T>(n: usize) -> Result<Self, LayoutError> {
434434
// Reduce the amount of code we need to monomorphize per `T`.
435-
return inner(mem::size_of::<T>(), Alignment::of::<T>(), n);
435+
// By multiplying by twice the size of the type, we can use `usize`
436+
// overflow checking to enforce the `isize::MAX` limits, and since
437+
// a single `T` can never be more than `isize::MAX` either, doubling
438+
// it can never overflow.
439+
return inner(mem::size_of::<T>() * 2, Alignment::of::<T>(), n);
436440

437441
#[inline]
438442
const fn inner(
439-
element_size: usize,
443+
element_size_times_two: usize,
440444
align: Alignment,
441445
n: usize,
442446
) -> Result<Layout, LayoutError> {
443-
// We need to check two things about the size:
444-
// - That the total size won't overflow a `usize`, and
445-
// - That the total size still fits in an `isize`.
446-
// By using division we can check them both with a single threshold.
447-
// That'd usually be a bad idea, but thankfully here the element size
448-
// and alignment are constants, so the compiler will fold all of it.
449-
if element_size != 0 && n > Layout::max_size_for_align(align) / element_size {
447+
let Some(total_size_times_two) = n.checked_mul(element_size_times_two) else {
450448
return Err(LayoutError);
451-
}
452-
453-
let array_size = element_size * n;
454-
455-
// SAFETY: We just checked above that the `array_size` will not
456-
// exceed `isize::MAX` even when rounded up to the alignment.
457-
// And `Alignment` guarantees it's a power of two.
458-
unsafe { Ok(Layout::from_size_align_unchecked(array_size, align.as_usize())) }
449+
};
450+
Layout::from_size_alignment(total_size_times_two >> 1, align)
459451
}
460452
}
461453
}

tests/codegen/layout-size-checks.rs

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,16 @@ type RGB48 = [u16; 3];
1111
// CHECK-LABEL: @layout_array_rgb48
1212
#[no_mangle]
1313
pub fn layout_array_rgb48(n: usize) -> Layout {
14-
// CHECK-NOT: llvm.umul.with.overflow.i64
15-
// CHECK: icmp ugt i64 %n, 1537228672809129301
16-
// CHECK-NOT: llvm.umul.with.overflow.i64
17-
// CHECK: mul nuw nsw i64 %n, 6
18-
// CHECK-NOT: llvm.umul.with.overflow.i64
14+
// CHECK-NOT: icmp
15+
// CHECK-NOT: mul
16+
// CHECK: %[[TUP:.+]] = tail call { i64, i1 } @llvm.umul.with.overflow.i64(i64 %n, i64 12)
17+
// CHECK-NOT: icmp
18+
// CHECK-NOT: mul
19+
// CHECK: %[[PROD:.+]] = extractvalue { i64, i1 } %[[TUP]], 0
20+
// CHECK-NEXT: lshr exact i64 %[[PROD]], 1
21+
// CHECK-NOT: icmp
22+
// CHECK-NOT: mul
23+
1924
Layout::array::<RGB48>(n).unwrap()
2025
}
2126

@@ -29,3 +34,5 @@ pub fn layout_array_i32(n: usize) -> Layout {
2934
// CHECK-NOT: llvm.umul.with.overflow.i64
3035
Layout::array::<i32>(n).unwrap()
3136
}
37+
38+
// CHECK: declare { i64, i1 } @llvm.umul.with.overflow.i64(i64, i64)

0 commit comments

Comments
 (0)