De-LLVM the unchecked shifts [MCP#693]

This is just one part of the MCP, but it's the one that IMHO removes the most noise from the standard library code.

Seems net simpler this way, since MIR already supported heterogeneous shifts anyway, and thus it's not more work for backends than before.

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,32 @@ pub fn coerce_unsized_into<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     }
 }
 
-pub fn cast_shift_expr_rhs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
+/// Shifts in MIR are all allowed to have mismatched LHS & RHS types.
+///
+/// This does all the appropriate conversions needed to pass it to the builder's
+/// shift methods, which are UB for out-of-range shifts.
+///
+/// If `is_unchecked` is false, this masks the RHS to ensure it stays in-bounds.
+/// 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 +335,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)
+}