From 3acb445f15d69c982c37df15a27d06b43ca6e388 Mon Sep 17 00:00:00 2001
From: SOFe <sofe2038@gmail.com>
Date: Sun, 30 Oct 2022 14:49:24 +0800
Subject: [PATCH] Added assert_unsafe_precondition! check for
 NonZeroXxx::from_mut_unchecked

---
 library/core/src/num/nonzero.rs | 33 ++++++++++++++++++++-------------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/library/core/src/num/nonzero.rs b/library/core/src/num/nonzero.rs
index 582c6796781..53ee6efcba7 100644
--- a/library/core/src/num/nonzero.rs
+++ b/library/core/src/num/nonzero.rs
@@ -171,18 +171,6 @@ macro_rules! nonzero_integer {
                 }
             }
 
-            /// Converts a primitive mutable reference to a non-zero mutable reference
-            /// if the referenced integer is not zero.
-            #[unstable(feature = "nonzero_from_mut", issue = "none")]
-            #[must_use]
-            #[inline]
-            pub fn from_mut(n: &mut $Int) -> Option<&mut Self> {
-                // SAFETY: Self is repr(transparent), and the value is non-zero.
-                // As long as the returned reference is alive,
-                // the user cannot `*n = 0` directly.
-                (*n != 0).then(|| unsafe { &mut *(n as *mut $Int as *mut Self) })
-            }
-
             /// Converts a primitive mutable reference to a non-zero mutable reference
             /// without checking whether the referenced value is non-zero.
             /// This results in undefined behavior if `*n` is zero.
@@ -194,7 +182,26 @@ macro_rules! nonzero_integer {
             #[inline]
             pub unsafe fn from_mut_unchecked(n: &mut $Int) -> &mut Self {
                 // SAFETY: Self is repr(transparent), and the value is assumed to be non-zero.
-                unsafe { &mut *(n as *mut $Int as *mut Self) }
+                unsafe {
+                    let n_alias = &mut *n;
+                    core::intrinsics::assert_unsafe_precondition!(
+                        concat!(stringify!($Ty), "::from_mut_unchecked requires the argument to dereference as non-zero"),
+                        (n_alias: &mut $Int) => *n_alias != 0
+                    );
+                    &mut *(n as *mut $Int as *mut Self)
+                }
+            }
+
+            /// Converts a primitive mutable reference to a non-zero mutable reference
+            /// if the referenced integer is not zero.
+            #[unstable(feature = "nonzero_from_mut", issue = "none")]
+            #[must_use]
+            #[inline]
+            pub fn from_mut(n: &mut $Int) -> Option<&mut Self> {
+                // SAFETY: Self is repr(transparent), and the value is non-zero.
+                // As long as the returned reference is alive,
+                // the user cannot `*n = 0` directly.
+                (*n != 0).then(|| unsafe { &mut *(n as *mut $Int as *mut Self) })
             }
 
             /// Returns the value as a primitive type.