Skip to content

Commit

Permalink
Greatly sped up checked_isqrt and isqrt methods
Browse files Browse the repository at this point in the history
* 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.
  • Loading branch information
ChaiTRex committed Jul 25, 2024
1 parent d180572 commit eb17af7
Show file tree
Hide file tree
Showing 4 changed files with 47 additions and 35 deletions.
36 changes: 27 additions & 9 deletions library/core/src/num/int_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}

Expand Down Expand Up @@ -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
}
}

Expand Down
1 change: 1 addition & 0 deletions library/core/src/num/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
34 changes: 11 additions & 23 deletions library/core/src/num/nonzero.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
// <https://en.wikipedia.org/wiki/Methods_of_computing_square_roots#Binary_numeral_system_(base_2)>
// which cites as source the following C code:
// <https://web.archive.org/web/20120306040058/http://medialab.freaknet.org/martin/src/sqrt/sqrt.c>.

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) }
}
};

Expand Down
11 changes: 8 additions & 3 deletions library/core/src/num/uint_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit eb17af7

Please sign in to comment.