diff --git a/library/core/src/num/int_macros.rs b/library/core/src/num/int_macros.rs index d40e02352a1d0..1b3dcf2287fc8 100644 --- a/library/core/src/num/int_macros.rs +++ b/library/core/src/num/int_macros.rs @@ -1580,7 +1580,18 @@ macro_rules! int_impl { if self < 0 { None } else { - Some((self as $UnsignedT).isqrt() as Self) + let result = crate::num::int_sqrt::$ActualT(self as $ActualT) as $SelfT; + + // SAFETY: Inform the optimizer that square roots of + // nonnegative integers are nonnegative and what the maximum + // result is. + unsafe { + crate::hint::assert_unchecked(result >= 0); + const MAX_RESULT: $SelfT = crate::num::int_sqrt::$ActualT($ActualT::MAX) as $SelfT; + crate::hint::assert_unchecked(result <= MAX_RESULT); + } + + Some(result) } } @@ -2766,14 +2777,21 @@ macro_rules! int_impl { without modifying the original"] #[inline] pub const fn isqrt(self) -> Self { - // I would like to implement it as - // ``` - // self.checked_isqrt().expect("argument of integer square root must be non-negative") - // ``` - // but `expect` is not yet stable as a `const fn`. - match self.checked_isqrt() { - Some(sqrt) => sqrt, - None => panic!("argument of integer square root must be non-negative"), + if self < 0 { + crate::num::int_sqrt::panic_for_negative_argument(); + } else { + let result = crate::num::int_sqrt::$ActualT(self as $ActualT) as $SelfT; + + // SAFETY: Inform the optimizer that square roots of + // nonnegative integers are nonnegative and what the maximum + // result is. + unsafe { + crate::hint::assert_unchecked(result >= 0); + const MAX_RESULT: $SelfT = crate::num::int_sqrt::$ActualT($ActualT::MAX) as $SelfT; + crate::hint::assert_unchecked(result <= MAX_RESULT); + } + + result } } diff --git a/library/core/src/num/mod.rs b/library/core/src/num/mod.rs index 4e8e0ecdde998..b059417357b80 100644 --- a/library/core/src/num/mod.rs +++ b/library/core/src/num/mod.rs @@ -43,6 +43,7 @@ mod uint_macros; // import uint_impl! mod error; mod int_log10; +mod int_sqrt; mod nonzero; mod overflow_panic; mod saturating; diff --git a/library/core/src/num/nonzero.rs b/library/core/src/num/nonzero.rs index d80d3241b1eee..6a0323136dc0e 100644 --- a/library/core/src/num/nonzero.rs +++ b/library/core/src/num/nonzero.rs @@ -1247,31 +1247,19 @@ macro_rules! nonzero_integer_signedness_dependent_methods { without modifying the original"] #[inline] pub const fn isqrt(self) -> Self { - // The algorithm is based on the one presented in - // - // which cites as source the following C code: - // . - - let mut op = self.get(); - let mut res = 0; - let mut one = 1 << (self.ilog2() & !1); - - while one != 0 { - if op >= res + one { - op -= res + one; - res = (res >> 1) + one; - } else { - res >>= 1; - } - one >>= 2; + let result = super::int_sqrt::$Int(self.get()); + + // SAFETY: Inform the optimizer that square roots of positive + // integers are positive and what the maximum result is. + unsafe { + hint::assert_unchecked(result > 0); + const MAX_RESULT: $Int = super::int_sqrt::$Int($Int::MAX); + hint::assert_unchecked(result <= MAX_RESULT); } - // SAFETY: The result fits in an integer with half as many bits. - // Inform the optimizer about it. - unsafe { hint::assert_unchecked(res < 1 << (Self::BITS / 2)) }; - - // SAFETY: The square root of an integer >= 1 is always >= 1. - unsafe { Self::new_unchecked(res) } + // SAFETY: The square root of a positive integer is always + // positive. + unsafe { Self::new_unchecked(result) } } }; diff --git a/library/core/src/num/uint_macros.rs b/library/core/src/num/uint_macros.rs index d50bcde01571c..756bac4813ab0 100644 --- a/library/core/src/num/uint_macros.rs +++ b/library/core/src/num/uint_macros.rs @@ -2588,10 +2588,15 @@ macro_rules! uint_impl { without modifying the original"] #[inline] pub const fn isqrt(self) -> Self { - match NonZero::new(self) { - Some(x) => x.isqrt().get(), - None => 0, + let result = crate::num::int_sqrt::$ActualT(self as $ActualT) as $SelfT; + + // SAFETY: Inform the optimizer of what the maximum result is. + unsafe { + const MAX_RESULT: $SelfT = crate::num::int_sqrt::$ActualT($ActualT::MAX) as $SelfT; + crate::hint::assert_unchecked(result <= MAX_RESULT); } + + result } /// Performs Euclidean division.