Maybe fix the regression?

Author: Centri3

compiler/rustc_mir_transform/src/simplify_pow_of_two.rs

@@ -3,7 +3,6 @@
 
 use crate::MirPass;
 use rustc_const_eval::interpret::{ConstValue, Scalar};
-use rustc_hir::definitions::{DefPathData, DisambiguatedDefPathData};
 use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, Ty, TyCtxt, UintTy};
@@ -33,18 +32,14 @@ impl<'tcx> MirPass<'tcx> for SimplifyPowOfTwo {
                 && let Some(def_id) = func.const_fn_def().map(|def| def.0)
                 && let def_path = tcx.def_path(def_id)
                 && tcx.crate_name(def_path.krate) == sym::core
-                // FIXME(Centri3): I feel like we should do this differently...
-                && let [
-                    DisambiguatedDefPathData { data: DefPathData::TypeNs(sym::num), disambiguator: 0 },
-                    DisambiguatedDefPathData { data: DefPathData::Impl, .. },
-                    DisambiguatedDefPathData { data: DefPathData::ValueNs(sym::pow), .. },
-                ] = &*def_path.data
                 && let [recv, exp] = args.as_slice()
                 && let Some(recv_const) = recv.constant()
                 && let ConstantKind::Val(
                     ConstValue::Scalar(Scalar::Int(recv_int)),
                     recv_ty,
                 ) = recv_const.literal
+                && recv_ty.is_integral()
+                && tcx.item_name(def_id) == sym::pow
                 && let Ok(recv_val) = match recv_ty.kind() {
                     ty::Int(_) => {
                         let result = recv_int.try_to_int(recv_int.size()).unwrap_or(-1).max(0);
@@ -63,8 +58,8 @@ impl<'tcx> MirPass<'tcx> for SimplifyPowOfTwo {
                 // `0` would be `1.pow()`, which we shouldn't try to optimize as it's
                 // already entirely optimized away
                 && power_used != 0.0
-                // Same here
-                && recv_val != 0
+                // `-inf` would be `0.pow()`
+                && power_used.is_finite()
             {
                 let power_used = power_used as u32;
                 let loc = Location { block: i, statement_index: bb.statements.len() };
@@ -110,7 +105,8 @@ impl<'tcx> MirPass<'tcx> for SimplifyPowOfTwo {
                 );
                 let shl_result = patch.new_temp(Ty::new_bool(tcx), span);
 
-                // Whether the shl will overflow, if so we return 0
+                // Whether the shl will overflow, if so we return 0. We can do this rather
+                // than doing a shr because only one bit is set on any power of two
                 patch.add_assign(
                     loc,
                     shl_result.into(),
@@ -130,12 +126,12 @@ impl<'tcx> MirPass<'tcx> for SimplifyPowOfTwo {
                     ),
                 );
 
-                let should_be_zero_bool = patch.new_temp(Ty::new_bool(tcx), span);
-                let should_be_zero = patch.new_temp(recv_ty, span);
+                let fine_bool = patch.new_temp(Ty::new_bool(tcx), span);
+                let fine = patch.new_temp(recv_ty, span);
 
                 patch.add_assign(
                     loc,
-                    should_be_zero_bool.into(),
+                    fine_bool.into(),
                     Rvalue::BinaryOp(
                         BinOp::BitOr,
                         Box::new((
@@ -147,94 +143,47 @@ impl<'tcx> MirPass<'tcx> for SimplifyPowOfTwo {
 
                 patch.add_assign(
                     loc,
-                    should_be_zero.into(),
-                    Rvalue::Cast(
-                        CastKind::IntToInt,
-                        Operand::Copy(should_be_zero_bool.into()),
-                        recv_ty,
-                    ),
+                    fine.into(),
+                    Rvalue::Cast(CastKind::IntToInt, Operand::Copy(fine_bool.into()), recv_ty),
                 );
 
-                let shl_exp_ty = patch.new_temp(exp_ty, span);
                 let shl = patch.new_temp(recv_ty, span);
 
                 patch.add_assign(
                     loc,
-                    shl_exp_ty.into(),
+                    shl.into(),
                     Rvalue::BinaryOp(
                         BinOp::Shl,
                         Box::new((
                             Operand::Constant(Box::new(Constant {
                                 span,
                                 user_ty: None,
-                                literal: ConstantKind::Val(ConstValue::from_u32(1), exp_ty),
+                                literal: ConstantKind::Val(
+                                    ConstValue::Scalar(Scalar::from_uint(1u128, recv_int.size())),
+                                    recv_ty,
+                                ),
                             })),
                             Operand::Copy(num_shl.into()),
                         )),
                     ),
                 );
 
-                patch.add_assign(
-                    loc,
-                    shl.into(),
-                    Rvalue::Cast(CastKind::IntToInt, Operand::Copy(shl_exp_ty.into()), recv_ty),
-                );
-
                 patch.add_assign(
                     loc,
                     *destination,
                     Rvalue::BinaryOp(
                         BinOp::MulUnchecked,
-                        Box::new((Operand::Copy(shl.into()), Operand::Copy(should_be_zero.into()))),
+                        Box::new((Operand::Copy(shl.into()), Operand::Copy(fine.into()))),
                     ),
                 );
 
-                // shl doesn't set the overflow flag on x86_64 or even in Rust, so shr to
-                // see if it overflowed. If it equals 1, it did not, but we also need to
-                // check `shl_result` to ensure that if this is a multiple of the type's
-                // size it won't wrap back over to 1
-                //
                 // FIXME(Centri3): Do we use `debug_assertions` or `overflow_checks` here?
                 if tcx.sess.opts.debug_assertions {
-                    let shr = patch.new_temp(recv_ty, span);
-                    let shl_eq_shr = patch.new_temp(Ty::new_bool(tcx), span);
-                    let overflowed = patch.new_temp(Ty::new_bool(tcx), span);
-
-                    patch.add_assign(
-                        loc,
-                        shr.into(),
-                        Rvalue::BinaryOp(
-                            BinOp::Shr,
-                            Box::new((Operand::Copy(shl.into()), Operand::Copy(num_shl.into()))),
-                        ),
-                    );
-
-                    patch.add_assign(
-                        loc,
-                        shl_eq_shr.into(),
-                        Rvalue::BinaryOp(
-                            BinOp::Eq,
-                            Box::new((Operand::Copy(shl.into()), Operand::Copy(shr.into()))),
-                        ),
-                    );
-
-                    patch.add_assign(
-                        loc,
-                        overflowed.into(),
-                        Rvalue::BinaryOp(
-                            BinOp::BitOr,
-                            Box::new((
-                                Operand::Copy(shl_eq_shr.into()),
-                                Operand::Copy(shl_result.into()),
-                            )),
-                        ),
-                    );
-
                     patch.patch_terminator(
                         i,
                         TerminatorKind::Assert {
-                            cond: Operand::Copy(overflowed.into()),
-                            expected: false,
+                            cond: Operand::Copy(fine_bool.into()),
+                            expected: true,
                             msg: Box::new(AssertMessage::Overflow(
                                 // For consistency with the previous error message, though
                                 // it's technically incorrect

tests/codegen/simplify-pow-of-two-debug-assertions.rs

@@ -0,0 +1,19 @@
+// compile-flags: -Copt-level=3 -Cdebug-assertions=true
+
+// CHECK-LABEL: @slow_2_u(
+#[no_mangle]
+fn slow_2_u(a: u32) -> u32 {
+    // CHECK: %_3 = icmp ult i32 %a, 32
+    // CHECK-NEXT: br i1 %_3, label %bb1, label %panic, !prof !{{[0-9]+}}
+    // CHECK-EMPTY:
+    // CHECK-NEXT: bb1:
+    // CHECK-NEXT: %_01 = shl nuw i32 1, %a
+    // CHECK-NEXT: ret i32 %_0
+    // CHECK-EMPTY:
+    // CHECK-NEXT: panic:
+    2u32.pow(a)
+}
+
+fn main() {
+    slow_2_u(2);
+}

tests/codegen/simplify-pow-of-two.rs

@@ -0,0 +1,16 @@
+// compile-flags: -Copt-level=3
+
+// CHECK-LABEL: @slow_2_u(
+#[no_mangle]
+fn slow_2_u(a: u32) -> u32 {
+    // CHECK: %_3 = icmp ult i32 %a, 32
+    // CHECK-NEXT: %_5 = zext i1 %_3 to i32
+    // CHECK-NEXT: %0 = and i32 %a, 31
+    // CHECK-NEXT: %_01 = shl nuw i32 %_5, %0
+    // CHECK-NEXT: ret i32 %_01
+    2u32.pow(a)
+}
+
+fn main() {
+    slow_2_u(2);
+}

tests/mir-opt/simplify_pow_of_two_no_overflow_checks.slow_256_i.SimplifyPowOfTwo.diff

@@ -9,20 +9,18 @@
 +     let mut _4: bool;
 +     let mut _5: bool;
 +     let mut _6: i32;
-+     let mut _7: u32;
-+     let mut _8: i32;
++     let mut _7: i32;
 
       bb0: {
           StorageLive(_2);
           _2 = _1;
--         _0 = core::num::<impl i32>::pow(const 256_i32, move _2) -> [return: bb1, unwind continue];
+-         _0 = core::num::<impl i32>::pow(const 256_i32, move _2) -> [return: bb1, unwind unreachable];
 +         _3 = CheckedMul(move _2, const 8_u32);
 +         _4 = Lt((_3.0: u32), const 32_u32);
 +         _5 = BitOr((_3.1: bool), _4);
 +         _6 = _5 as i32 (IntToInt);
-+         _7 = Shl(const 1_u32, (_3.0: u32));
-+         _8 = _7 as i32 (IntToInt);
-+         _0 = MulUnchecked(_8, _6);
++         _7 = Shl(const 1_i32, (_3.0: u32));
++         _0 = MulUnchecked(_7, _6);
 +         goto -> bb1;
       }
 

tests/mir-opt/simplify_pow_of_two_no_overflow_checks.slow_256_u.SimplifyPowOfTwo.diff

@@ -10,19 +10,17 @@
 +     let mut _5: bool;
 +     let mut _6: u32;
 +     let mut _7: u32;
-+     let mut _8: u32;
 
       bb0: {
           StorageLive(_2);
           _2 = _1;
--         _0 = core::num::<impl u32>::pow(const 256_u32, move _2) -> [return: bb1, unwind continue];
+-         _0 = core::num::<impl u32>::pow(const 256_u32, move _2) -> [return: bb1, unwind unreachable];
 +         _3 = CheckedMul(move _2, const 8_u32);
 +         _4 = Lt((_3.0: u32), const 32_u32);
 +         _5 = BitOr((_3.1: bool), _4);
 +         _6 = _5 as u32 (IntToInt);
 +         _7 = Shl(const 1_u32, (_3.0: u32));
-+         _8 = _7 as u32 (IntToInt);
-+         _0 = MulUnchecked(_8, _6);
++         _0 = MulUnchecked(_7, _6);
 +         goto -> bb1;
       }
 

tests/mir-opt/simplify_pow_of_two_no_overflow_checks.slow_2_i.SimplifyPowOfTwo.diff

@@ -9,20 +9,18 @@
 +     let mut _4: bool;
 +     let mut _5: bool;
 +     let mut _6: i32;
-+     let mut _7: u32;
-+     let mut _8: i32;
++     let mut _7: i32;
 
       bb0: {
           StorageLive(_2);
           _2 = _1;
--         _0 = core::num::<impl i32>::pow(const 2_i32, move _2) -> [return: bb1, unwind continue];
+-         _0 = core::num::<impl i32>::pow(const 2_i32, move _2) -> [return: bb1, unwind unreachable];
 +         _3 = CheckedMul(move _2, const 1_u32);
 +         _4 = Lt((_3.0: u32), const 32_u32);
 +         _5 = BitOr((_3.1: bool), _4);
 +         _6 = _5 as i32 (IntToInt);
-+         _7 = Shl(const 1_u32, (_3.0: u32));
-+         _8 = _7 as i32 (IntToInt);
-+         _0 = MulUnchecked(_8, _6);
++         _7 = Shl(const 1_i32, (_3.0: u32));
++         _0 = MulUnchecked(_7, _6);
 +         goto -> bb1;
       }
 

tests/mir-opt/simplify_pow_of_two_no_overflow_checks.slow_2_u.SimplifyPowOfTwo.diff

@@ -10,19 +10,17 @@
 +     let mut _5: bool;
 +     let mut _6: u32;
 +     let mut _7: u32;
-+     let mut _8: u32;
 
       bb0: {
           StorageLive(_2);
           _2 = _1;
--         _0 = core::num::<impl u32>::pow(const 2_u32, move _2) -> [return: bb1, unwind continue];
+-         _0 = core::num::<impl u32>::pow(const 2_u32, move _2) -> [return: bb1, unwind unreachable];
 +         _3 = CheckedMul(move _2, const 1_u32);
 +         _4 = Lt((_3.0: u32), const 32_u32);
 +         _5 = BitOr((_3.1: bool), _4);
 +         _6 = _5 as u32 (IntToInt);
 +         _7 = Shl(const 1_u32, (_3.0: u32));
-+         _8 = _7 as u32 (IntToInt);
-+         _0 = MulUnchecked(_8, _6);
++         _0 = MulUnchecked(_7, _6);
 +         goto -> bb1;
       }
 

tests/mir-opt/simplify_pow_of_two_no_overflow_checks.slow_4_i.SimplifyPowOfTwo.diff

@@ -9,20 +9,18 @@
 +     let mut _4: bool;
 +     let mut _5: bool;
 +     let mut _6: i32;
-+     let mut _7: u32;
-+     let mut _8: i32;
++     let mut _7: i32;
 
       bb0: {
           StorageLive(_2);
           _2 = _1;
--         _0 = core::num::<impl i32>::pow(const 4_i32, move _2) -> [return: bb1, unwind continue];
+-         _0 = core::num::<impl i32>::pow(const 4_i32, move _2) -> [return: bb1, unwind unreachable];
 +         _3 = CheckedMul(move _2, const 2_u32);
 +         _4 = Lt((_3.0: u32), const 32_u32);
 +         _5 = BitOr((_3.1: bool), _4);
 +         _6 = _5 as i32 (IntToInt);
-+         _7 = Shl(const 1_u32, (_3.0: u32));
-+         _8 = _7 as i32 (IntToInt);
-+         _0 = MulUnchecked(_8, _6);
++         _7 = Shl(const 1_i32, (_3.0: u32));
++         _0 = MulUnchecked(_7, _6);
 +         goto -> bb1;
       }
 

tests/mir-opt/simplify_pow_of_two_no_overflow_checks.slow_4_u.SimplifyPowOfTwo.diff

@@ -10,19 +10,17 @@
 +     let mut _5: bool;
 +     let mut _6: u32;
 +     let mut _7: u32;
-+     let mut _8: u32;
 
       bb0: {
           StorageLive(_2);
           _2 = _1;
--         _0 = core::num::<impl u32>::pow(const 4_u32, move _2) -> [return: bb1, unwind continue];
+-         _0 = core::num::<impl u32>::pow(const 4_u32, move _2) -> [return: bb1, unwind unreachable];
 +         _3 = CheckedMul(move _2, const 2_u32);
 +         _4 = Lt((_3.0: u32), const 32_u32);
 +         _5 = BitOr((_3.1: bool), _4);
 +         _6 = _5 as u32 (IntToInt);
 +         _7 = Shl(const 1_u32, (_3.0: u32));
-+         _8 = _7 as u32 (IntToInt);
-+         _0 = MulUnchecked(_8, _6);
++         _0 = MulUnchecked(_7, _6);
 +         goto -> bb1;
       }
 

tests/mir-opt/simplify_pow_of_two_overflow_checks.slow_256_i.SimplifyPowOfTwo.diff

@@ -9,27 +9,19 @@
 +     let mut _4: bool;
 +     let mut _5: bool;
 +     let mut _6: i32;
-+     let mut _7: u32;
-+     let mut _8: i32;
-+     let mut _9: i32;
-+     let mut _10: bool;
-+     let mut _11: bool;
++     let mut _7: i32;
 
       bb0: {
           StorageLive(_2);
           _2 = _1;
--         _0 = core::num::<impl i32>::pow(const 256_i32, move _2) -> [return: bb1, unwind continue];
+-         _0 = core::num::<impl i32>::pow(const 256_i32, move _2) -> [return: bb1, unwind unreachable];
 +         _3 = CheckedMul(move _2, const 8_u32);
 +         _4 = Lt((_3.0: u32), const 32_u32);
 +         _5 = BitOr((_3.1: bool), _4);
 +         _6 = _5 as i32 (IntToInt);
-+         _7 = Shl(const 1_u32, (_3.0: u32));
-+         _8 = _7 as i32 (IntToInt);
-+         _0 = MulUnchecked(_8, _6);
-+         _9 = Shr(_8, (_3.0: u32));
-+         _10 = Eq(_8, _9);
-+         _11 = BitOr(_10, _4);
-+         assert(!_11, "attempt to compute `{} * {}`, which would overflow", const 1_u32, (_3.0: u32)) -> [success: bb1, unwind continue];
++         _7 = Shl(const 1_i32, (_3.0: u32));
++         _0 = MulUnchecked(_7, _6);
++         assert(_5, "attempt to compute `{} * {}`, which would overflow", const 1_u32, (_3.0: u32)) -> [success: bb1, unwind unreachable];
       }
 
       bb1: {