Rollup merge of rust-lang#128166 - ChaiTRex:isqrt, r=tgross35

Improved `checked_isqrt` and `isqrt` methods

### Improved tests of `isqrt` and `checked_isqrt` implementations

* Inputs chosen more thoroughly and systematically.
* Checks that `isqrt` and `checked_isqrt` have equivalent results for signed types, either equivalent numerically or equivalent as a panic and a `None`.
* Checks that `isqrt` has numerically-equivalent results for unsigned types and their `NonZero` counterparts.

### Added benchmarks for `isqrt` implementations

### Greatly sped up `checked_isqrt` and `isqrt` methods

* Uses a lookup table for 8-bit integers and then the Karatsuba square root algorithm for larger integers.
* Includes optimization hints that give the compiler the exact numeric range of results.

### Feature tracking issue

`isqrt` is an unstable feature tracked at rust-lang#116226.

<details><summary>Benchmarked improvements</summary>

### Command used to benchmark

    ./x bench library/core -- int_sqrt

### Before

    benchmarks:
        num::int_sqrt::i128::isqrt           439591.65/iter  +/- 6652.70
        num::int_sqrt::i16::isqrt              5302.97/iter   +/- 160.93
        num::int_sqrt::i32::isqrt             62999.11/iter  +/- 2022.05
        num::int_sqrt::i64::isqrt            125248.81/iter  +/- 1674.43
        num::int_sqrt::i8::isqrt                123.56/iter     +/- 1.87
        num::int_sqrt::isize::isqrt          125356.56/iter  +/- 1017.03
        num::int_sqrt::non_zero_u128::isqrt  437443.75/iter  +/- 3535.43
        num::int_sqrt::non_zero_u16::isqrt     8604.58/iter    +/- 94.76
        num::int_sqrt::non_zero_u32::isqrt    62933.33/iter   +/- 517.30
        num::int_sqrt::non_zero_u64::isqrt   125076.38/iter +/- 11340.61
        num::int_sqrt::non_zero_u8::isqrt       221.51/iter     +/- 1.58
        num::int_sqrt::non_zero_usize::isqrt 136005.21/iter  +/- 2020.35
        num::int_sqrt::u128::isqrt           439014.55/iter  +/- 3920.45
        num::int_sqrt::u16::isqrt              8575.08/iter   +/- 148.06
        num::int_sqrt::u32::isqrt             63008.89/iter   +/- 803.67
        num::int_sqrt::u64::isqrt            125088.09/iter   +/- 879.29
        num::int_sqrt::u8::isqrt                230.18/iter     +/- 2.04
        num::int_sqrt::usize::isqrt          125237.51/iter  +/- 4747.83
### After

    benchmarks:
        num::int_sqrt::i128::isqrt           105184.89/iter +/- 1171.38
        num::int_sqrt::i16::isqrt              1910.26/iter   +/- 78.50
        num::int_sqrt::i32::isqrt             34260.34/iter  +/- 960.84
        num::int_sqrt::i64::isqrt             45939.19/iter +/- 2525.65
        num::int_sqrt::i8::isqrt                 22.87/iter    +/- 0.45
        num::int_sqrt::isize::isqrt           45884.17/iter  +/- 595.49
        num::int_sqrt::non_zero_u128::isqrt  106344.27/iter  +/- 780.99
        num::int_sqrt::non_zero_u16::isqrt     2790.19/iter   +/- 53.43
        num::int_sqrt::non_zero_u32::isqrt    33613.99/iter  +/- 362.96
        num::int_sqrt::non_zero_u64::isqrt    46235.42/iter  +/- 429.69
        num::int_sqrt::non_zero_u8::isqrt        31.78/iter    +/- 0.75
        num::int_sqrt::non_zero_usize::isqrt  46208.75/iter  +/- 375.27
        num::int_sqrt::u128::isqrt           106385.94/iter +/- 1649.95
        num::int_sqrt::u16::isqrt              2747.69/iter   +/- 28.72
        num::int_sqrt::u32::isqrt             33627.09/iter  +/- 475.68
        num::int_sqrt::u64::isqrt             46182.29/iter  +/- 311.16
        num::int_sqrt::u8::isqrt                 33.10/iter    +/- 0.30
        num::int_sqrt::usize::isqrt           46165.00/iter  +/- 388.41

</details>

Author: matthiaskrgr

library/core/benches/lib.rs

@@ -8,6 +8,7 @@
 #![feature(iter_array_chunks)]
 #![feature(iter_next_chunk)]
 #![feature(iter_advance_by)]
+#![feature(isqrt)]
 
 extern crate test;
 

library/core/benches/num/int_sqrt/mod.rs

@@ -0,0 +1,62 @@
+use rand::Rng;
+use test::{black_box, Bencher};
+
+macro_rules! int_sqrt_bench {
+    ($t:ty, $predictable:ident, $random:ident, $random_small:ident, $random_uniform:ident) => {
+        #[bench]
+        fn $predictable(bench: &mut Bencher) {
+            bench.iter(|| {
+                for n in 0..(<$t>::BITS / 8) {
+                    for i in 1..=(100 as $t) {
+                        let x = black_box(i << (n * 8));
+                        black_box(x.isqrt());
+                    }
+                }
+            });
+        }
+
+        #[bench]
+        fn $random(bench: &mut Bencher) {
+            let mut rng = crate::bench_rng();
+            /* Exponentially distributed random numbers from the whole range of the type.  */
+            let numbers: Vec<$t> =
+                (0..256).map(|_| rng.gen::<$t>() >> rng.gen_range(0..<$t>::BITS)).collect();
+            bench.iter(|| {
+                for x in &numbers {
+                    black_box(black_box(x).isqrt());
+                }
+            });
+        }
+
+        #[bench]
+        fn $random_small(bench: &mut Bencher) {
+            let mut rng = crate::bench_rng();
+            /* Exponentially distributed random numbers from the range 0..256.  */
+            let numbers: Vec<$t> =
+                (0..256).map(|_| (rng.gen::<u8>() >> rng.gen_range(0..u8::BITS)) as $t).collect();
+            bench.iter(|| {
+                for x in &numbers {
+                    black_box(black_box(x).isqrt());
+                }
+            });
+        }
+
+        #[bench]
+        fn $random_uniform(bench: &mut Bencher) {
+            let mut rng = crate::bench_rng();
+            /* Exponentially distributed random numbers from the whole range of the type.  */
+            let numbers: Vec<$t> = (0..256).map(|_| rng.gen::<$t>()).collect();
+            bench.iter(|| {
+                for x in &numbers {
+                    black_box(black_box(x).isqrt());
+                }
+            });
+        }
+    };
+}
+
+int_sqrt_bench! {u8, u8_sqrt_predictable, u8_sqrt_random, u8_sqrt_random_small, u8_sqrt_uniform}
+int_sqrt_bench! {u16, u16_sqrt_predictable, u16_sqrt_random, u16_sqrt_random_small, u16_sqrt_uniform}
+int_sqrt_bench! {u32, u32_sqrt_predictable, u32_sqrt_random, u32_sqrt_random_small, u32_sqrt_uniform}
+int_sqrt_bench! {u64, u64_sqrt_predictable, u64_sqrt_random, u64_sqrt_random_small, u64_sqrt_uniform}
+int_sqrt_bench! {u128, u128_sqrt_predictable, u128_sqrt_random, u128_sqrt_random_small, u128_sqrt_uniform}

library/core/benches/num/mod.rs

@@ -2,6 +2,7 @@ mod dec2flt;
 mod flt2dec;
 mod int_log;
 mod int_pow;
+mod int_sqrt;
 
 use std::str::FromStr;
 

library/core/src/num/int_macros.rs

@@ -1579,7 +1579,33 @@ macro_rules! int_impl {
             if self < 0 {
                 None
             } else {
-                Some((self as $UnsignedT).isqrt() as Self)
+                // SAFETY: Input is nonnegative in this `else` branch.
+                let result = unsafe {
+                    crate::num::int_sqrt::$ActualT(self as $ActualT) as $SelfT
+                };
+
+                // Inform the optimizer what the range of outputs is. If
+                // testing `core` crashes with no panic message and a
+                // `num::int_sqrt::i*` test failed, it's because your edits
+                // caused these assertions to become false.
+                //
+                // SAFETY: Integer square root is a monotonically nondecreasing
+                // function, which means that increasing the input will never
+                // cause the output to decrease. Thus, since the input for
+                // nonnegative signed integers is bounded by
+                // `[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 {
+                        crate::num::int_sqrt::$ActualT(<$ActualT>::MAX) as $SelfT
+                    };
+
+                    crate::hint::assert_unchecked(result >= 0);
+                    crate::hint::assert_unchecked(result <= MAX_RESULT);
+                }
+
+                Some(result)
             }
         }
 
@@ -2800,15 +2826,11 @@ macro_rules! int_impl {
         #[must_use = "this returns the result of the operation, \
                       without modifying the original"]
         #[inline]
+        #[track_caller]
         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"),
+                None => crate::num::int_sqrt::panic_for_negative_argument(),
             }
         }