De-LLVM the unchecked shifts [MCP#693]

compiler/rustc_codegen_ssa/src/base.rs

@@ -5,7 +5,7 @@ use crate::back::write::{
     compute_per_cgu_lto_type, start_async_codegen, submit_codegened_module_to_llvm,
     submit_post_lto_module_to_llvm, submit_pre_lto_module_to_llvm, ComputedLtoType, OngoingCodegen,
 };
-use crate::common::{IntPredicate, RealPredicate, TypeKind};
+use crate::common::{self, IntPredicate, RealPredicate, TypeKind};
 use crate::errors;
 use crate::meth;
 use crate::mir;
@@ -33,7 +33,7 @@ use rustc_middle::mir::mono::{CodegenUnit, CodegenUnitNameBuilder, MonoItem};
 use rustc_middle::query::Providers;
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, TyAndLayout};
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
-use rustc_session::config::{self, CrateType, EntryFnType, OutputType};
+use rustc_session::config::{self, CrateType, EntryFnType, OptLevel, OutputType};
 use rustc_session::Session;
 use rustc_span::symbol::sym;
 use rustc_span::Symbol;
@@ -300,14 +300,35 @@ pub fn coerce_unsized_into<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     }
 }
 
-pub fn cast_shift_expr_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
+/// Returns `rhs` sufficiently masked, truncated, and/or extended so that
+/// it can be used to shift `lhs`.
+///
+/// Shifts in MIR are all allowed to have mismatched LHS & RHS types.
+/// The shift methods in `BuilderMethods`, however, are fully homogeneous
+/// (both parameters and the return type are all the same type).
+///
+/// If `is_unchecked` is false, this masks the RHS to ensure it stays in-bounds,
+/// as the `BuilderMethods` shifts are UB for out-of-bounds shift amounts.
+/// For 32- and 64-bit types, this matches the semantics
+/// of Java. (See related discussion on #1877 and #10183.)
+///
+/// If `is_unchecked` is true, this does no masking, and adds sufficient `assume`
+/// calls or operation flags to preserve as much freedom to optimize as possible.
+pub fn build_shift_expr_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     bx: &mut Bx,
     lhs: Bx::Value,
-    rhs: Bx::Value,
+    mut rhs: Bx::Value,
+    is_unchecked: bool,
 ) -> Bx::Value {
     // Shifts may have any size int on the rhs
     let mut rhs_llty = bx.cx().val_ty(rhs);
     let mut lhs_llty = bx.cx().val_ty(lhs);
+
+    let mask = common::shift_mask_val(bx, lhs_llty, rhs_llty, false);
+    if !is_unchecked {
+        rhs = bx.and(rhs, mask);
+    }
+
     if bx.cx().type_kind(rhs_llty) == TypeKind::Vector {
         rhs_llty = bx.cx().element_type(rhs_llty)
     }
@@ -317,6 +338,12 @@ pub fn cast_shift_expr_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     let rhs_sz = bx.cx().int_width(rhs_llty);
     let lhs_sz = bx.cx().int_width(lhs_llty);
     if lhs_sz < rhs_sz {
+        if is_unchecked && bx.sess().opts.optimize != OptLevel::No {
+            // FIXME: Use `trunc nuw` once that's available
+            let inrange = bx.icmp(IntPredicate::IntULE, rhs, mask);
+            bx.assume(inrange);
+        }
+
         bx.trunc(rhs, lhs_llty)
     } else if lhs_sz > rhs_sz {
         // We zero-extend even if the RHS is signed. So e.g. `(x: i32) << -1i8` will zero-extend the

compiler/rustc_codegen_ssa/src/common.rs

@@ -3,10 +3,9 @@
 use rustc_hir::LangItem;
 use rustc_middle::mir;
 use rustc_middle::ty::Instance;
-use rustc_middle::ty::{self, layout::TyAndLayout, Ty, TyCtxt};
+use rustc_middle::ty::{self, layout::TyAndLayout, TyCtxt};
 use rustc_span::Span;
 
-use crate::base;
 use crate::traits::*;
 
 #[derive(Copy, Clone)]
@@ -128,44 +127,6 @@ pub fn build_langcall<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     (bx.fn_abi_of_instance(instance, ty::List::empty()), bx.get_fn_addr(instance), instance)
 }
 
-// To avoid UB from LLVM, these two functions mask RHS with an
-// appropriate mask unconditionally (i.e., the fallback behavior for
-// all shifts). For 32- and 64-bit types, this matches the semantics
-// of Java. (See related discussion on #1877 and #10183.)
-
-pub fn build_masked_lshift<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
-    bx: &mut Bx,
-    lhs: Bx::Value,
-    rhs: Bx::Value,
-) -> Bx::Value {
-    let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
-    // #1877, #10183: Ensure that input is always valid
-    let rhs = shift_mask_rhs(bx, rhs);
-    bx.shl(lhs, rhs)
-}
-
-pub fn build_masked_rshift<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
-    bx: &mut Bx,
-    lhs_t: Ty<'tcx>,
-    lhs: Bx::Value,
-    rhs: Bx::Value,
-) -> Bx::Value {
-    let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
-    // #1877, #10183: Ensure that input is always valid
-    let rhs = shift_mask_rhs(bx, rhs);
-    let is_signed = lhs_t.is_signed();
-    if is_signed { bx.ashr(lhs, rhs) } else { bx.lshr(lhs, rhs) }
-}
-
-fn shift_mask_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
-    bx: &mut Bx,
-    rhs: Bx::Value,
-) -> Bx::Value {
-    let rhs_llty = bx.val_ty(rhs);
-    let shift_val = shift_mask_val(bx, rhs_llty, rhs_llty, false);
-    bx.and(rhs, shift_val)
-}
-
 pub fn shift_mask_val<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     bx: &mut Bx,
     llty: Bx::Type,

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -3,7 +3,7 @@ use super::place::PlaceRef;
 use super::{FunctionCx, LocalRef};
 
 use crate::base;
-use crate::common::{self, IntPredicate};
+use crate::common::IntPredicate;
 use crate::traits::*;
 use crate::MemFlags;
 
@@ -860,14 +860,12 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     bx.inbounds_gep(llty, lhs, &[rhs])
                 }
             }
-            mir::BinOp::Shl => common::build_masked_lshift(bx, lhs, rhs),
-            mir::BinOp::ShlUnchecked => {
-                let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
+            mir::BinOp::Shl | mir::BinOp::ShlUnchecked => {
+                let rhs = base::build_shift_expr_rhs(bx, lhs, rhs, op == mir::BinOp::ShlUnchecked);
                 bx.shl(lhs, rhs)
             }
-            mir::BinOp::Shr => common::build_masked_rshift(bx, input_ty, lhs, rhs),
-            mir::BinOp::ShrUnchecked => {
-                let rhs = base::cast_shift_expr_rhs(bx, lhs, rhs);
+            mir::BinOp::Shr | mir::BinOp::ShrUnchecked => {
+                let rhs = base::build_shift_expr_rhs(bx, lhs, rhs, op == mir::BinOp::ShrUnchecked);
                 if is_signed { bx.ashr(lhs, rhs) } else { bx.lshr(lhs, rhs) }
             }
             mir::BinOp::Ne

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -454,9 +454,8 @@ pub fn check_intrinsic_type(
             sym::unchecked_div | sym::unchecked_rem | sym::exact_div => {
                 (1, 0, vec![param(0), param(0)], param(0))
             }
-            sym::unchecked_shl | sym::unchecked_shr | sym::rotate_left | sym::rotate_right => {
-                (1, 0, vec![param(0), param(0)], param(0))
-            }
+            sym::unchecked_shl | sym::unchecked_shr => (2, 0, vec![param(0), param(1)], param(0)),
+            sym::rotate_left | sym::rotate_right => (1, 0, vec![param(0), param(0)], param(0)),
             sym::unchecked_add | sym::unchecked_sub | sym::unchecked_mul => {
                 (1, 0, vec![param(0), param(0)], param(0))
             }

library/core/src/intrinsics.rs

@@ -2224,18 +2224,20 @@ extern "rust-intrinsic" {
     /// Safe wrappers for this intrinsic are available on the integer
     /// primitives via the `checked_shl` method. For example,
     /// [`u32::checked_shl`]
+    #[cfg(not(bootstrap))]
     #[rustc_const_stable(feature = "const_int_unchecked", since = "1.40.0")]
     #[rustc_nounwind]
-    pub fn unchecked_shl<T: Copy>(x: T, y: T) -> T;
+    pub fn unchecked_shl<T: Copy, U: Copy>(x: T, y: U) -> T;
     /// Performs an unchecked right shift, resulting in undefined behavior when
     /// `y < 0` or `y >= N`, where N is the width of T in bits.
     ///
     /// Safe wrappers for this intrinsic are available on the integer
     /// primitives via the `checked_shr` method. For example,
     /// [`u32::checked_shr`]
+    #[cfg(not(bootstrap))]
     #[rustc_const_stable(feature = "const_int_unchecked", since = "1.40.0")]
     #[rustc_nounwind]
-    pub fn unchecked_shr<T: Copy>(x: T, y: T) -> T;
+    pub fn unchecked_shr<T: Copy, U: Copy>(x: T, y: U) -> T;
 
     /// Returns the result of an unchecked addition, resulting in
     /// undefined behavior when `x + y > T::MAX` or `x + y < T::MIN`.

library/core/src/num/int_macros.rs

@@ -1253,10 +1253,18 @@ macro_rules! int_impl {
         #[inline(always)]
         #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
         pub const unsafe fn unchecked_shl(self, rhs: u32) -> Self {
-            // SAFETY: the caller must uphold the safety contract for
-            // `unchecked_shl`.
-            // Any legal shift amount is losslessly representable in the self type.
-            unsafe { intrinsics::unchecked_shl(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
+            #[cfg(bootstrap)]
+            {
+                // For bootstrapping, just use built-in primitive shift.
+                // panicking is a legal manifestation of UB
+                self << rhs
+            }
+            #[cfg(not(bootstrap))]
+            {
+                // SAFETY: the caller must uphold the safety contract for
+                // `unchecked_shl`.
+                unsafe { intrinsics::unchecked_shl(self, rhs) }
+            }
         }
 
         /// Checked shift right. Computes `self >> rhs`, returning `None` if `rhs` is
@@ -1336,10 +1344,18 @@ macro_rules! int_impl {
         #[inline(always)]
         #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
         pub const unsafe fn unchecked_shr(self, rhs: u32) -> Self {
-            // SAFETY: the caller must uphold the safety contract for
-            // `unchecked_shr`.
-            // Any legal shift amount is losslessly representable in the self type.
-            unsafe { intrinsics::unchecked_shr(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
+            #[cfg(bootstrap)]
+            {
+                // For bootstrapping, just use built-in primitive shift.
+                // panicking is a legal manifestation of UB
+                self >> rhs
+            }
+            #[cfg(not(bootstrap))]
+            {
+                // SAFETY: the caller must uphold the safety contract for
+                // `unchecked_shr`.
+                unsafe { intrinsics::unchecked_shr(self, rhs) }
+            }
         }
 
         /// Checked absolute value. Computes `self.abs()`, returning `None` if

library/core/src/num/mod.rs

@@ -286,17 +286,6 @@ macro_rules! widening_impl {
     };
 }
 
-macro_rules! conv_rhs_for_unchecked_shift {
-    ($SelfT:ty, $x:expr) => {{
-        // If the `as` cast will truncate, ensure we still tell the backend
-        // that the pre-truncation value was also small.
-        if <$SelfT>::BITS < 32 {
-            intrinsics::assume($x <= (<$SelfT>::MAX as u32));
-        }
-        $x as $SelfT
-    }};
-}
-
 impl i8 {
     int_impl! {
         Self = i8,

library/core/src/num/uint_macros.rs

@@ -1313,10 +1313,18 @@ macro_rules! uint_impl {
         #[inline(always)]
         #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
         pub const unsafe fn unchecked_shl(self, rhs: u32) -> Self {
-            // SAFETY: the caller must uphold the safety contract for
-            // `unchecked_shl`.
-            // Any legal shift amount is losslessly representable in the self type.
-            unsafe { intrinsics::unchecked_shl(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
+            #[cfg(bootstrap)]
+            {
+                // For bootstrapping, just use built-in primitive shift.
+                // panicking is a legal manifestation of UB
+                self << rhs
+            }
+            #[cfg(not(bootstrap))]
+            {
+                // SAFETY: the caller must uphold the safety contract for
+                // `unchecked_shl`.
+                unsafe { intrinsics::unchecked_shl(self, rhs) }
+            }
         }
 
         /// Checked shift right. Computes `self >> rhs`, returning `None`
@@ -1396,10 +1404,18 @@ macro_rules! uint_impl {
         #[inline(always)]
         #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
         pub const unsafe fn unchecked_shr(self, rhs: u32) -> Self {
-            // SAFETY: the caller must uphold the safety contract for
-            // `unchecked_shr`.
-            // Any legal shift amount is losslessly representable in the self type.
-            unsafe { intrinsics::unchecked_shr(self, conv_rhs_for_unchecked_shift!($SelfT, rhs)) }
+            #[cfg(bootstrap)]
+            {
+                // For bootstrapping, just use built-in primitive shift.
+                // panicking is a legal manifestation of UB
+                self >> rhs
+            }
+            #[cfg(not(bootstrap))]
+            {
+                // SAFETY: the caller must uphold the safety contract for
+                // `unchecked_shr`.
+                unsafe { intrinsics::unchecked_shr(self, rhs) }
+            }
         }
 
         /// Checked exponentiation. Computes `self.pow(exp)`, returning `None` if

library/core/src/ptr/mod.rs

@@ -1781,9 +1781,19 @@ pub(crate) const unsafe fn align_offset<T: Sized>(p: *const T, a: usize) -> usiz
     // FIXME(#75598): Direct use of these intrinsics improves codegen significantly at opt-level <=
     // 1, where the method versions of these operations are not inlined.
     use intrinsics::{
-        assume, cttz_nonzero, exact_div, mul_with_overflow, unchecked_rem, unchecked_shl,
-        unchecked_shr, unchecked_sub, wrapping_add, wrapping_mul, wrapping_sub,
+        assume, cttz_nonzero, exact_div, mul_with_overflow, unchecked_rem, unchecked_sub,
+        wrapping_add, wrapping_mul, wrapping_sub,
     };
+    #[cfg(bootstrap)]
+    const unsafe fn unchecked_shl(value: usize, shift: usize) -> usize {
+        value << shift
+    }
+    #[cfg(bootstrap)]
+    const unsafe fn unchecked_shr(value: usize, shift: usize) -> usize {
+        value >> shift
+    }
+    #[cfg(not(bootstrap))]
+    use intrinsics::{unchecked_shl, unchecked_shr};
 
     /// Calculate multiplicative modular inverse of `x` modulo `m`.
     ///

tests/codegen/unchecked_shifts.rs

@@ -2,6 +2,7 @@
 
 #![crate_type = "lib"]
 #![feature(unchecked_shifts)]
+#![feature(core_intrinsics)]
 
 // CHECK-LABEL: @unchecked_shl_unsigned_same
 #[no_mangle]
@@ -19,7 +20,7 @@ pub unsafe fn unchecked_shl_unsigned_smaller(a: u16, b: u32) -> u16 {
     // This uses -DAG to avoid failing on irrelevant reorderings,
     // like emitting the truncation earlier.
 
-    // CHECK-DAG: %[[INRANGE:.+]] = icmp ult i32 %b, 65536
+    // CHECK-DAG: %[[INRANGE:.+]] = icmp ult i32 %b, 16
     // CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
     // CHECK-DAG: %[[TRUNC:.+]] = trunc i32 %b to i16
     // CHECK-DAG: shl i16 %a, %[[TRUNC]]
@@ -51,7 +52,7 @@ pub unsafe fn unchecked_shr_signed_smaller(a: i16, b: u32) -> i16 {
     // This uses -DAG to avoid failing on irrelevant reorderings,
     // like emitting the truncation earlier.
 
-    // CHECK-DAG: %[[INRANGE:.+]] = icmp ult i32 %b, 32768
+    // CHECK-DAG: %[[INRANGE:.+]] = icmp ult i32 %b, 16
     // CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
     // CHECK-DAG: %[[TRUNC:.+]] = trunc i32 %b to i16
     // CHECK-DAG: ashr i16 %a, %[[TRUNC]]
@@ -66,3 +67,47 @@ pub unsafe fn unchecked_shr_signed_bigger(a: i64, b: u32) -> i64 {
     // CHECK: ashr i64 %a, %[[EXT]]
     a.unchecked_shr(b)
 }
+
+// CHECK-LABEL: @unchecked_shr_u128_i8
+#[no_mangle]
+pub unsafe fn unchecked_shr_u128_i8(a: u128, b: i8) -> u128 {
+    // CHECK-NOT: assume
+    // CHECK: %[[EXT:.+]] = zext{{( nneg)?}} i8 %b to i128
+    // CHECK: lshr i128 %a, %[[EXT]]
+    std::intrinsics::unchecked_shr(a, b)
+}
+
+// CHECK-LABEL: @unchecked_shl_i128_u8
+#[no_mangle]
+pub unsafe fn unchecked_shl_i128_u8(a: i128, b: u8) -> i128 {
+    // CHECK-NOT: assume
+    // CHECK: %[[EXT:.+]] = zext{{( nneg)?}} i8 %b to i128
+    // CHECK: shl i128 %a, %[[EXT]]
+    std::intrinsics::unchecked_shl(a, b)
+}
+
+// CHECK-LABEL: @unchecked_shl_u8_i128
+#[no_mangle]
+pub unsafe fn unchecked_shl_u8_i128(a: u8, b: i128) -> u8 {
+    // This uses -DAG to avoid failing on irrelevant reorderings,
+    // like emitting the truncation earlier.
+
+    // CHECK-DAG: %[[INRANGE:.+]] = icmp ult i128 %b, 8
+    // CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
+    // CHECK-DAG: %[[TRUNC:.+]] = trunc i128 %b to i8
+    // CHECK-DAG: shl i8 %a, %[[TRUNC]]
+    std::intrinsics::unchecked_shl(a, b)
+}
+
+// CHECK-LABEL: @unchecked_shr_i8_u128
+#[no_mangle]
+pub unsafe fn unchecked_shr_i8_u128(a: i8, b: u128) -> i8 {
+    // This uses -DAG to avoid failing on irrelevant reorderings,
+    // like emitting the truncation earlier.
+
+    // CHECK-DAG: %[[INRANGE:.+]] = icmp ult i128 %b, 8
+    // CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
+    // CHECK-DAG: %[[TRUNC:.+]] = trunc i128 %b to i8
+    // CHECK-DAG: ashr i8 %a, %[[TRUNC]]
+    std::intrinsics::unchecked_shr(a, b)
+}