From 8f432d4ae602cbf629974910a002998c1aa0f653 Mon Sep 17 00:00:00 2001
From: AngelicosPhosphoros <angelicos.phosphoros@protonmail.com>
Date: Sat, 30 Dec 2023 20:51:53 +0100
Subject: [PATCH] Add assume into `NonZeroIntX::get`

LLVM currently don't support range metadata for function arguments so it fails to optimize non zero integers using their invariant if they are provided using by-value function arguments.

Related to https://github.com/rust-lang/rust/issues/119422
Related to https://github.com/llvm/llvm-project/issues/76628
Related to https://github.com/rust-lang/rust/issues/49572
---
 library/core/src/num/nonzero.rs      | 22 ++++++--
 tests/codegen/issues/issue-119422.rs | 83 ++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+), 4 deletions(-)
 create mode 100644 tests/codegen/issues/issue-119422.rs

diff --git a/library/core/src/num/nonzero.rs b/library/core/src/num/nonzero.rs
index f5ecf501ce990..2df38ab5848af 100644
--- a/library/core/src/num/nonzero.rs
+++ b/library/core/src/num/nonzero.rs
@@ -104,6 +104,18 @@ macro_rules! nonzero_integers {
                 #[inline]
                 #[rustc_const_stable(feature = "const_nonzero_get", since = "1.34.0")]
                 pub const fn get(self) -> $Int {
+                    // FIXME: Remove this after LLVM supports `!range` metadata for function
+                    // arguments https://github.com/llvm/llvm-project/issues/76628
+                    //
+                    // Rustc can set range metadata only if it loads `self` from
+                    // memory somewhere. If the value of `self` was from by-value argument
+                    // of some not-inlined function, LLVM don't have range metadata
+                    // to understand that the value cannot be zero.
+
+                    // SAFETY: It is an invariant of this type.
+                    unsafe {
+                        intrinsics::assume(self.0 != 0);
+                    }
                     self.0
                 }
 
@@ -114,7 +126,9 @@ macro_rules! nonzero_integers {
                 #[doc = concat!("Converts a `", stringify!($Ty), "` into an `", stringify!($Int), "`")]
                 #[inline]
                 fn from(nonzero: $Ty) -> Self {
-                    nonzero.0
+                    // Call nonzero to keep information range information
+                    // from get method.
+                    nonzero.get()
                 }
             }
 
@@ -233,7 +247,7 @@ macro_rules! nonzero_leading_trailing_zeros {
                 #[inline]
                 pub const fn leading_zeros(self) -> u32 {
                     // SAFETY: since `self` cannot be zero, it is safe to call `ctlz_nonzero`.
-                    unsafe { intrinsics::ctlz_nonzero(self.0 as $Uint) as u32 }
+                    unsafe { intrinsics::ctlz_nonzero(self.get() as $Uint) as u32 }
                 }
 
                 /// Returns the number of trailing zeros in the binary representation
@@ -257,7 +271,7 @@ macro_rules! nonzero_leading_trailing_zeros {
                 #[inline]
                 pub const fn trailing_zeros(self) -> u32 {
                     // SAFETY: since `self` cannot be zero, it is safe to call `cttz_nonzero`.
-                    unsafe { intrinsics::cttz_nonzero(self.0 as $Uint) as u32 }
+                    unsafe { intrinsics::cttz_nonzero(self.get() as $Uint) as u32 }
                 }
 
             }
@@ -515,7 +529,7 @@ macro_rules! nonzero_unsigned_operations {
                               without modifying the original"]
                 #[inline]
                 pub const fn ilog10(self) -> u32 {
-                    super::int_log10::$Int(self.0)
+                    super::int_log10::$Int(self.get())
                 }
 
                 /// Calculates the middle point of `self` and `rhs`.
diff --git a/tests/codegen/issues/issue-119422.rs b/tests/codegen/issues/issue-119422.rs
new file mode 100644
index 0000000000000..9c99d96317d77
--- /dev/null
+++ b/tests/codegen/issues/issue-119422.rs
@@ -0,0 +1,83 @@
+//! This test checks that compiler don't generate useless compares to zeros
+//! for NonZero integer types.
+
+// compile-flags: -O --edition=2021 -Zmerge-functions=disabled
+// only-64bit (because the LLVM type of i64 for usize shows up)
+
+#![crate_type = "lib"]
+
+use core::num::*;
+use core::ptr::NonNull;
+
+// CHECK-LABEL: @check_non_null
+#[no_mangle]
+pub fn check_non_null(x: NonNull<u8>) -> bool {
+    // CHECK: ret i1 false
+    x.as_ptr().is_null()
+}
+
+// CHECK-LABEL: @equals_zero_is_false_u8
+#[no_mangle]
+pub fn equals_zero_is_false_u8(x: NonZeroU8) -> bool {
+    // CHECK-NOT: br
+    // CHECK: ret i1 false
+    // CHECK-NOT: br
+    x.get() == 0
+}
+
+// CHECK-LABEL: @not_equals_zero_is_true_u8
+#[no_mangle]
+pub fn not_equals_zero_is_true_u8(x: NonZeroU8) -> bool {
+    // CHECK-NOT: br
+    // CHECK: ret i1 true
+    // CHECK-NOT: br
+    x.get() != 0
+}
+
+// CHECK-LABEL: @equals_zero_is_false_i8
+#[no_mangle]
+pub fn equals_zero_is_false_i8(x: NonZeroI8) -> bool {
+    // CHECK-NOT: br
+    // CHECK: ret i1 false
+    // CHECK-NOT: br
+    x.get() == 0
+}
+
+// CHECK-LABEL: @not_equals_zero_is_true_i8
+#[no_mangle]
+pub fn not_equals_zero_is_true_i8(x: NonZeroI8) -> bool {
+    // CHECK-NOT: br
+    // CHECK: ret i1 true
+    // CHECK-NOT: br
+    x.get() != 0
+}
+
+// CHECK-LABEL: @usize_try_from_u32
+#[no_mangle]
+pub fn usize_try_from_u32(x: NonZeroU32) -> NonZeroUsize {
+    // CHECK-NOT: br
+    // CHECK: zext i32 %{{.*}} to i64
+    // CHECK-NOT: br
+    // CHECK: ret i64
+    x.try_into().unwrap()
+}
+
+// CHECK-LABEL: @isize_try_from_i32
+#[no_mangle]
+pub fn isize_try_from_i32(x: NonZeroI32) -> NonZeroIsize {
+    // CHECK-NOT: br
+    // CHECK: sext i32 %{{.*}} to i64
+    // CHECK-NOT: br
+    // CHECK: ret i64
+    x.try_into().unwrap()
+}
+
+// CHECK-LABEL: @u64_from_nonzero_is_not_zero
+#[no_mangle]
+pub fn u64_from_nonzero_is_not_zero(x: NonZeroU64)->bool {
+    // CHECK-NOT: br
+    // CHECK: ret i1 false
+    // CHECK-NOT: br
+    let v: u64 = x.into();
+    v == 0
+}