Add "integer square root" method to integer primitive types

library/core/src/num/int_macros.rs

@@ -898,6 +898,29 @@ macro_rules! int_impl {
             acc.checked_mul(base)
         }
 
+        /// Returns the square root of the number, rounded down.
+        ///
+        /// Returns `None` if `self` is negative.
+        ///
+        /// # Examples
+        ///
+        /// Basic usage:
+        /// ```
+        #[doc = concat!("assert_eq!(10", stringify!($SelfT), ".checked_isqrt(), Some(3));")]
+        /// ```
+        #[stable(feature = "isqrt", since = "1.73.0")]
+        #[rustc_const_stable(feature = "isqrt", since = "1.73.0")]
+        #[must_use = "this returns the result of the operation, \
+                      without modifying the original"]
+        #[inline]
+        pub const fn checked_isqrt(self) -> Option<Self> {
+            if self < 0 {
+                None
+            } else {
+                Some((self as $UnsignedT).isqrt() as Self)
+            }
+        }
+
         /// Saturating integer addition. Computes `self + rhs`, saturating at the numeric
         /// bounds instead of overflowing.
         ///
@@ -2061,6 +2084,35 @@ macro_rules! int_impl {
             acc * base
         }
 
+        /// Returns the square root of the number, rounded down.
+        ///
+        /// # Panics
+        ///
+        /// This function will panic if `self` is negative.
+        ///
+        /// # Examples
+        ///
+        /// Basic usage:
+        /// ```
+        #[doc = concat!("assert_eq!(10", stringify!($SelfT), ".isqrt(), 3);")]
+        /// ```
+        #[stable(feature = "isqrt", since = "1.73.0")]
+        #[rustc_const_stable(feature = "isqrt", since = "1.73.0")]
+        #[must_use = "this returns the result of the operation, \
+                      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"),
+            }
+        }
+
         /// Calculates the quotient of Euclidean division of `self` by `rhs`.
         ///
         /// This computes the integer `q` such that `self = q * rhs + r`, with

library/core/src/num/uint_macros.rs

@@ -1979,6 +1979,41 @@ macro_rules! uint_impl {
             acc * base
         }
 
+        /// Returns the square root of the number, rounded down.
+        ///
+        /// # Examples
+        ///
+        /// Basic usage:
+        /// ```
+        #[doc = concat!("assert_eq!(10", stringify!($SelfT), ".isqrt(), 3);")]
+        /// ```
+        #[stable(feature = "isqrt", since = "1.73.0")]
+        #[rustc_const_stable(feature = "isqrt", since = "1.73.0")]
+        #[must_use = "this returns the result of the operation, \
+                      without modifying the original"]
+        #[inline]
+        pub const fn isqrt(self) -> Self {
+            if self < 2 {
+                return self;
+            }
+
+            let mut x: Self = self;
+            let mut c: Self = 0;
+            let mut d: Self = 1 << (self.ilog2() & !1);
+
+            while (d != 0) {
+                if x >= c + d {
+                    x -= c + d;
+                    c = (c >> 1) + d;
+                } else {
+                    c >>= 1;
+                }
+                d >>= 2;
+            }
+
+            c
+        }
+
         /// Performs Euclidean division.
         ///
         /// Since, for the positive integers, all common

library/core/tests/num/int_macros.rs

@@ -290,6 +290,34 @@ macro_rules! int_module {
                 assert_eq!(r.saturating_pow(0), 1 as $T);
             }
 
+            #[test]
+            fn test_isqrt() {
+                assert_eq!($T::MIN.checked_isqrt(), None);
+                assert_eq!((-1 as $T).checked_isqrt(), None);
+                assert_eq!((0 as $T).isqrt(), 0 as $T);
+                assert_eq!((1 as $T).isqrt(), 1 as $T);
+                assert_eq!((2 as $T).isqrt(), 1 as $T);
+                assert_eq!((99 as $T).isqrt(), 9 as $T);
+                assert_eq!((100 as $T).isqrt(), 10 as $T);
+
+                let n_max: $T = (1024 * 1024).min($T::MAX as u128) as $T;
+                for n in 0..=n_max {
+                    let isqrt: $T = n.isqrt();
+
+                    assert!(isqrt.pow(2) <= n);
+                    let (square, overflow) = (isqrt + 1).overflowing_pow(2);
+                    assert!(overflow || square > n);
+                }
+
+                for n in ($T::MAX - 127)..=$T::MAX {
+                    let isqrt: $T = n.isqrt();
+
+                    assert!(isqrt.pow(2) <= n);
+                    let (square, overflow) = (isqrt + 1).overflowing_pow(2);
+                    assert!(overflow || square > n);
+                }
+            }
+
             #[test]
             fn test_div_floor() {
                 let a: $T = 8;

library/core/tests/num/uint_macros.rs

@@ -206,6 +206,31 @@ macro_rules! uint_module {
                 assert_eq!(r.saturating_pow(2), MAX);
             }
 
+            #[test]
+            fn test_isqrt() {
+                assert_eq!((0 as $T).isqrt(), 0 as $T);
+                assert_eq!((1 as $T).isqrt(), 1 as $T);
+                assert_eq!((2 as $T).isqrt(), 1 as $T);
+                assert_eq!((99 as $T).isqrt(), 9 as $T);
+                assert_eq!((100 as $T).isqrt(), 10 as $T);
+                assert_eq!($T::MAX.isqrt(), (1 << ($T::BITS / 2)) - 1);
+
+                let n_max: $T = (1024 * 1024).min($T::MAX as u128) as $T;
+                for n in 0..=n_max {
+                    let isqrt: $T = n.isqrt();
+
+                    assert!(isqrt.pow(2) <= n);
+                    assert!(isqrt + 1 == (1 as $T) << ($T::BITS / 2) || (isqrt + 1).pow(2) > n);
+                }
+
+                for n in ($T::MAX - 255)..=$T::MAX {
+                    let isqrt: $T = n.isqrt();
+
+                    assert!(isqrt.pow(2) <= n);
+                    assert!(isqrt + 1 == (1 as $T) << ($T::BITS / 2) || (isqrt + 1).pow(2) > n);
+                }
+            }
+
             #[test]
             fn test_div_floor() {
                 assert_eq!((8 as $T).div_floor(3), 2);