diff --git a/library/core/src/num/imp/int_sqrt.rs b/library/core/src/num/imp/int_sqrt.rs index 152841793d56..c2bb4bf6078f 100644 --- a/library/core/src/num/imp/int_sqrt.rs +++ b/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. diff --git a/library/core/src/num/int_macros.rs b/library/core/src/num/int_macros.rs index aba0eaf8d484..a7fc7fcc60a5 100644 --- a/library/core/src/num/int_macros.rs +++ b/library/core/src/num/int_macros.rs @@ -1845,10 +1845,9 @@ pub const fn checked_isqrt(self) -> Option { 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 @@ -1862,15 +1861,9 @@ pub const fn checked_isqrt(self) -> Option { // `[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) } } diff --git a/library/core/src/num/uint_macros.rs b/library/core/src/num/uint_macros.rs index 27dbe6d3f1d8..2ed81c8862e1 100644 --- a/library/core/src/num/uint_macros.rs +++ b/library/core/src/num/uint_macros.rs @@ -3595,7 +3595,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*` @@ -3606,10 +3606,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