optimize: pow when base is a power of two

if base == 2 ** k, then
   (2 ** k) ** n
== 2 ** (k * n)
== 1 << (k * n)

Author: Kmeakin

library/core/src/num/uint_macros.rs

@@ -2078,6 +2078,17 @@ macro_rules! uint_impl {
             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 {
@@ -3246,6 +3257,18 @@ macro_rules! uint_impl {
             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 (shift, overflow1) = k.overflowing_mul(exp);
+                let (ret, overflow2) = ((1 as Self).unbounded_shl(shift), shift >= Self::BITS);
+                return (ret, overflow1 | overflow2)
+            }
+
             if intrinsics::is_val_statically_known(exp) {
                 while exp > 1 {
                     if (exp & 1) == 1 {
@@ -3301,6 +3324,20 @@ macro_rules! uint_impl {
             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();
+                let shift = k * exp;
+                // Panic on overflow if `-C overflow-checks` is enabled.
+                // Otherwise will be optimized out
+                let _overflow_check = (1 as Self) << shift;
+                return (1 as Self).unbounded_shl(shift)
+            }
+
             if intrinsics::is_val_statically_known(exp) {
                 while exp > 1 {
                     if (exp & 1) == 1 {

tests/codegen-llvm/pow_known_base.rs

@@ -0,0 +1,36 @@
+//@ 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: %[[SHIFT_AMOUNT:.+]] = and i32 %exp, 31
+    // CHECK: %[[POW2:.+]] = shl nuw i32 1, %[[SHIFT_AMOUNT]]
+    // CHECK: ret i32 %[[POW2]]
+    2u32.pow(exp)
+}
+
+// 4 ** n == 2 ** (2 * n) == 1 << (2 * n)
+// CHECK-LABEL: @pow4
+#[no_mangle]
+pub fn pow4(exp: u32) -> u32 {
+    // CHECK: %[[EXP2:.+]] = shl i32 %exp, 1
+    // CHECK: %[[SHIFT_AMOUNT:.+]] = and i32 %[[EXP2]], 30
+    // CHECK: %[[POW4:.+]] = shl nuw nsw i32 1, %[[SHIFT_AMOUNT]]
+    // CHECK: ret i32 %[[POW4]]
+    4u32.pow(exp)
+}
+
+// 16 ** n == 2 ** (4 * n) == 1 << (4 * n)
+// CHECK-LABEL: @pow16
+#[no_mangle]
+pub fn pow16(exp: u32) -> u32 {
+    // CHECK: %[[EXP2:.+]] = shl i32 %exp, 2
+    // CHECK: %[[SHIFT_AMOUNT:.+]] = and i32 %[[EXP2]], 28
+    // CHECK: %[[POW16:.+]] = shl nuw nsw i32 1, %[[SHIFT_AMOUNT]]
+    // CHECK: ret i32 %[[POW16]]
+    16u32.pow(exp)
+}

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
 

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