Skip to content

Commit eb17af7

Browse files
committed
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.
1 parent d180572 commit eb17af7

File tree

4 files changed

+47
-35
lines changed

4 files changed

+47
-35
lines changed

library/core/src/num/int_macros.rs

+27-9
Original file line numberDiff line numberDiff line change
@@ -1580,7 +1580,18 @@ macro_rules! int_impl {
15801580
if self < 0 {
15811581
None
15821582
} else {
1583-
Some((self as $UnsignedT).isqrt() as Self)
1583+
let result = crate::num::int_sqrt::$ActualT(self as $ActualT) as $SelfT;
1584+
1585+
// SAFETY: Inform the optimizer that square roots of
1586+
// nonnegative integers are nonnegative and what the maximum
1587+
// result is.
1588+
unsafe {
1589+
crate::hint::assert_unchecked(result >= 0);
1590+
const MAX_RESULT: $SelfT = crate::num::int_sqrt::$ActualT($ActualT::MAX) as $SelfT;
1591+
crate::hint::assert_unchecked(result <= MAX_RESULT);
1592+
}
1593+
1594+
Some(result)
15841595
}
15851596
}
15861597

@@ -2766,14 +2777,21 @@ macro_rules! int_impl {
27662777
without modifying the original"]
27672778
#[inline]
27682779
pub const fn isqrt(self) -> Self {
2769-
// I would like to implement it as
2770-
// ```
2771-
// self.checked_isqrt().expect("argument of integer square root must be non-negative")
2772-
// ```
2773-
// but `expect` is not yet stable as a `const fn`.
2774-
match self.checked_isqrt() {
2775-
Some(sqrt) => sqrt,
2776-
None => panic!("argument of integer square root must be non-negative"),
2780+
if self < 0 {
2781+
crate::num::int_sqrt::panic_for_negative_argument();
2782+
} else {
2783+
let result = crate::num::int_sqrt::$ActualT(self as $ActualT) as $SelfT;
2784+
2785+
// SAFETY: Inform the optimizer that square roots of
2786+
// nonnegative integers are nonnegative and what the maximum
2787+
// result is.
2788+
unsafe {
2789+
crate::hint::assert_unchecked(result >= 0);
2790+
const MAX_RESULT: $SelfT = crate::num::int_sqrt::$ActualT($ActualT::MAX) as $SelfT;
2791+
crate::hint::assert_unchecked(result <= MAX_RESULT);
2792+
}
2793+
2794+
result
27772795
}
27782796
}
27792797

library/core/src/num/mod.rs

+1
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ mod uint_macros; // import uint_impl!
4343

4444
mod error;
4545
mod int_log10;
46+
mod int_sqrt;
4647
mod nonzero;
4748
mod overflow_panic;
4849
mod saturating;

library/core/src/num/nonzero.rs

+11-23
Original file line numberDiff line numberDiff line change
@@ -1247,31 +1247,19 @@ macro_rules! nonzero_integer_signedness_dependent_methods {
12471247
without modifying the original"]
12481248
#[inline]
12491249
pub const fn isqrt(self) -> Self {
1250-
// The algorithm is based on the one presented in
1251-
// <https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Binary_numeral_system_(base_2)>
1252-
// which cites as source the following C code:
1253-
// <https://web.archive.org/web/20120306040058/http://medialab.freaknet.org/martin/src/sqrt/sqrt.c>.
1254-
1255-
let mut op = self.get();
1256-
let mut res = 0;
1257-
let mut one = 1 << (self.ilog2() & !1);
1258-
1259-
while one != 0 {
1260-
if op >= res + one {
1261-
op -= res + one;
1262-
res = (res >> 1) + one;
1263-
} else {
1264-
res >>= 1;
1265-
}
1266-
one >>= 2;
1250+
let result = super::int_sqrt::$Int(self.get());
1251+
1252+
// SAFETY: Inform the optimizer that square roots of positive
1253+
// integers are positive and what the maximum result is.
1254+
unsafe {
1255+
hint::assert_unchecked(result > 0);
1256+
const MAX_RESULT: $Int = super::int_sqrt::$Int($Int::MAX);
1257+
hint::assert_unchecked(result <= MAX_RESULT);
12671258
}
12681259

1269-
// SAFETY: The result fits in an integer with half as many bits.
1270-
// Inform the optimizer about it.
1271-
unsafe { hint::assert_unchecked(res < 1 << (Self::BITS / 2)) };
1272-
1273-
// SAFETY: The square root of an integer >= 1 is always >= 1.
1274-
unsafe { Self::new_unchecked(res) }
1260+
// SAFETY: The square root of a positive integer is always
1261+
// positive.
1262+
unsafe { Self::new_unchecked(result) }
12751263
}
12761264
};
12771265

library/core/src/num/uint_macros.rs

+8-3
Original file line numberDiff line numberDiff line change
@@ -2588,10 +2588,15 @@ macro_rules! uint_impl {
25882588
without modifying the original"]
25892589
#[inline]
25902590
pub const fn isqrt(self) -> Self {
2591-
match NonZero::new(self) {
2592-
Some(x) => x.isqrt().get(),
2593-
None => 0,
2591+
let result = crate::num::int_sqrt::$ActualT(self as $ActualT) as $SelfT;
2592+
2593+
// SAFETY: Inform the optimizer of what the maximum result is.
2594+
unsafe {
2595+
const MAX_RESULT: $SelfT = crate::num::int_sqrt::$ActualT($ActualT::MAX) as $SelfT;
2596+
crate::hint::assert_unchecked(result <= MAX_RESULT);
25942597
}
2598+
2599+
result
25952600
}
25962601

25972602
/// Performs Euclidean division.

0 commit comments

Comments
 (0)