From eaaed00ff53bc121dacdc3fc229b793f83d72e78 Mon Sep 17 00:00:00 2001 From: ivan-shrimp <70307174+ivan-shrimp@users.noreply.github.com> Date: Fri, 19 Jul 2024 18:10:41 +0800 Subject: [PATCH 1/2] add `NonZero::isqrt` --- library/core/src/num/mod.rs | 1 - library/core/src/num/nonzero.rs | 60 ++++++++++++++++++++++++++--- library/core/src/num/uint_macros.rs | 47 +++++----------------- 3 files changed, 63 insertions(+), 45 deletions(-) diff --git a/library/core/src/num/mod.rs b/library/core/src/num/mod.rs index 034af6a0d5731..4e8e0ecdde998 100644 --- a/library/core/src/num/mod.rs +++ b/library/core/src/num/mod.rs @@ -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; diff --git a/library/core/src/num/nonzero.rs b/library/core/src/num/nonzero.rs index 0c6f06dc017e7..2a5821c708f12 100644 --- a/library/core/src/num/nonzero.rs +++ b/library/core/src/num/nonzero.rs @@ -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> 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 + // + // which cites as source the following C code: + // . + + 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 result is positive. + unsafe { Self::new_unchecked(res) } + } }; // Associated items for signed nonzero types only. ( - Self = $Ty:ident, Primitive = signed $Int:ident, UnsignedPrimitive = $Uint:ty, ) => { diff --git a/library/core/src/num/uint_macros.rs b/library/core/src/num/uint_macros.rs index ad72c29758bd7..d50bcde01571c 100644 --- a/library/core/src/num/uint_macros.rs +++ b/library/core/src/num/uint_macros.rs @@ -1226,10 +1226,9 @@ macro_rules! uint_impl { without modifying the original"] #[inline] pub const fn checked_ilog2(self) -> Option { - 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 { - 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; - } - - // The algorithm is based on the one presented in - // - // which cites as source the following C code: - // . - - 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; + match NonZero::new(self) { + Some(x) => x.isqrt().get(), + None => 0, } - - // 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. From 90bba8beb5cba6cbd0ff0fc92ed0a8b8ef537f8c Mon Sep 17 00:00:00 2001 From: ivan-shrimp <70307174+ivan-shrimp@users.noreply.github.com> Date: Fri, 19 Jul 2024 19:16:33 +0800 Subject: [PATCH 2/2] improve safety comment --- library/core/src/num/nonzero.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/core/src/num/nonzero.rs b/library/core/src/num/nonzero.rs index 2a5821c708f12..d80d3241b1eee 100644 --- a/library/core/src/num/nonzero.rs +++ b/library/core/src/num/nonzero.rs @@ -1270,7 +1270,7 @@ macro_rules! nonzero_integer_signedness_dependent_methods { // Inform the optimizer about it. unsafe { hint::assert_unchecked(res < 1 << (Self::BITS / 2)) }; - // SAFETY: The result is positive. + // SAFETY: The square root of an integer >= 1 is always >= 1. unsafe { Self::new_unchecked(res) } } };