mikros/rust
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Rollup merge of #126199 - ivan-shrimp:nonzero_isqrt, r=tgross35

Browse Source 
Add `isqrt` to `NonZero<uN>`

Implements [#70887 (comment)](https://github.com/rust-lang/rust/issues/116226#issuecomment-2144225174), with the following signature:

```rust
impl NonZero<uN> {
    const fn isqrt(self) -> Self;
}
```

Unintended benefits include one fewer panicking branch in `ilog2` for LLVM to optimize away, and one fewer `assume_unchecked` as `NonZero` already does that.

The fast path for `self == 1` is dropped, but the current implementation is very slow anyways compared to hardware. Performance improvements can always come later.

(I didn't add the function to `NonZero<iN>`, since _every_ existing `NonZero` method is non-panicking, and it might be nice to leave it that way.)
This commit is contained in:
Matthias Krüger 2024-07-19 17:06:49 +02:00 committed by GitHub
commit 355cdd3bc1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 63 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
library/core/src/num/mod.rs
View File

@ -3,7 +3,6 @@
#![stable(feature = "rust1", since = "1.0.0")]
use crate::ascii;
use crate::hint;
use crate::intrinsics;
use crate::mem;
use crate::str::FromStr;

60
library/core/src/num/nonzero.rs
View File

@ -3,6 +3,7 @@
use crate::cmp::Ordering;
use crate::fmt;
use crate::hash::{Hash, Hasher};
use crate::hint;
use crate::intrinsics;
use crate::marker::{Freeze, StructuralPartialEq};
use crate::ops::{BitOr, BitOrAssign, Div, DivAssign, Neg, Rem, RemAssign};
@ -604,7 +605,6 @@ macro_rules! nonzero_integer {
}
nonzero_integer_signedness_dependent_methods! {
Self = $Ty,
Primitive = $signedness $Int,
UnsignedPrimitive = $Uint,
}
@ -823,7 +823,7 @@ macro_rules! nonzero_integer {
}
}
nonzero_integer_signedness_dependent_impls!($Ty $signedness $Int);
nonzero_integer_signedness_dependent_impls!($signedness $Int);
};
(Self = $Ty:ident, Primitive = unsigned $Int:ident $(,)?) => {
@ -849,7 +849,7 @@ macro_rules! nonzero_integer {
macro_rules! nonzero_integer_signedness_dependent_impls {
// Impls for unsigned nonzero types only.
($Ty:ident unsigned $Int:ty) => {
(unsigned $Int:ty) => {
#[stable(feature = "nonzero_div", since = "1.51.0")]
impl Div<NonZero<$Int>> for $Int {
type Output = $Int;
@ -897,7 +897,7 @@ macro_rules! nonzero_integer_signedness_dependent_impls {
}
};
// Impls for signed nonzero types only.
($Ty:ident signed $Int:ty) => {
(signed $Int:ty) => {
#[stable(feature = "signed_nonzero_neg", since = "1.71.0")]
impl Neg for NonZero<$Int> {
type Output = Self;
@ -918,7 +918,6 @@ macro_rules! nonzero_integer_signedness_dependent_impls {
macro_rules! nonzero_integer_signedness_dependent_methods {
// Associated items for unsigned nonzero types only.
(
Self = $Ty:ident,
Primitive = unsigned $Int:ident,
UnsignedPrimitive = $Uint:ty,
) => {
@ -1224,11 +1223,60 @@ macro_rules! nonzero_integer_signedness_dependent_methods {
intrinsics::ctpop(self.get()) < 2
}
/// Returns the square root of the number, rounded down.
///
/// # Examples
///
/// Basic usage:
/// ```
/// #![feature(isqrt)]
/// # use std::num::NonZero;
/// #
/// # fn main() { test().unwrap(); }
/// # fn test() -> Option<()> {
#[doc = concat!("let ten = NonZero::new(10", stringify!($Int), ")?;")]
#[doc = concat!("let three = NonZero::new(3", stringify!($Int), ")?;")]
///
/// assert_eq!(ten.isqrt(), three);
/// # Some(())
/// # }
#[unstable(feature = "isqrt", issue = "116226")]
#[rustc_const_unstable(feature = "isqrt", issue = "116226")]
#[must_use = "this returns the result of the operation, \
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;
}
// 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) }
}
};
// Associated items for signed nonzero types only.
(
Self = $Ty:ident,
Primitive = signed $Int:ident,
UnsignedPrimitive = $Uint:ty,
) => {

47
library/core/src/num/uint_macros.rs
View File

@ -1226,10 +1226,9 @@ macro_rules! uint_impl {
without modifying the original"]
#[inline]
pub const fn checked_ilog2(self) -> Option<u32> {
if let Some(x) = NonZero::new(self) {
Some(x.ilog2())
} else {
None
match NonZero::new(self) {
Some(x) => Some(x.ilog2()),
None => None,
}
}
@ -1248,10 +1247,9 @@ macro_rules! uint_impl {
without modifying the original"]
#[inline]
pub const fn checked_ilog10(self) -> Option<u32> {
if let Some(x) = NonZero::new(self) {
Some(x.ilog10())
} else {
None
match NonZero::new(self) {
Some(x) => Some(x.ilog10()),
None => None,
}
}
@ -2590,37 +2588,10 @@ macro_rules! uint_impl {
without modifying the original"]
#[inline]
pub const fn isqrt(self) -> Self {
if self < 2 {
return self;
match NonZero::new(self) {
Some(x) => x.isqrt().get(),
None => 0,
}
// 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;
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;
}
// SAFETY: the result is positive and fits in an integer with half as many bits.
// Inform the optimizer about it.
unsafe {
hint::assert_unchecked(0 < res);
hint::assert_unchecked(res < 1 << (Self::BITS / 2));
}
res
}
/// Performs Euclidean division.