Skip to content

Commit

Permalink
Expanded power of two opt
Browse files Browse the repository at this point in the history
  • Loading branch information
@NCGThompson
NCGThompson committed Jan 11, 2024
1 parent 3c0b2d3 commit cdc2f08
Show file tree
Hide file tree
Showing 2 changed files with 229 additions and 121 deletions.
191 changes: 129 additions & 62 deletions library/core/src/num/int_macros.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -901,26 +901,46 @@ macro_rules! int_impl {
#[rustc_const_stable(feature = "const_int_pow", since = "1.50.0")]
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[rustc_allow_const_fn_unstable(is_val_statically_known)]
#[inline]
pub const fn checked_pow(self, mut exp: u32) -> Option<Self> {
if exp == 0 {
return Some(1);
}
let mut base = self;
let mut acc: Self = 1;
// SAFETY: This path has the same behavior as the other.
if unsafe { intrinsics::is_val_statically_known(self) }
&& self.unsigned_abs().is_power_of_two()
{
// SAFETY: We just checked this is a power of two. and above zero.
let power_used = unsafe { intrinsics::cttz_nonzero(self.wrapping_abs()) as u32 };
let num_shl = power_used.saturating_mul(exp);
let res = try_opt!((1 as Self).checked_shl(num_shl));

let sign = self.is_negative() && exp & 1 != 0;
if !sign && res == Self::MIN {
None
} else if sign {
Some(res.wrapping_neg())
} else {
Some(res)
}
} else {
if exp == 0 {
return Some(1);
}
let mut base = self;
let mut acc: Self = 1;

while exp > 1 {
if (exp & 1) == 1 {
acc = try_opt!(acc.checked_mul(base));
while exp > 1 {
if (exp & 1) == 1 {
acc = try_opt!(acc.checked_mul(base));
}
exp /= 2;
base = try_opt!(base.checked_mul(base));
}
exp /= 2;
base = try_opt!(base.checked_mul(base));
// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc.checked_mul(base)
}
// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc.checked_mul(base)
}

/// Returns the square root of the number, rounded down.
Expand Down Expand Up @@ -1537,27 +1557,48 @@ macro_rules! int_impl {
#[rustc_const_stable(feature = "const_int_pow", since = "1.50.0")]
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[rustc_allow_const_fn_unstable(is_val_statically_known)]
#[inline]
pub const fn wrapping_pow(self, mut exp: u32) -> Self {
if exp == 0 {
return 1;
}
let mut base = self;
let mut acc: Self = 1;
// SAFETY: This path has the same behavior as the other.
if unsafe { intrinsics::is_val_statically_known(self) }
&& self.unsigned_abs().is_power_of_two()
{
// SAFETY: We just checked this is a power of two. and above zero.
let power_used = unsafe { intrinsics::cttz_nonzero(self.wrapping_abs()) as u32 };
let num_shl = power_used.saturating_mul(exp);
let res = match (1 as Self).checked_shl(num_shl) {
Some(x) => x,
None => 0
};

while exp > 1 {
if (exp & 1) == 1 {
acc = acc.wrapping_mul(base);
let sign = self.is_negative() && exp & 1 != 0;
if sign {
res.wrapping_neg()
} else {
res
}
exp /= 2;
base = base.wrapping_mul(base);
}
} else {
if exp == 0 {
return 1;
}
let mut base = self;
let mut acc: Self = 1;

// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc.wrapping_mul(base)
while exp > 1 {
if (exp & 1) == 1 {
acc = acc.wrapping_mul(base);
}
exp /= 2;
base = base.wrapping_mul(base);
}

// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
acc.wrapping_mul(base)
}
}

/// Calculates `self` + `rhs`
Expand Down Expand Up @@ -2039,36 +2080,58 @@ macro_rules! int_impl {
#[rustc_const_stable(feature = "const_int_pow", since = "1.50.0")]
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[rustc_allow_const_fn_unstable(is_val_statically_known)]
#[inline]
pub const fn overflowing_pow(self, mut exp: u32) -> (Self, bool) {
if exp == 0 {
return (1,false);
}
let mut base = self;
let mut acc: Self = 1;
let mut overflown = false;
// Scratch space for storing results of overflowing_mul.
let mut r;

while exp > 1 {
if (exp & 1) == 1 {
r = acc.overflowing_mul(base);
acc = r.0;
// SAFETY: This path has the same behavior as the other.
if unsafe { intrinsics::is_val_statically_known(self) }
&& self.unsigned_abs().is_power_of_two()
{
// SAFETY: We just checked this is a power of two. and above zero.
let power_used = unsafe { intrinsics::cttz_nonzero(self.wrapping_abs()) as u32 };
let num_shl = power_used.saturating_mul(exp);
let res = match (1 as Self).checked_shl(num_shl) {
Some(x) => x,
None => 0
};

let sign = self.is_negative() && exp & 1 != 0;
let overflow = res == 0 || (sign && res == Self::MIN);
if sign {
(res.wrapping_neg(), overflow)
} else {
(res, overflow)
}
} else {
if exp == 0 {
return (1,false);
}
let mut base = self;
let mut acc: Self = 1;
let mut overflown = false;
// Scratch space for storing results of overflowing_mul.
let mut r;

while exp > 1 {
if (exp & 1) == 1 {
r = acc.overflowing_mul(base);
acc = r.0;
overflown |= r.1;
}
exp /= 2;
r = base.overflowing_mul(base);
base = r.0;
overflown |= r.1;
}
exp /= 2;
r = base.overflowing_mul(base);
base = r.0;
overflown |= r.1;
}

// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
r = acc.overflowing_mul(base);
r.1 |= overflown;
r
// since exp!=0, finally the exp must be 1.
// Deal with the final bit of the exponent separately, since
// squaring the base afterwards is not necessary and may cause a
// needless overflow.
r = acc.overflowing_mul(base);
r.1 |= overflown;
r
}
}

/// Raises self to the power of `exp`, using exponentiation by squaring.
Expand All @@ -2086,30 +2149,34 @@ macro_rules! int_impl {
#[rustc_const_stable(feature = "const_int_pow", since = "1.50.0")]
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[rustc_allow_const_fn_unstable(is_val_statically_known)]
#[inline]
#[rustc_inherit_overflow_checks]
#[track_caller] // Hides the hackish overflow check for powers of two.
#[rustc_allow_const_fn_unstable(is_val_statically_known)]
pub const fn pow(self, mut exp: u32) -> Self {
// SAFETY: This path has the same behavior as the other.
if unsafe { intrinsics::is_val_statically_known(self) }
&& self > 0
&& (self & (self - 1) == 0)
&& self.unsigned_abs().is_power_of_two()
{
// SAFETY: We just checked this is a power of two. and above zero.
let power_used = unsafe { intrinsics::cttz_nonzero(self) as u32 };
let power_used = unsafe { intrinsics::cttz_nonzero(self.wrapping_abs()) as u32 };
let num_shl = power_used.saturating_mul(exp);
let res = match (1 as Self).checked_shl(num_shl) {
Some(x) => x,
None => 0
};

let sign = self.is_negative() && exp & 1 != 0;
#[allow(arithmetic_overflow)]
if res <= 0 {
if res == 0 || (!sign && res == Self::MIN) {
// So it panics.
_ = Self::MAX * Self::MAX;
}
res
if sign {
res.wrapping_neg()
} else {
res
}
} else {
if exp == 0 {
return 1;
Expand Down
Loading

0 comments on commit cdc2f08

Please sign in to comment.