Also get `add nuw` from `uN::checked_add`

4 files changed, 47 insertions, 32 deletions

diff --git a/library/core/src/num/uint_macros.rs b/library/core/src/num/uint_macros.rs
index 00450c2cda3..42f23f7ed71 100644
--- a/library/core/src/num/uint_macros.rs
+++ b/library/core/src/num/uint_macros.rs
@@ -455,8 +455,19 @@ macro_rules! uint_impl {
                       without modifying the original"]
         #[inline]
         pub const fn checked_add(self, rhs: Self) -> Option<Self> {
-            let (a, b) = self.overflowing_add(rhs);
-            if unlikely!(b) { None } else { Some(a) }
+            // This used to use `overflowing_add`, but that means it ends up being
+            // a `wrapping_add`, losing some optimization opportunities. Notably,
+            // phrasing it this way helps `.checked_add(1)` optimize to a check
+            // against `MAX` and a `add nuw`.
+            // Per <https://github.com/rust-lang/rust/pull/124114#issuecomment-2066173305>,
+            // LLVM is happy to re-form the intrinsic later if useful.
+
+            if unlikely!(intrinsics::add_with_overflow(self, rhs).1) {
+                None
+            } else {
+                // SAFETY: Just checked it doesn't overflow
+                Some(unsafe { intrinsics::unchecked_add(self, rhs) })
+            }
         }
 
         /// Strict integer addition. Computes `self + rhs`, panicking
diff --git a/tests/codegen/checked_math.rs b/tests/codegen/checked_math.rs
index 41016e3b7be..75df5866d6e 100644
--- a/tests/codegen/checked_math.rs
+++ b/tests/codegen/checked_math.rs
@@ -84,3 +84,17 @@ pub fn checked_shr_signed(a: i32, b: u32) -> Option<i32> {
     // CHECK: ret { i32, i32 } %[[R1]]
     a.checked_shr(b)
 }
+
+// CHECK-LABEL: @checked_add_one_unwrap_unsigned
+// CHECK-SAME: (i32 noundef %x)
+#[no_mangle]
+pub fn checked_add_one_unwrap_unsigned(x: u32) -> u32 {
+    // CHECK: %[[IS_MAX:.+]] = icmp eq i32 %x, -1
+    // CHECK: br i1 %[[IS_MAX]], label %[[NONE_BB:.+]], label %[[SOME_BB:.+]]
+    // CHECK: [[SOME_BB]]:
+    // CHECK: %[[R:.+]] = add nuw i32 %x, 1
+    // CHECK: ret i32 %[[R]]
+    // CHECK: [[NONE_BB]]:
+    // CHECK: call {{.+}}unwrap_failed
+    x.checked_add(1).unwrap()
+}
diff --git a/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-abort.mir b/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-abort.mir
index e31a8cb6937..69c11ebcacc 100644
--- a/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-abort.mir
+++ b/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-abort.mir
@@ -11,15 +11,9 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
         scope 3 (inlined <u16 as Step>::forward_checked) {
             scope 4 {
                 scope 6 (inlined core::num::<impl u16>::checked_add) {
+                    let mut _5: (u16, bool);
+                    let mut _6: bool;
                     let mut _7: bool;
-                    scope 7 {
-                    }
-                    scope 8 (inlined core::num::<impl u16>::overflowing_add) {
-                        let mut _5: (u16, bool);
-                        let _6: bool;
-                        scope 9 {
-                        }
-                    }
                 }
             }
             scope 5 (inlined convert::num::ptr_try_from_impls::<impl TryFrom<usize> for u16>::try_from) {
@@ -27,11 +21,11 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
                 let mut _4: u16;
             }
         }
-        scope 10 (inlined Option::<u16>::is_none) {
-            scope 11 (inlined Option::<u16>::is_some) {
+        scope 7 (inlined Option::<u16>::is_none) {
+            scope 8 (inlined Option::<u16>::is_some) {
             }
         }
-        scope 12 (inlined core::num::<impl u16>::wrapping_add) {
+        scope 9 (inlined core::num::<impl u16>::wrapping_add) {
         }
     }
 
@@ -45,12 +39,11 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
     bb1: {
         _4 = _2 as u16 (IntToInt);
         StorageDead(_3);
+        StorageLive(_7);
         StorageLive(_6);
         StorageLive(_5);
         _5 = AddWithOverflow(_1, _4);
         _6 = (_5.1: bool);
-        StorageDead(_5);
-        StorageLive(_7);
         _7 = unlikely(move _6) -> [return: bb2, unwind unreachable];
     }
 
@@ -59,14 +52,16 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
     }
 
     bb3: {
-        StorageDead(_7);
+        StorageDead(_5);
         StorageDead(_6);
+        StorageDead(_7);
         goto -> bb7;
     }
 
     bb4: {
-        StorageDead(_7);
+        StorageDead(_5);
         StorageDead(_6);
+        StorageDead(_7);
         goto -> bb6;
     }
 
diff --git a/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-unwind.mir b/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-unwind.mir
index 8cc9be27e21..e6ea6c51001 100644
--- a/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-unwind.mir
+++ b/tests/mir-opt/pre-codegen/checked_ops.step_forward.PreCodegen.after.panic-unwind.mir
@@ -11,15 +11,9 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
         scope 3 (inlined <u16 as Step>::forward_checked) {
             scope 4 {
                 scope 6 (inlined core::num::<impl u16>::checked_add) {
+                    let mut _5: (u16, bool);
+                    let mut _6: bool;
                     let mut _7: bool;
-                    scope 7 {
-                    }
-                    scope 8 (inlined core::num::<impl u16>::overflowing_add) {
-                        let mut _5: (u16, bool);
-                        let _6: bool;
-                        scope 9 {
-                        }
-                    }
                 }
             }
             scope 5 (inlined convert::num::ptr_try_from_impls::<impl TryFrom<usize> for u16>::try_from) {
@@ -27,11 +21,11 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
                 let mut _4: u16;
             }
         }
-        scope 10 (inlined Option::<u16>::is_none) {
-            scope 11 (inlined Option::<u16>::is_some) {
+        scope 7 (inlined Option::<u16>::is_none) {
+            scope 8 (inlined Option::<u16>::is_some) {
             }
         }
-        scope 12 (inlined core::num::<impl u16>::wrapping_add) {
+        scope 9 (inlined core::num::<impl u16>::wrapping_add) {
         }
     }
 
@@ -45,12 +39,11 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
     bb1: {
         _4 = _2 as u16 (IntToInt);
         StorageDead(_3);
+        StorageLive(_7);
         StorageLive(_6);
         StorageLive(_5);
         _5 = AddWithOverflow(_1, _4);
         _6 = (_5.1: bool);
-        StorageDead(_5);
-        StorageLive(_7);
         _7 = unlikely(move _6) -> [return: bb2, unwind unreachable];
     }
 
@@ -59,14 +52,16 @@ fn step_forward(_1: u16, _2: usize) -> u16 {
     }
 
     bb3: {
-        StorageDead(_7);
+        StorageDead(_5);
         StorageDead(_6);
+        StorageDead(_7);
         goto -> bb7;
     }
 
     bb4: {
-        StorageDead(_7);
+        StorageDead(_5);
         StorageDead(_6);
+        StorageDead(_7);
         goto -> bb6;
     }