Rollup merge of #155265 - Apersoma:isqrt-smarter, r=jhpratt,tgross35

Improved assumptions relating to isqrt

Improved various assumptions relating to values yielded by `isqrt`.
Does not solve but does improve rust-lang/rust#132763.

Re-openeing of rust-lang/rust#154115

Added assumptions are:
* if `x` is nonzero then `x.isqrt()`  is nonzero
* `x.isqrt() <= x`
* `x.isqrt() * x.isqrt() <= x`

library/core/src/num/imp/int_sqrt.rs

@@ -41,36 +41,6 @@ pub(in crate::num) const fn u8(n: u8) -> u8 {
     U8_ISQRT_WITH_REMAINDER[n as usize].0
 }
 
-/// Generates an `i*` function that returns the [integer square root](
-/// https://en.wikipedia.org/wiki/Integer_square_root) of any **nonnegative**
-/// input of a specific signed integer type.
-macro_rules! signed_fn {
-    ($SignedT:ident, $UnsignedT:ident) => {
-        /// Returns the [integer square root](
-        /// https://en.wikipedia.org/wiki/Integer_square_root) of any
-        /// **nonnegative**
-        #[doc = concat!("[`", stringify!($SignedT), "`](prim@", stringify!($SignedT), ")")]
-        /// input.
-        ///
-        /// # Safety
-        ///
-        /// This results in undefined behavior when the input is negative.
-        #[must_use = "this returns the result of the operation, \
-                      without modifying the original"]
-        #[inline]
-        pub(in crate::num) const unsafe fn $SignedT(n: $SignedT) -> $SignedT {
-            debug_assert!(n >= 0, "Negative input inside `isqrt`.");
-            $UnsignedT(n as $UnsignedT) as $SignedT
-        }
-    };
-}
-
-signed_fn!(i8, u8);
-signed_fn!(i16, u16);
-signed_fn!(i32, u32);
-signed_fn!(i64, u64);
-signed_fn!(i128, u128);
-
 /// Generates a `u*` function that returns the [integer square root](
 /// https://en.wikipedia.org/wiki/Integer_square_root) of any input of
 /// a specific unsigned integer type.

library/core/src/num/int_macros.rs

@@ -1933,10 +1933,9 @@ pub const fn checked_isqrt(self) -> Option<Self> {
             if self < 0 {
                 None
             } else {
-                // SAFETY: Input is nonnegative in this `else` branch.
-                let result = unsafe {
-                    imp::int_sqrt::$ActualT(self as $ActualT) as $SelfT
-                };
+                // The upper bound of `$UnsignedT::MAX.isqrt()` told to the compiler
+                // in the unsigned function also tells it that `result >= 0`
+                let result = self.cast_unsigned().isqrt().cast_signed();
 
                 // Inform the optimizer what the range of outputs is. If
                 // testing `core` crashes with no panic message and a
@@ -1950,15 +1949,9 @@ pub const fn checked_isqrt(self) -> Option<Self> {
                 // `[0, <$ActualT>::MAX]`, sqrt(n) will be bounded by
                 // `[sqrt(0), sqrt(<$ActualT>::MAX)]`.
                 unsafe {
-                    // SAFETY: `<$ActualT>::MAX` is nonnegative.
-                    const MAX_RESULT: $SelfT = unsafe {
-                        imp::int_sqrt::$ActualT(<$ActualT>::MAX) as $SelfT
-                    };
-
-                    crate::hint::assert_unchecked(result >= 0);
+                    const MAX_RESULT: $SelfT = <$SelfT>::MAX.cast_unsigned().isqrt().cast_signed();
                     crate::hint::assert_unchecked(result <= MAX_RESULT);
                 }
-
                 Some(result)
             }
         }

library/core/src/num/uint_macros.rs

@@ -3682,7 +3682,7 @@ pub const fn pow(self, mut exp: u32) -> Self {
                       without modifying the original"]
         #[inline]
         pub const fn isqrt(self) -> Self {
-            let result = imp::int_sqrt::$ActualT(self as $ActualT) as $SelfT;
+            let result = imp::int_sqrt::$ActualT(self as $ActualT) as Self;
 
             // Inform the optimizer what the range of outputs is. If testing
             // `core` crashes with no panic message and a `num::int_sqrt::u*`
@@ -3693,10 +3693,29 @@ pub const fn isqrt(self) -> Self {
             // function, which means that increasing the input will never
             // cause the output to decrease. Thus, since the input for unsigned
             // integers is bounded by `[0, <$ActualT>::MAX]`, sqrt(n) will be
-            // bounded by `[sqrt(0), sqrt(<$ActualT>::MAX)]`.
+            // bounded by `[sqrt(0), sqrt(<$ActualT>::MAX)]` and bounding the
+            // input by `[1, <$ActualT>::MAX]` bounds sqrt(n) by
+            // `[sqrt(1), sqrt(<$ActualT>::MAX)]`.
             unsafe {
                 const MAX_RESULT: $SelfT = imp::int_sqrt::$ActualT(<$ActualT>::MAX) as $SelfT;
-                crate::hint::assert_unchecked(result <= MAX_RESULT);
+                crate::hint::assert_unchecked(result <= MAX_RESULT)
+            }
+
+            if self >= 1 {
+                // SAFETY: The above statements about monotonicity also apply here.
+                // Since the input in this branch is bounded by `[1, <$ActualT>::MAX]`,
+                // sqrt(n) is bounded by `[sqrt(1), sqrt(<$ActualT>::MAX)]`, and
+                // `sqrt(1) == 1`.
+                unsafe { crate::hint::assert_unchecked(result >= 1) }
+            }
+
+            // SAFETY: the isqrt implementation returns the square root and rounds down,
+            // meaning `result * result <= self`. This implies `result <= self`.
+            // The compiler needs both to optimize for both.
+            // `result * result <= self` implies the multiplication will not overflow.
+            unsafe {
+                crate::hint::assert_unchecked(result.unchecked_mul(result) <= self);
+                crate::hint::assert_unchecked(result <= self);
             }
 
             result