From 28c2d32ca7ca6014a52001153a770e8f362e5344 Mon Sep 17 00:00:00 2001
From: Scott McMurray <scottmcm@users.noreply.github.com>
Date: Fri, 24 Nov 2023 02:55:22 -0800
Subject: [PATCH] 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...
---
 library/core/src/alloc/layout.rs    | 26 +++++++++-----------------
 tests/codegen/layout-size-checks.rs | 17 ++++++++++++-----
 2 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/library/core/src/alloc/layout.rs b/library/core/src/alloc/layout.rs
index 65946e09ff9ba..aa01d41274a8a 100644
--- a/library/core/src/alloc/layout.rs
+++ b/library/core/src/alloc/layout.rs
@@ -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)
         }
     }
 }
diff --git a/tests/codegen/layout-size-checks.rs b/tests/codegen/layout-size-checks.rs
index d067cc10a948c..91015ab1bcc01 100644
--- a/tests/codegen/layout-size-checks.rs
+++ b/tests/codegen/layout-size-checks.rs
@@ -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()
 }
 
@@ -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)