Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add "integer square root" method to integer primitive types #116176

Merged
merged 10 commits into from
Sep 29, 2023
52 changes: 52 additions & 0 deletions library/core/src/num/int_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
FedericoStra marked this conversation as resolved.
Show resolved Hide resolved
#[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.
///
Expand Down Expand Up @@ -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")]
FedericoStra marked this conversation as resolved.
Show resolved Hide resolved
#[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
Expand Down
35 changes: 35 additions & 0 deletions library/core/src/num/uint_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
FedericoStra marked this conversation as resolved.
Show resolved Hide resolved
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[inline]
pub const fn isqrt(self) -> Self {
dtolnay marked this conversation as resolved.
Show resolved Hide resolved
if self < 2 {
FedericoStra marked this conversation as resolved.
Show resolved Hide resolved
return self;
}

let mut x: Self = self;
let mut c: Self = 0;
let mut d: Self = 1 << (self.ilog2() & !1);
FedericoStra marked this conversation as resolved.
Show resolved Hide resolved

while (d != 0) {
FedericoStra marked this conversation as resolved.
Show resolved Hide resolved
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
Expand Down
28 changes: 28 additions & 0 deletions library/core/tests/num/int_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
dtolnay marked this conversation as resolved.
Show resolved Hide resolved
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;
Expand Down
25 changes: 25 additions & 0 deletions library/core/tests/num/uint_macros.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
Loading