Skip to content

Optimize checked_ilog and pow when base is a power of two #147250

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.

Sign up for GitHub

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

Open
Kmeakin wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Optimize checked_ilog and pow when base is a power of two #147250

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7c61ed7
optimize: `checked_ilog` when `base` is a power of two
Kmeakin Oct 2, 2025
7ae2563
refactor: Increase test coverage for `pow`
Kmeakin Nov 22, 2025
c52c2ee
refactor: Deduplicate `pow` implementations
Kmeakin Nov 22, 2025
55e1407
optimize: `pow` if `exp` is statically known
Kmeakin Nov 22, 2025
b0aca14
optimize: `pow` when `base` is a power of two
Kmeakin Oct 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions library/core/src/num/imp/overflow_panic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,3 +49,9 @@ pub(in crate::num) const fn shr() -> ! {
pub(in crate::num) const fn shl() -> ! {
panic!("attempt to shift left with overflow")
}

#[cold]
#[track_caller]
pub(in crate::num) const fn pow() -> ! {
panic!("attempt to exponentiate with overflow")
}
195 changes: 89 additions & 106 deletions library/core/src/num/int_macros.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1761,11 +1761,41 @@ macro_rules! int_impl {
without modifying the original"]
#[inline]
pub const fn checked_pow(self, mut exp: u32) -> Option<Self> {
let mut base = self;
let mut acc: Self = 1;

if intrinsics::is_val_statically_known(base) {
if base.unsigned_abs().is_power_of_two() {
let k = base.unsigned_abs().ilog2();
let shift = try_opt!(k.checked_mul(exp));
let magnitude = try_opt!((1 as Self).checked_shl(shift));
return if base < 0 && (exp % 2) == 1 {
Some(magnitude.wrapping_neg())
} else {
Some(magnitude)
}
}
}

if exp == 0 {
return Some(1);
}
let mut base = self;
let mut acc: Self = 1;

if intrinsics::is_val_statically_known(exp) {
while exp > 1 {
if (exp & 1) == 1 {
acc = try_opt!(acc.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.
return acc.checked_mul(base);
}

loop {
if (exp & 1) == 1 {
Expand Down Expand Up @@ -1807,23 +1837,10 @@ macro_rules! int_impl {
without modifying the original"]
#[inline]
#[track_caller]
pub const fn strict_pow(self, mut exp: u32) -> Self {
if exp == 0 {
return 1;
}
let mut base = self;
let mut acc: Self = 1;

loop {
if (exp & 1) == 1 {
acc = acc.strict_mul(base);
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return acc;
}
}
exp /= 2;
base = base.strict_mul(base);
pub const fn strict_pow(self, exp: u32) -> Self {
match self.checked_pow(exp) {
Some(x) => x,
None => imp::overflow_panic::pow(),
}
}

Expand Down Expand Up @@ -2438,43 +2455,9 @@ macro_rules! int_impl {
#[must_use = "this returns the result of the operation, \
without modifying the original"]
#[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;

if intrinsics::is_val_statically_known(exp) {
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.
acc.wrapping_mul(base)
} else {
// This is faster than the above when the exponent is not known
// at compile time. We can't use the same code for the constant
// exponent case because LLVM is currently unable to unroll
// this loop.
loop {
if (exp & 1) == 1 {
acc = acc.wrapping_mul(base);
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return acc;
}
}
exp /= 2;
base = base.wrapping_mul(base);
}
}
pub const fn wrapping_pow(self, exp: u32) -> Self {
let (a, _) = self.overflowing_pow(exp);
a
}

/// Calculates `self` + `rhs`.
Expand Down Expand Up @@ -3039,30 +3022,62 @@ macro_rules! int_impl {
without modifying the original"]
#[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;
let mut overflow = false;
let mut tmp_overflow;

if intrinsics::is_val_statically_known(base) {
if base.unsigned_abs().is_power_of_two() {
let k = base.unsigned_abs().ilog2();
let Some(shift) = k.checked_mul(exp) else {
return (0, true)
};
let magnitude = (1 as Self).unbounded_shl(shift);
return if base < 0 && (exp % 2) == 1 {
(magnitude.wrapping_neg(), shift >= Self::BITS)
} else {
(magnitude, shift >= Self::BITS)
}
}
}

if exp == 0 {
return (1, false);
}

if intrinsics::is_val_statically_known(exp) {
while exp > 1 {
if (exp & 1) == 1 {
(acc, tmp_overflow) = acc.overflowing_mul(base);
overflow |= tmp_overflow;
}
exp /= 2;
(base, tmp_overflow) = base.overflowing_mul(base);
overflow |= tmp_overflow;
}

// 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, tmp_overflow) = acc.overflowing_mul(base);
overflow |= tmp_overflow;
return (acc, overflow);
}

loop {
if (exp & 1) == 1 {
r = acc.overflowing_mul(base);
(acc, tmp_overflow) = acc.overflowing_mul(base);
overflow |= tmp_overflow;
// since exp!=0, finally the exp must be 1.
if exp == 1 {
r.1 |= overflown;
return r;
return (acc, overflow);
}
acc = r.0;
overflown |= r.1;
}
exp /= 2;
r = base.overflowing_mul(base);
base = r.0;
overflown |= r.1;
(base, tmp_overflow) = base.overflowing_mul(base);
overflow |= tmp_overflow;
}
}

Expand All @@ -3082,43 +3097,11 @@ macro_rules! int_impl {
without modifying the original"]
#[inline]
#[rustc_inherit_overflow_checks]
pub const fn pow(self, mut exp: u32) -> Self {
if exp == 0 {
return 1;
}
let mut base = self;
let mut acc = 1;

if intrinsics::is_val_statically_known(exp) {
while exp > 1 {
if (exp & 1) == 1 {
acc = acc * base;
}
exp /= 2;
base = base * 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 * base
pub const fn pow(self, exp: u32) -> Self {
if intrinsics::overflow_checks() {
self.strict_pow(exp)
} else {
// This is faster than the above when the exponent is not known
// at compile time. We can't use the same code for the constant
// exponent case because LLVM is currently unable to unroll
// this loop.
loop {
if (exp & 1) == 1 {
acc = acc * base;
// since exp!=0, finally the exp must be 1.
if exp == 1 {
return acc;
}
}
exp /= 2;
base = base * base;
}
self.wrapping_pow(exp)
}
}

Expand Down
Loading
Loading