Implement checked Shl/Shr at MIR building.

7 files changed, 109 insertions, 107 deletions

diff --git a/compiler/rustc_codegen_cranelift/src/num.rs b/compiler/rustc_codegen_cranelift/src/num.rs
index c058ece96d8..fbecdab158c 100644
--- a/compiler/rustc_codegen_cranelift/src/num.rs
+++ b/compiler/rustc_codegen_cranelift/src/num.rs
@@ -170,14 +170,6 @@ pub(crate) fn codegen_checked_int_binop<'tcx>(
     in_lhs: CValue<'tcx>,
     in_rhs: CValue<'tcx>,
 ) -> CValue<'tcx> {
-    if bin_op != BinOp::Shl && bin_op != BinOp::Shr {
-        assert_eq!(
-            in_lhs.layout().ty,
-            in_rhs.layout().ty,
-            "checked int binop requires lhs and rhs of same type"
-        );
-    }
-
     let lhs = in_lhs.load_scalar(fx);
     let rhs = in_rhs.load_scalar(fx);
 
@@ -271,21 +263,6 @@ pub(crate) fn codegen_checked_int_binop<'tcx>(
                 _ => unreachable!("invalid non-integer type {}", ty),
             }
         }
-        BinOp::Shl => {
-            let val = fx.bcx.ins().ishl(lhs, rhs);
-            let ty = fx.bcx.func.dfg.value_type(val);
-            let max_shift = i64::from(ty.bits()) - 1;
-            let has_overflow = fx.bcx.ins().icmp_imm(IntCC::UnsignedGreaterThan, rhs, max_shift);
-            (val, has_overflow)
-        }
-        BinOp::Shr => {
-            let val =
-                if !signed { fx.bcx.ins().ushr(lhs, rhs) } else { fx.bcx.ins().sshr(lhs, rhs) };
-            let ty = fx.bcx.func.dfg.value_type(val);
-            let max_shift = i64::from(ty.bits()) - 1;
-            let has_overflow = fx.bcx.ins().icmp_imm(IntCC::UnsignedGreaterThan, rhs, max_shift);
-            (val, has_overflow)
-        }
         _ => bug!("binop {:?} on checked int/uint lhs: {:?} rhs: {:?}", bin_op, in_lhs, in_rhs),
     };
 
diff --git a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
index 3d856986fb4..13c4fa132d8 100644
--- a/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/rvalue.rs
@@ -663,17 +663,6 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 };
                 bx.checked_binop(oop, input_ty, lhs, rhs)
             }
-            mir::BinOp::Shl | mir::BinOp::Shr => {
-                let lhs_llty = bx.cx().val_ty(lhs);
-                let rhs_llty = bx.cx().val_ty(rhs);
-                let invert_mask = common::shift_mask_val(bx, lhs_llty, rhs_llty, true);
-                let outer_bits = bx.and(rhs, invert_mask);
-
-                let of = bx.icmp(IntPredicate::IntNE, outer_bits, bx.cx().const_null(rhs_llty));
-                let val = self.codegen_scalar_binop(bx, op, lhs, rhs, input_ty);
-
-                (val, of)
-            }
             _ => bug!("Operator `{:?}` is not a checkable operator", op),
         };
 
diff --git a/compiler/rustc_const_eval/src/transform/validate.rs b/compiler/rustc_const_eval/src/transform/validate.rs
index 068491646f4..785dde93e31 100644
--- a/compiler/rustc_const_eval/src/transform/validate.rs
+++ b/compiler/rustc_const_eval/src/transform/validate.rs
@@ -553,15 +553,6 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                             );
                         }
                     }
-                    Shl | Shr => {
-                        for x in [a, b] {
-                            check_kinds!(
-                                x,
-                                "Cannot perform checked shift on non-integer type {:?}",
-                                ty::Uint(..) | ty::Int(..)
-                            )
-                        }
-                    }
                     _ => self.fail(location, format!("There is no checked version of {:?}", op)),
                 }
             }
diff --git a/compiler/rustc_middle/src/ty/util.rs b/compiler/rustc_middle/src/ty/util.rs
index 90270e0ee9d..754b6efd692 100644
--- a/compiler/rustc_middle/src/ty/util.rs
+++ b/compiler/rustc_middle/src/ty/util.rs
@@ -59,22 +59,13 @@ impl<'tcx> fmt::Display for Discr<'tcx> {
     }
 }
 
-fn int_size_and_signed<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty<'tcx>) -> (Size, bool) {
-    let (int, signed) = match *ty.kind() {
-        ty::Int(ity) => (Integer::from_int_ty(&tcx, ity), true),
-        ty::Uint(uty) => (Integer::from_uint_ty(&tcx, uty), false),
-        _ => bug!("non integer discriminant"),
-    };
-    (int.size(), signed)
-}
-
 impl<'tcx> Discr<'tcx> {
     /// Adds `1` to the value and wraps around if the maximum for the type is reached.
     pub fn wrap_incr(self, tcx: TyCtxt<'tcx>) -> Self {
         self.checked_add(tcx, 1).0
     }
     pub fn checked_add(self, tcx: TyCtxt<'tcx>, n: u128) -> (Self, bool) {
-        let (size, signed) = int_size_and_signed(tcx, self.ty);
+        let (size, signed) = self.ty.int_size_and_signed(tcx);
         let (val, oflo) = if signed {
             let min = size.signed_int_min();
             let max = size.signed_int_max();
@@ -922,12 +913,21 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for OpaqueTypeExpander<'tcx> {
 }
 
 impl<'tcx> Ty<'tcx> {
+    pub fn int_size_and_signed(self, tcx: TyCtxt<'tcx>) -> (Size, bool) {
+        let (int, signed) = match *self.kind() {
+            ty::Int(ity) => (Integer::from_int_ty(&tcx, ity), true),
+            ty::Uint(uty) => (Integer::from_uint_ty(&tcx, uty), false),
+            _ => bug!("non integer discriminant"),
+        };
+        (int.size(), signed)
+    }
+
     /// Returns the maximum value for the given numeric type (including `char`s)
     /// or returns `None` if the type is not numeric.
     pub fn numeric_max_val(self, tcx: TyCtxt<'tcx>) -> Option<ty::Const<'tcx>> {
         let val = match self.kind() {
             ty::Int(_) | ty::Uint(_) => {
-                let (size, signed) = int_size_and_signed(tcx, self);
+                let (size, signed) = self.int_size_and_signed(tcx);
                 let val =
                     if signed { size.signed_int_max() as u128 } else { size.unsigned_int_max() };
                 Some(val)
@@ -948,7 +948,7 @@ impl<'tcx> Ty<'tcx> {
     pub fn numeric_min_val(self, tcx: TyCtxt<'tcx>) -> Option<ty::Const<'tcx>> {
         let val = match self.kind() {
             ty::Int(_) | ty::Uint(_) => {
-                let (size, signed) = int_size_and_signed(tcx, self);
+                let (size, signed) = self.int_size_and_signed(tcx);
                 let val = if signed { size.truncate(size.signed_int_min() as u128) } else { 0 };
                 Some(val)
             }
diff --git a/compiler/rustc_mir_build/src/build/expr/as_rvalue.rs b/compiler/rustc_mir_build/src/build/expr/as_rvalue.rs
index a4e48c1545d..259e0f8f75c 100644
--- a/compiler/rustc_mir_build/src/build/expr/as_rvalue.rs
+++ b/compiler/rustc_mir_build/src/build/expr/as_rvalue.rs
@@ -9,6 +9,7 @@ use crate::build::expr::category::{Category, RvalueFunc};
 use crate::build::{BlockAnd, BlockAndExtension, Builder, NeedsTemporary};
 use rustc_hir::lang_items::LangItem;
 use rustc_middle::middle::region;
+use rustc_middle::mir::interpret::Scalar;
 use rustc_middle::mir::AssertKind;
 use rustc_middle::mir::Place;
 use rustc_middle::mir::*;
@@ -519,30 +520,68 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     ) -> BlockAnd<Rvalue<'tcx>> {
         let source_info = self.source_info(span);
         let bool_ty = self.tcx.types.bool;
-        if self.check_overflow && op.is_checkable() && ty.is_integral() {
-            let result_tup = self.tcx.mk_tup(&[ty, bool_ty]);
-            let result_value = self.temp(result_tup, span);
+        let rvalue = match op {
+            BinOp::Add | BinOp::Sub | BinOp::Mul if self.check_overflow && ty.is_integral() => {
+                let result_tup = self.tcx.mk_tup(&[ty, bool_ty]);
+                let result_value = self.temp(result_tup, span);
 
-            self.cfg.push_assign(
-                block,
-                source_info,
-                result_value,
-                Rvalue::CheckedBinaryOp(op, Box::new((lhs.to_copy(), rhs.to_copy()))),
-            );
-            let val_fld = Field::new(0);
-            let of_fld = Field::new(1);
+                self.cfg.push_assign(
+                    block,
+                    source_info,
+                    result_value,
+                    Rvalue::CheckedBinaryOp(op, Box::new((lhs.to_copy(), rhs.to_copy()))),
+                );
+                let val_fld = Field::new(0);
+                let of_fld = Field::new(1);
+
+                let tcx = self.tcx;
+                let val = tcx.mk_place_field(result_value, val_fld, ty);
+                let of = tcx.mk_place_field(result_value, of_fld, bool_ty);
 
-            let tcx = self.tcx;
-            let val = tcx.mk_place_field(result_value, val_fld, ty);
-            let of = tcx.mk_place_field(result_value, of_fld, bool_ty);
+                let err = AssertKind::Overflow(op, lhs, rhs);
+                block = self.assert(block, Operand::Move(of), false, err, span);
 
-            let err = AssertKind::Overflow(op, lhs, rhs);
+                Rvalue::Use(Operand::Move(val))
+            }
+            BinOp::Shl | BinOp::Shr if self.check_overflow && ty.is_integral() => {
+                // Consider that the shift overflows if `rhs < 0` or `rhs >= bits`.
+                // This can be encoded as a single operation as `(rhs & -bits) != 0`.
+                let (size, _) = ty.int_size_and_signed(self.tcx);
+                let bits = size.bits();
+                debug_assert!(bits.is_power_of_two());
+                let mask = !((bits - 1) as u128);
+
+                let rhs_ty = rhs.ty(&self.local_decls, self.tcx);
+                let (rhs_size, _) = rhs_ty.int_size_and_signed(self.tcx);
+                let mask = Operand::const_from_scalar(
+                    self.tcx,
+                    rhs_ty,
+                    Scalar::from_uint(rhs_size.truncate(mask), rhs_size),
+                    span,
+                );
 
-            block = self.assert(block, Operand::Move(of), false, err, span);
+                let outer_bits = self.temp(rhs_ty, span);
+                self.cfg.push_assign(
+                    block,
+                    source_info,
+                    outer_bits,
+                    Rvalue::BinaryOp(BinOp::BitAnd, Box::new((rhs.to_copy(), mask))),
+                );
 
-            block.and(Rvalue::Use(Operand::Move(val)))
-        } else {
-            if ty.is_integral() && (op == BinOp::Div || op == BinOp::Rem) {
+                let overflows = self.temp(bool_ty, span);
+                let zero = self.zero_literal(span, rhs_ty);
+                self.cfg.push_assign(
+                    block,
+                    source_info,
+                    overflows,
+                    Rvalue::BinaryOp(BinOp::Ne, Box::new((Operand::Move(outer_bits), zero))),
+                );
+
+                let overflow_err = AssertKind::Overflow(op, lhs.to_copy(), rhs.to_copy());
+                block = self.assert(block, Operand::Move(overflows), false, overflow_err, span);
+                Rvalue::BinaryOp(op, Box::new((lhs, rhs)))
+            }
+            BinOp::Div | BinOp::Rem if ty.is_integral() => {
                 // Checking division and remainder is more complex, since we 1. always check
                 // and 2. there are two possible failure cases, divide-by-zero and overflow.
 
@@ -601,10 +640,12 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
 
                     block = self.assert(block, Operand::Move(of), false, overflow_err, span);
                 }
-            }
 
-            block.and(Rvalue::BinaryOp(op, Box::new((lhs, rhs))))
-        }
+                Rvalue::BinaryOp(op, Box::new((lhs, rhs)))
+            }
+            _ => Rvalue::BinaryOp(op, Box::new((lhs, rhs))),
+        };
+        block.and(rvalue)
     }
 
     fn build_zero_repeat(
diff --git a/tests/mir-opt/issue_101973.inner.ConstProp.diff b/tests/mir-opt/issue_101973.inner.ConstProp.diff
index 6db8e4d2664..fb0b3866e69 100644
--- a/tests/mir-opt/issue_101973.inner.ConstProp.diff
+++ b/tests/mir-opt/issue_101973.inner.ConstProp.diff
@@ -12,13 +12,14 @@
       let mut _7: u32;                     // in scope 0 at $DIR/issue_101973.rs:+1:31: +1:52
       let mut _8: u32;                     // in scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
       let mut _9: u32;                     // in scope 0 at $DIR/issue_101973.rs:+1:33: +1:39
-      let mut _10: (u32, bool);            // in scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
-      let mut _11: (u32, bool);            // in scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
+      let mut _10: i32;                    // in scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
+      let mut _11: bool;                   // in scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
+      let mut _12: i32;                    // in scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
+      let mut _13: bool;                   // in scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
       scope 1 (inlined imm8) {             // at $DIR/issue_101973.rs:14:5: 14:17
           debug x => _1;                   // in scope 1 at $DIR/issue_101973.rs:5:13: 5:14
-          let mut _12: u32;                // in scope 1 at $DIR/issue_101973.rs:7:12: 7:27
-          let mut _13: u32;                // in scope 1 at $DIR/issue_101973.rs:7:12: 7:20
-          let mut _14: (u32, bool);        // in scope 1 at $DIR/issue_101973.rs:7:12: 7:20
+          let mut _14: u32;                // in scope 1 at $DIR/issue_101973.rs:7:12: 7:27
+          let mut _15: u32;                // in scope 1 at $DIR/issue_101973.rs:7:12: 7:20
           scope 2 {
               debug out => _4;             // in scope 2 at $DIR/issue_101973.rs:6:9: 6:16
           }
@@ -32,43 +33,46 @@
           StorageLive(_2);                 // scope 0 at $DIR/issue_101973.rs:+1:5: +1:65
           StorageLive(_3);                 // scope 0 at $DIR/issue_101973.rs:+1:5: +1:58
           StorageLive(_4);                 // scope 0 at $DIR/issue_101973.rs:+1:5: +1:17
-          StorageLive(_12);                // scope 2 at $DIR/issue_101973.rs:7:12: 7:27
-          StorageLive(_13);                // scope 2 at $DIR/issue_101973.rs:7:12: 7:20
-          _14 = CheckedShr(_1, const 0_i32); // scope 2 at $DIR/issue_101973.rs:7:12: 7:20
-          assert(!move (_14.1: bool), "attempt to shift right by `{}`, which would overflow", const 0_i32) -> bb3; // scope 2 at $DIR/issue_101973.rs:7:12: 7:20
+          StorageLive(_14);                // scope 2 at $DIR/issue_101973.rs:7:12: 7:27
+          StorageLive(_15);                // scope 2 at $DIR/issue_101973.rs:7:12: 7:20
+          _15 = Shr(_1, const 0_i32);      // scope 2 at $DIR/issue_101973.rs:7:12: 7:20
+          _14 = BitAnd(move _15, const 255_u32); // scope 2 at $DIR/issue_101973.rs:7:12: 7:27
+          StorageDead(_15);                // scope 2 at $DIR/issue_101973.rs:7:26: 7:27
+          _4 = BitOr(const 0_u32, move _14); // scope 2 at $DIR/issue_101973.rs:7:5: 7:27
+          StorageDead(_14);                // scope 2 at $DIR/issue_101973.rs:7:26: 7:27
+          StorageLive(_6);                 // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
+          StorageLive(_7);                 // scope 0 at $DIR/issue_101973.rs:+1:31: +1:52
+          StorageLive(_8);                 // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
+-         _10 = BitAnd(const 8_i32, const -32_i32); // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
+-         _11 = Ne(move _10, const 0_i32); // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
+-         assert(!move _11, "attempt to shift right by `{}`, which would overflow", const 8_i32) -> bb1; // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
++         _10 = const 0_i32;               // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
++         _11 = const false;               // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
++         assert(!const false, "attempt to shift right by `{}`, which would overflow", const 8_i32) -> bb1; // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
       }
   
       bb1: {
-          _8 = move (_10.0: u32);          // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
+          _8 = Shr(_1, const 8_i32);       // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
           _7 = BitAnd(move _8, const 15_u32); // scope 0 at $DIR/issue_101973.rs:+1:31: +1:52
           StorageDead(_8);                 // scope 0 at $DIR/issue_101973.rs:+1:51: +1:52
-          _11 = CheckedShl(_7, const 1_i32); // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
-          assert(!move (_11.1: bool), "attempt to shift left by `{}`, which would overflow", const 1_i32) -> bb2; // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
+-         _12 = BitAnd(const 1_i32, const -32_i32); // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
+-         _13 = Ne(move _12, const 0_i32); // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
+-         assert(!move _13, "attempt to shift left by `{}`, which would overflow", const 1_i32) -> bb2; // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
++         _12 = const 0_i32;               // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
++         _13 = const false;               // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
++         assert(!const false, "attempt to shift left by `{}`, which would overflow", const 1_i32) -> bb2; // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
       }
   
       bb2: {
-          _6 = move (_11.0: u32);          // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
+          _6 = Shl(move _7, const 1_i32);  // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
           StorageDead(_7);                 // scope 0 at $DIR/issue_101973.rs:+1:56: +1:57
-          _3 = rotate_right::<u32>(_4, _6) -> bb4; // scope 3 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
+          _3 = rotate_right::<u32>(_4, _6) -> bb3; // scope 3 at $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
                                            // mir::Constant
                                            // + span: $SRC_DIR/core/src/num/uint_macros.rs:LL:COL
                                            // + literal: Const { ty: extern "rust-intrinsic" fn(u32, u32) -> u32 {rotate_right::<u32>}, val: Value(<ZST>) }
       }
   
       bb3: {
-          _13 = move (_14.0: u32);         // scope 2 at $DIR/issue_101973.rs:7:12: 7:20
-          _12 = BitAnd(move _13, const 255_u32); // scope 2 at $DIR/issue_101973.rs:7:12: 7:27
-          StorageDead(_13);                // scope 2 at $DIR/issue_101973.rs:7:26: 7:27
-          _4 = BitOr(const 0_u32, move _12); // scope 2 at $DIR/issue_101973.rs:7:5: 7:27
-          StorageDead(_12);                // scope 2 at $DIR/issue_101973.rs:7:26: 7:27
-          StorageLive(_6);                 // scope 0 at $DIR/issue_101973.rs:+1:31: +1:57
-          StorageLive(_7);                 // scope 0 at $DIR/issue_101973.rs:+1:31: +1:52
-          StorageLive(_8);                 // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
-          _10 = CheckedShr(_1, const 8_i32); // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
-          assert(!move (_10.1: bool), "attempt to shift right by `{}`, which would overflow", const 8_i32) -> bb1; // scope 0 at $DIR/issue_101973.rs:+1:32: +1:45
-      }
-  
-      bb4: {
           StorageDead(_6);                 // scope 0 at $DIR/issue_101973.rs:+1:57: +1:58
           StorageDead(_4);                 // scope 0 at $DIR/issue_101973.rs:+1:57: +1:58
           _2 = move _3 as i32 (IntToInt);  // scope 0 at $DIR/issue_101973.rs:+1:5: +1:65
diff --git a/tests/ui/consts/offset_from_ub.stderr b/tests/ui/consts/offset_from_ub.stderr
index fff4729689f..6530084a585 100644
--- a/tests/ui/consts/offset_from_ub.stderr
+++ b/tests/ui/consts/offset_from_ub.stderr
@@ -39,19 +39,19 @@ error[E0080]: evaluation of constant value failed
   --> $DIR/offset_from_ub.rs:53:14
    |
 LL |     unsafe { ptr_offset_from(end_ptr, start_ptr) }
-   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ out-of-bounds offset_from: alloc18 has size 4, so pointer to 10 bytes starting at offset 0 is out-of-bounds
+   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ out-of-bounds offset_from: alloc17 has size 4, so pointer to 10 bytes starting at offset 0 is out-of-bounds
 
 error[E0080]: evaluation of constant value failed
   --> $DIR/offset_from_ub.rs:62:14
    |
 LL |     unsafe { ptr_offset_from(start_ptr, end_ptr) }
-   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ out-of-bounds offset_from: alloc21 has size 4, so pointer to 10 bytes starting at offset 0 is out-of-bounds
+   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ out-of-bounds offset_from: alloc20 has size 4, so pointer to 10 bytes starting at offset 0 is out-of-bounds
 
 error[E0080]: evaluation of constant value failed
   --> $DIR/offset_from_ub.rs:70:14
    |
 LL |     unsafe { ptr_offset_from(end_ptr, end_ptr) }
-   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ out-of-bounds offset_from: alloc24 has size 4, so pointer at offset 10 is out-of-bounds
+   |              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ out-of-bounds offset_from: alloc23 has size 4, so pointer at offset 10 is out-of-bounds
 
 error[E0080]: evaluation of constant value failed
   --> $DIR/offset_from_ub.rs:79:14