From 0866f366270d523cac95f9fe58837c97487396c3 Mon Sep 17 00:00:00 2001
From: EFanZh <efanzh@gmail.com>
Date: Mon, 1 Jan 2024 21:51:04 +0800
Subject: [PATCH] Assert the underlying values of `NonZero*` values are not
 zero

---
 library/core/src/num/nonzero.rs |  12 +++-
 tests/codegen/nonzero.rs        | 106 ++++++++++++++++++++++++++++++++
 2 files changed, 117 insertions(+), 1 deletion(-)
 create mode 100644 tests/codegen/nonzero.rs

diff --git a/library/core/src/num/nonzero.rs b/library/core/src/num/nonzero.rs
index f5ecf501ce990..6f37079b9b492 100644
--- a/library/core/src/num/nonzero.rs
+++ b/library/core/src/num/nonzero.rs
@@ -1,6 +1,7 @@
 //! Definitions of integer that is known not to equal zero.
 
 use crate::fmt;
+use crate::hint;
 use crate::ops::{BitOr, BitOrAssign, Div, Neg, Rem};
 use crate::str::FromStr;
 
@@ -104,6 +105,14 @@ macro_rules! nonzero_integers {
                 #[inline]
                 #[rustc_const_stable(feature = "const_nonzero_get", since = "1.34.0")]
                 pub const fn get(self) -> $Int {
+                    // SAFETY: nonzero types do not contain zero values.
+                    unsafe {
+                        // We can’t use `intrinsics::assume(self.0 != 0)` here because it is not const-stable.
+                        if self.0 == 0 {
+                            hint::unreachable_unchecked();
+                        }
+                    }
+
                     self.0
                 }
 
@@ -114,7 +123,8 @@ macro_rules! nonzero_integers {
                 #[doc = concat!("Converts a `", stringify!($Ty), "` into an `", stringify!($Int), "`")]
                 #[inline]
                 fn from(nonzero: $Ty) -> Self {
-                    nonzero.0
+                    // Utilize the non-zero assertion from the `get` function.
+                    nonzero.get()
                 }
             }
 
diff --git a/tests/codegen/nonzero.rs b/tests/codegen/nonzero.rs
new file mode 100644
index 0000000000000..728d3c20f553e
--- /dev/null
+++ b/tests/codegen/nonzero.rs
@@ -0,0 +1,106 @@
+// compile-flags: -C opt-level=1 -Z merge-functions=disabled
+
+#![crate_type = "lib"]
+
+use core::num::{
+    NonZeroI128, NonZeroI16, NonZeroI32, NonZeroI64, NonZeroI8, NonZeroIsize, NonZeroU128,
+    NonZeroU16, NonZeroU32, NonZeroU64, NonZeroU8, NonZeroUsize,
+};
+
+extern crate core;
+
+// CHECK-LABEL: void @non_zero_i8_is_not_zero(i8
+#[no_mangle]
+fn non_zero_i8_is_not_zero(x: NonZeroI8) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(i8::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_i16_is_not_zero(i16
+#[no_mangle]
+fn non_zero_i16_is_not_zero(x: NonZeroI16) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(i16::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_i32_is_not_zero(i32
+#[no_mangle]
+fn non_zero_i32_is_not_zero(x: NonZeroI32) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(i32::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_i64_is_not_zero(i64
+#[no_mangle]
+fn non_zero_i64_is_not_zero(x: NonZeroI64) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(i64::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_i128_is_not_zero(i128
+#[no_mangle]
+fn non_zero_i128_is_not_zero(x: NonZeroI128) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(i128::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_isize_is_not_zero(i
+#[no_mangle]
+fn non_zero_isize_is_not_zero(x: NonZeroIsize) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(isize::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_u8_is_not_zero(i8
+#[no_mangle]
+fn non_zero_u8_is_not_zero(x: NonZeroU8) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(u8::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_u16_is_not_zero(i16
+#[no_mangle]
+fn non_zero_u16_is_not_zero(x: NonZeroU16) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(u16::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_u32_is_not_zero(i32
+#[no_mangle]
+fn non_zero_u32_is_not_zero(x: NonZeroU32) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(u32::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_u64_is_not_zero(i64
+#[no_mangle]
+fn non_zero_u64_is_not_zero(x: NonZeroU64) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(u64::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_u128_is_not_zero(i128
+#[no_mangle]
+fn non_zero_u128_is_not_zero(x: NonZeroU128) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(u128::from(x), 0);
+}
+
+// CHECK-LABEL: void @non_zero_usize_is_not_zero(i
+#[no_mangle]
+fn non_zero_usize_is_not_zero(x: NonZeroUsize) {
+    // CHECK-NOT: br
+    assert_ne!(x.get(), 0);
+    assert_ne!(usize::from(x), 0);
+}