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
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
d9cc57e
optimize: `checked_ilog` when `base` is a power of two
Kmeakin Oct 2, 2025
20f67bc
refactor: Increase test coverage for `pow`
Kmeakin Nov 22, 2025
97d25c8
refactor: Deduplicate `pow` implementations
Kmeakin Nov 22, 2025
759e81a
optimize: `pow` if `exp` is statically known
Kmeakin Nov 22, 2025
c2f36c5
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/overflow_panic.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,12 @@ pub(super) const fn rem() -> ! {
panic!("attempt to calculate the remainder with overflow")
}

#[cold]
#[track_caller]
pub(super) const fn pow() -> ! {
panic!("attempt to calculate the power with overflow")
}

#[cold]
#[track_caller]
pub(super) const fn neg() -> ! {
Expand Down
219 changes: 128 additions & 91 deletions library/core/src/num/uint_macros.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1555,6 +1555,10 @@ macro_rules! uint_impl {
///
/// ```
#[doc = concat!("assert_eq!(5", stringify!($SelfT), ".checked_ilog(5), Some(1));")]
#[doc = concat!("assert_eq!(4", stringify!($SelfT), ".checked_ilog(5), Some(0));")]
#[doc = concat!("assert_eq!(5", stringify!($SelfT), ".checked_ilog(0), None);")]
#[doc = concat!("assert_eq!(5", stringify!($SelfT), ".checked_ilog(1), None);")]
#[doc = concat!("assert_eq!(0", stringify!($SelfT), ".checked_ilog(1), None);")]
/// ```
#[stable(feature = "int_log", since = "1.67.0")]
#[rustc_const_stable(feature = "int_log", since = "1.67.0")]
Expand All @@ -1569,9 +1573,14 @@ macro_rules! uint_impl {
// applied by the compiler. If you want those specific bases,
// use `.checked_ilog2()` or `.checked_ilog10()` directly.
if core::intrinsics::is_val_statically_known(base) {
if base == 2 {
return self.checked_ilog2();
} else if base == 10 {
// change of base:
// if base == 2 ** k, then
// log(base, n) == log(2, n) / k
if base.is_power_of_two() && base > 1 {
let k = base.ilog2();
return Some(try_opt!(self.checked_ilog2()) / k);
}
if base == 10 {
return self.checked_ilog10();
}
}
Expand Down Expand Up @@ -2063,11 +2072,39 @@ macro_rules! uint_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) && base.is_power_of_two() {
// change of base:
// if base == 2 ** k, then
// (2 ** k) ** n
// == 2 ** (k * n)
// == 1 << (k * n)
let k = base.ilog2();
let shift = try_opt!(k.checked_mul(exp));
return (1 as Self).checked_shl(shift);
}

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);
}

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

loop {
if (exp & 1) == 1 {
Expand Down Expand Up @@ -2108,23 +2145,10 @@ macro_rules! uint_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) {
None => overflow_panic::pow(),
Some(a) => a,
}
}

Expand Down Expand Up @@ -2584,43 +2608,9 @@ macro_rules! uint_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 @@ -3260,30 +3250,60 @@ macro_rules! uint_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) && base.is_power_of_two() {
// change of base:
// if base == 2 ** k, then
// (2 ** k) ** n
// == 2 ** (k * n)
// == 1 << (k * n)
let k = base.ilog2();
let Some(shift) = k.checked_mul(exp) else {
return (0, true)
};
return ((1 as Self).unbounded_shl(shift), shift >= Self::BITS)
}

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);
}

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

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 @@ -3301,12 +3321,25 @@ macro_rules! uint_impl {
#[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(base) && base.is_power_of_two() {
// change of base:
// if base == 2 ** k, then
// (2 ** k) ** n
// == 2 ** (k * n)
// == 1 << (k * n)
let k = base.ilog2();
// Panic on overflow if `-C overflow-checks` is enabled.
// Otherwise will be optimized out
let _overflow_check = (1 as Self) << (k * exp);
let Some(shift) = k.checked_mul(exp) else {
return 0
};
return (1 as Self).unbounded_shl(shift)
}

if intrinsics::is_val_statically_known(exp) {
while exp > 1 {
if (exp & 1) == 1 {
Expand All @@ -3320,23 +3353,27 @@ macro_rules! uint_impl {
// 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
} 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;
}
return acc * base;
}

if exp == 0 {
return 1;
}

// 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;
}
exp /= 2;
base = base * base;
}
}

Expand Down
Loading
Loading