optimize: pow when base is a power of two

Kmeakin · Kmeakin · commit c2f36c5b55b2 · 2025-11-23T17:30:51.000Z
if base == 2 ** k, then
   (2 ** k) ** n
== 2 ** (k * n)
== 1 &lt;&lt; (k * n)
diff --git a/library/core/src/num/uint_macros.rs b/library/core/src/num/uint_macros.rs
@@ -2072,12 +2072,20 @@ macro_rules! uint_impl {
                       without modifying the original"]
         #[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;
 
+            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 {
@@ -2094,6 +2102,10 @@ macro_rules! uint_impl {
                 return acc.checked_mul(base);
             }
 
+            if exp == 0 {
+                return Some(1);
+            }
+
             loop {
                 if (exp & 1) == 1 {
                     acc = try_opt!(acc.checked_mul(base));
@@ -3238,14 +3250,24 @@ 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 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 {
@@ -3266,6 +3288,10 @@ macro_rules! uint_impl {
                 return (acc, overflow);
             }
 
+            if exp == 0 {
+                return (1, false);
+            }
+
             loop {
                 if (exp & 1) == 1 {
                     (acc, tmp_overflow) = acc.overflowing_mul(base);
@@ -3295,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 {
@@ -3317,6 +3356,10 @@ macro_rules! uint_impl {
                 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
diff --git a/library/coretests/tests/num/uint_macros.rs b/library/coretests/tests/num/uint_macros.rs
@@ -271,13 +271,19 @@ macro_rules! uint_module {
             assert_eq!(from_str::<$T>("x"), None);
         }
 
-        // Not const because `catch_unwind` is not const yet.
+        // Not const because overflow always panics during const evaluation, and
+        // `catch_unwind` is not const yet.
         #[test]
-        fn strict_pow() {
+        fn panicking_pow() {
             use std::panic::catch_unwind;
 
             {
                 const R: $T = 0;
+                assert_eq!(R.pow(0), 1 as $T);
+                assert_eq!(R.pow(1), 0 as $T);
+                assert_eq!(R.pow(2), 0 as $T);
+                assert_eq!(R.pow(128), 0 as $T);
+
                 assert_eq!(R.strict_pow(0), 1 as $T);
                 assert_eq!(R.strict_pow(1), 0 as $T);
                 assert_eq!(R.strict_pow(2), 0 as $T);
@@ -286,6 +292,11 @@ macro_rules! uint_module {
 
             {
                 const R: $T = 1;
+                assert_eq!(R.pow(0), 1 as $T);
+                assert_eq!(R.pow(1), 1 as $T);
+                assert_eq!(R.pow(2), 1 as $T);
+                assert_eq!(R.pow(128), 1 as $T);
+
                 assert_eq!(R.strict_pow(0), 1 as $T);
                 assert_eq!(R.strict_pow(1), 1 as $T);
                 assert_eq!(R.strict_pow(2), 1 as $T);
@@ -294,48 +305,24 @@ macro_rules! uint_module {
 
             {
                 const R: $T = 2;
+                assert_eq!(R.pow(0), 1 as $T);
+                assert_eq!(R.pow(1), 2 as $T);
+                assert_eq!(R.pow(2), 4 as $T);
+                assert_eq!(R.pow(128), 0 as $T);
+                assert_eq!(R.pow(129), 0 as $T);
+
                 assert_eq!(R.strict_pow(0), 1 as $T);
                 assert_eq!(R.strict_pow(1), 2 as $T);
                 assert_eq!(R.strict_pow(2), 4 as $T);
                 assert!(catch_unwind(|| R.strict_pow(128)).is_err());
             }
 
+            // If `k * exp` wraps around into a valid shift amount, `pow` should still return 0/panic
             {
-                const R: $T = $T::MAX;
-                assert_eq!(R.strict_pow(0), 1 as $T);
-                assert_eq!(R.strict_pow(1), R as $T);
-                assert!(catch_unwind(|| R.strict_pow(2)).is_err());
-                assert!(catch_unwind(|| R.strict_pow(128)).is_err());
-            }
-        }
-
-        // Not const because overflow always panics during const evaluation, and
-        // `catch_unwind` is not const yet.
-        #[test]
-        fn pow() {
-            {
-                const R: $T = 0;
-                assert_eq!(R.pow(0), 1 as $T);
-                assert_eq!(R.pow(1), 0 as $T);
-                assert_eq!(R.pow(2), 0 as $T);
-                assert_eq!(R.pow(128), 0 as $T);
-            }
-
-            {
-                const R: $T = 1;
-                assert_eq!(R.pow(0), 1 as $T);
-                assert_eq!(R.pow(1), 1 as $T);
-                assert_eq!(R.pow(2), 1 as $T);
-                assert_eq!(R.pow(128), 1 as $T);
-            }
-
-            {
-                const R: $T = 2;
-                assert_eq!(R.pow(0), 1 as $T);
-                assert_eq!(R.pow(1), 2 as $T);
-                assert_eq!(R.pow(2), 4 as $T);
-                assert_eq!(R.pow(128), 0 as $T);
-                assert_eq!(R.pow(129), 0 as $T);
+                const R: $T = 4;
+                const HALF: u32 = u32::MAX / 2 + 1;
+                assert_eq!(R.pow(HALF), 0 as $T);
+                assert!(catch_unwind(|| R.strict_pow(HALF)).is_err());
             }
 
             {
@@ -345,6 +332,11 @@ macro_rules! uint_module {
                 assert_eq!(R.pow(2), 1 as $T);
                 assert_eq!(R.pow(128), 1 as $T);
                 assert_eq!(R.pow(129), R as $T);
+
+                assert_eq!(R.strict_pow(0), 1 as $T);
+                assert_eq!(R.strict_pow(1), R as $T);
+                assert!(catch_unwind(|| R.strict_pow(2)).is_err());
+                assert!(catch_unwind(|| R.strict_pow(128)).is_err());
             }
         }
 
@@ -443,6 +435,17 @@ macro_rules! uint_module {
                     assert_eq_const_safe!($T: R.saturating_pow(129), $T::MAX);
                 }
 
+                // overflow in the shift caclculation should result in the final
+                // result being 0 rather than accidentally succeeding due to a
+                // shift within the word size
+                // ie `4 ** 0x8000_0000` should give 0 rather than 1 << 0
+                {
+                    const R: $T = 4;
+                    const HALF: u32 = u32::MAX / 2 + 1;
+                    assert_eq_const_safe!($T: R.wrapping_pow(HALF), 0 as $T);
+                    assert_eq_const_safe!(($T, bool): R.overflowing_pow(HALF), (0 as $T, true));
+                }
+
                 {
                     const R: $T = $T::MAX;
                     assert_eq_const_safe!($T: R.wrapping_pow(0), 1 as $T);
diff --git a/tests/codegen-llvm/pow_known_base.rs b/tests/codegen-llvm/pow_known_base.rs
@@ -0,0 +1,43 @@
+//@ compile-flags: -Copt-level=3
+// Test that `pow` can use a faster implementation when `base` is a
+// known power of two
+
+#![crate_type = "lib"]
+
+// CHECK-LABEL: @pow2
+#[no_mangle]
+pub fn pow2(exp: u32) -> u32 {
+    // CHECK: %[[OVERFLOW:.+]] = icmp ugt i32 %exp, 31
+    // CHECK: %[[POW:.+]] = shl nuw i32 1, %exp
+    // CHECK: %[[RET:.+]] = select i1 %[[OVERFLOW]], i32 0, i32 %[[POW]]
+    // CHECK: ret i32 %[[RET]]
+    2u32.pow(exp)
+}
+
+// 4 ** n == 2 ** (2 * n) == 1 << (2 * n)
+// CHECK-LABEL: @pow4
+#[no_mangle]
+pub fn pow4(exp: u32) -> u32 {
+    // CHECK: %[[SHIFT_AMOUNT:.+]] = shl i32 %exp, 1
+    // CHECK: %[[ICMP1:.+]] = icmp slt i32 %exp, 0
+    // CHECK: %[[ICMP2:.+]] = icmp ugt i32 %[[SHIFT_AMOUNT]], 31
+    // CHECK: %[[OVERFLOW:.+]] = or i1 %[[ICMP1]], %[[ICMP2]]
+    // CHECK: %[[POW:.+]] = shl nuw i32 1, %[[SHIFT_AMOUNT]]
+    // CHECK: %[[RET:.+]] = select i1 %[[OVERFLOW]], i32 0, i32 %[[POW]]
+    // CHECK: ret i32 %[[RET]]
+    4u32.pow(exp)
+}
+
+// 16 ** n == 2 ** (4 * n) == 1 << (4 * n)
+// CHECK-LABEL: @pow16
+#[no_mangle]
+pub fn pow16(exp: u32) -> u32 {
+    // CHECK: %[[SHIFT_AMOUNT:.+]] = shl i32 %exp, 2
+    // CHECK: %[[ICMP1:.+]] = icmp ugt i32 %exp, 1073741823
+    // CHECK: %[[ICMP2:.+]] = icmp ugt i32 %[[SHIFT_AMOUNT]], 31
+    // CHECK: %[[OVERFLOW:.+]] = or i1 %[[ICMP1]], %[[ICMP2]]
+    // CHECK: %[[POW:.+]] = shl nuw i32 1, %[[SHIFT_AMOUNT]]
+    // CHECK: %[[RET:.+]] = select i1 %[[OVERFLOW]], i32 0, i32 %[[POW]]
+    // CHECK: ret i32 %[[RET]]
+    16u32.pow(exp)
+}
diff --git a/tests/ui/numbers-arithmetic/overflowing-pow-signed.rs b/tests/ui/numbers-arithmetic/overflowing-pow-signed.rs
@@ -1,6 +1,6 @@
 //@ run-fail
 //@ regex-error-pattern: thread 'main'.*panicked
-//@ error-pattern: attempt to multiply with overflow
+//@ regex-error-pattern: attempt to (multiply|shift left) with overflow
 //@ needs-subprocess
 //@ compile-flags: -C debug-assertions
 
diff --git a/tests/ui/numbers-arithmetic/overflowing-pow-unsigned.rs b/tests/ui/numbers-arithmetic/overflowing-pow-unsigned.rs
@@ -1,6 +1,6 @@
 //@ run-fail
 //@ regex-error-pattern: thread 'main'.*panicked
-//@ error-pattern: attempt to multiply with overflow
+//@ regex-error-pattern: attempt to (multiply|shift left) with overflow
 //@ needs-subprocess
 //@ compile-flags: -C debug-assertions
 
