Emit `trunc nuw` for unchecked shifts and `to_immediate_scalar`

- For shifts this shrinks the IR by no longer needing an `assume` while still providing the UB information
- Having this on the `i8`→`i1` truncations will hopefully help with some places that have to load `i8`s or pass those in LLVM structs without range information

10 files changed, 77 insertions, 50 deletions

diff --git a/compiler/rustc_codegen_gcc/src/builder.rs b/compiler/rustc_codegen_gcc/src/builder.rs
index b8e37b60480..76846692459 100644
--- a/compiler/rustc_codegen_gcc/src/builder.rs
+++ b/compiler/rustc_codegen_gcc/src/builder.rs
@@ -1694,7 +1694,7 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
 
     fn to_immediate_scalar(&mut self, val: Self::Value, scalar: abi::Scalar) -> Self::Value {
         if scalar.is_bool() {
-            return self.trunc(val, self.cx().type_i1());
+            return self.unchecked_utrunc(val, self.cx().type_i1());
         }
         val
     }
diff --git a/compiler/rustc_codegen_llvm/src/builder.rs b/compiler/rustc_codegen_llvm/src/builder.rs
index e10a4d63407..18bb6a609aa 100644
--- a/compiler/rustc_codegen_llvm/src/builder.rs
+++ b/compiler/rustc_codegen_llvm/src/builder.rs
@@ -29,13 +29,13 @@ use smallvec::SmallVec;
 use tracing::{debug, instrument};
 
 use crate::abi::FnAbiLlvmExt;
-use crate::attributes;
 use crate::common::Funclet;
 use crate::context::{CodegenCx, SimpleCx};
 use crate::llvm::{self, AtomicOrdering, AtomicRmwBinOp, BasicBlock, False, Metadata, True};
 use crate::type_::Type;
 use crate::type_of::LayoutLlvmExt;
 use crate::value::Value;
+use crate::{attributes, llvm_util};
 
 #[must_use]
 pub(crate) struct GenericBuilder<'a, 'll, CX: Borrow<SimpleCx<'ll>>> {
@@ -606,7 +606,7 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
 
     fn to_immediate_scalar(&mut self, val: Self::Value, scalar: abi::Scalar) -> Self::Value {
         if scalar.is_bool() {
-            return self.trunc(val, self.cx().type_i1());
+            return self.unchecked_utrunc(val, self.cx().type_i1());
         }
         val
     }
@@ -942,6 +942,30 @@ impl<'a, 'll, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'll, 'tcx> {
         unsafe { llvm::LLVMBuildTrunc(self.llbuilder, val, dest_ty, UNNAMED) }
     }
 
+    fn unchecked_utrunc(&mut self, val: &'ll Value, dest_ty: &'ll Type) -> &'ll Value {
+        let trunc = self.trunc(val, dest_ty);
+        if llvm_util::get_version() >= (19, 0, 0) {
+            unsafe {
+                if llvm::LLVMIsATruncInst(trunc).is_some() {
+                    llvm::LLVMSetNUW(trunc, True);
+                }
+            }
+        }
+        trunc
+    }
+
+    fn unchecked_strunc(&mut self, val: &'ll Value, dest_ty: &'ll Type) -> &'ll Value {
+        let trunc = self.trunc(val, dest_ty);
+        if llvm_util::get_version() >= (19, 0, 0) {
+            unsafe {
+                if llvm::LLVMIsATruncInst(trunc).is_some() {
+                    llvm::LLVMSetNSW(trunc, True);
+                }
+            }
+        }
+        trunc
+    }
+
     fn sext(&mut self, val: &'ll Value, dest_ty: &'ll Type) -> &'ll Value {
         unsafe { llvm::LLVMBuildSExt(self.llbuilder, val, dest_ty, UNNAMED) }
     }
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
index 3b0187b9d37..1e105e86189 100644
--- a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
+++ b/compiler/rustc_codegen_llvm/src/llvm/ffi.rs
@@ -1165,6 +1165,7 @@ unsafe extern "C" {
 
     // Operations on instructions
     pub(crate) fn LLVMIsAInstruction(Val: &Value) -> Option<&Value>;
+    pub(crate) fn LLVMIsATruncInst(Val: &Value) -> Option<&Value>;
     pub(crate) fn LLVMGetFirstBasicBlock(Fn: &Value) -> &BasicBlock;
     pub(crate) fn LLVMGetOperand(Val: &Value, Index: c_uint) -> Option<&Value>;
 
diff --git a/compiler/rustc_codegen_ssa/src/base.rs b/compiler/rustc_codegen_ssa/src/base.rs
index e800492dad0..a989e9c64c0 100644
--- a/compiler/rustc_codegen_ssa/src/base.rs
+++ b/compiler/rustc_codegen_ssa/src/base.rs
@@ -24,7 +24,7 @@ use rustc_middle::query::Providers;
 use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutOf, TyAndLayout};
 use rustc_middle::ty::{self, Instance, Ty, TyCtxt};
 use rustc_session::Session;
-use rustc_session::config::{self, CrateType, EntryFnType, OptLevel, OutputType};
+use rustc_session::config::{self, CrateType, EntryFnType, OutputType};
 use rustc_span::{DUMMY_SP, Symbol, sym};
 use rustc_trait_selection::infer::{BoundRegionConversionTime, TyCtxtInferExt};
 use rustc_trait_selection::traits::{ObligationCause, ObligationCtxt};
@@ -364,13 +364,7 @@ pub(crate) fn build_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)
+        if is_unchecked { bx.unchecked_utrunc(rhs, lhs_llty) } else { 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
         // RHS to `255i32`. But then we mask the shift amount to be within the size of the LHS
diff --git a/compiler/rustc_codegen_ssa/src/traits/builder.rs b/compiler/rustc_codegen_ssa/src/traits/builder.rs
index 48ae000f2c6..345db313022 100644
--- a/compiler/rustc_codegen_ssa/src/traits/builder.rs
+++ b/compiler/rustc_codegen_ssa/src/traits/builder.rs
@@ -340,6 +340,17 @@ pub trait BuilderMethods<'a, 'tcx>:
     }
 
     fn trunc(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value;
+    /// Produces the same value as [`Self::trunc`] (and defaults to that),
+    /// but is UB unless the *zero*-extending the result can reproduce `val`.
+    fn unchecked_utrunc(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value {
+        self.trunc(val, dest_ty)
+    }
+    /// Produces the same value as [`Self::trunc`] (and defaults to that),
+    /// but is UB unless the *sign*-extending the result can reproduce `val`.
+    fn unchecked_strunc(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value {
+        self.trunc(val, dest_ty)
+    }
+
     fn sext(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value;
     fn fptoui_sat(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value;
     fn fptosi_sat(&mut self, val: Self::Value, dest_ty: Self::Type) -> Self::Value;
diff --git a/tests/codegen/intrinsics/transmute-niched.rs b/tests/codegen/intrinsics/transmute-niched.rs
index 5ded0e192fa..8ff5cc8ee4f 100644
--- a/tests/codegen/intrinsics/transmute-niched.rs
+++ b/tests/codegen/intrinsics/transmute-niched.rs
@@ -170,7 +170,7 @@ pub unsafe fn check_bool_from_ordering(x: std::cmp::Ordering) -> bool {
     // OPT: call void @llvm.assume(i1 %2)
     // CHECK-NOT: icmp
     // CHECK-NOT: assume
-    // CHECK: %[[R:.+]] = trunc i8 %x to i1
+    // CHECK: %[[R:.+]] = trunc{{( nuw)?}} i8 %x to i1
     // CHECK: ret i1 %[[R]]
 
     transmute(x)
diff --git a/tests/codegen/intrinsics/transmute.rs b/tests/codegen/intrinsics/transmute.rs
index 5b40a6a12c0..ff297b27065 100644
--- a/tests/codegen/intrinsics/transmute.rs
+++ b/tests/codegen/intrinsics/transmute.rs
@@ -11,6 +11,9 @@ use std::intrinsics::mir::*;
 use std::intrinsics::{transmute, transmute_unchecked};
 use std::mem::MaybeUninit;
 
+// FIXME(LLVM18REMOVED): `trunc nuw` doesn't exist in LLVM 18, so once we no
+// longer support it the optional flag checks can be changed to required.
+
 pub enum ZstNever {}
 
 #[repr(align(2))]
@@ -153,7 +156,7 @@ pub unsafe fn check_from_newtype(x: Scalar64) -> u64 {
 pub unsafe fn check_aggregate_to_bool(x: Aggregate8) -> bool {
     // CHECK: %x = alloca [1 x i8], align 1
     // CHECK: %[[BYTE:.+]] = load i8, ptr %x, align 1
-    // CHECK: %[[BOOL:.+]] = trunc i8 %[[BYTE]] to i1
+    // CHECK: %[[BOOL:.+]] = trunc{{( nuw)?}} i8 %[[BYTE]] to i1
     // CHECK: ret i1 %[[BOOL]]
     transmute(x)
 }
@@ -171,7 +174,7 @@ pub unsafe fn check_aggregate_from_bool(x: bool) -> Aggregate8 {
 #[no_mangle]
 pub unsafe fn check_byte_to_bool(x: u8) -> bool {
     // CHECK-NOT: alloca
-    // CHECK: %[[R:.+]] = trunc i8 %x to i1
+    // CHECK: %[[R:.+]] = trunc{{( nuw)?}} i8 %x to i1
     // CHECK: ret i1 %[[R]]
     transmute(x)
 }
@@ -284,7 +287,7 @@ pub unsafe fn check_long_array_more_aligned(x: [u8; 100]) -> [u32; 25] {
 #[no_mangle]
 pub unsafe fn check_pair_with_bool(x: (u8, bool)) -> (bool, i8) {
     // CHECK-NOT: alloca
-    // CHECK: trunc i8 %x.0 to i1
+    // CHECK: trunc{{( nuw)?}} i8 %x.0 to i1
     // CHECK: zext i1 %x.1 to i8
     transmute(x)
 }
@@ -338,7 +341,7 @@ pub unsafe fn check_heterogeneous_integer_pair(x: (i32, bool)) -> (bool, u32) {
     // CHECK: store i8 %[[WIDER]]
 
     // CHECK: %[[BYTE:.+]] = load i8
-    // CHECK: trunc i8 %[[BYTE:.+]] to i1
+    // CHECK: trunc{{( nuw)?}} i8 %[[BYTE:.+]] to i1
     // CHECK: load i32
     transmute(x)
 }
diff --git a/tests/codegen/transmute-scalar.rs b/tests/codegen/transmute-scalar.rs
index 43da7c1781e..c080259a917 100644
--- a/tests/codegen/transmute-scalar.rs
+++ b/tests/codegen/transmute-scalar.rs
@@ -26,7 +26,7 @@ pub fn bool_to_byte(b: bool) -> u8 {
 }
 
 // CHECK-LABEL: define{{.*}}zeroext i1 @byte_to_bool(i8{{.*}} %byte)
-// CHECK: %_0 = trunc i8 %byte to i1
+// CHECK: %_0 = trunc{{( nuw)?}} i8 %byte to i1
 // CHECK-NEXT: ret i1 %_0
 #[no_mangle]
 pub unsafe fn byte_to_bool(byte: u8) -> bool {
diff --git a/tests/codegen/unchecked_shifts.rs b/tests/codegen/unchecked_shifts.rs
index b27eb73c0cc..a0d8990d29f 100644
--- a/tests/codegen/unchecked_shifts.rs
+++ b/tests/codegen/unchecked_shifts.rs
@@ -1,4 +1,10 @@
-//@ compile-flags: -Copt-level=3
+//@ revisions: LLVM18 LLVM19PLUS
+//@ compile-flags: -Copt-level=3 -C no-prepopulate-passes
+//@[LLVM18] exact-llvm-major-version: 18
+//@[LLVM19PLUS] min-llvm-version: 19
+
+// This runs mir-opts to inline the standard library call, but doesn't run LLVM
+// optimizations so it doesn't need to worry about them adding more flags.
 
 #![crate_type = "lib"]
 #![feature(unchecked_shifts)]
@@ -17,12 +23,9 @@ pub unsafe fn unchecked_shl_unsigned_same(a: u32, b: u32) -> u32 {
 // CHECK-LABEL: @unchecked_shl_unsigned_smaller
 #[no_mangle]
 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, 16
-    // CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
-    // CHECK-DAG: %[[TRUNC:.+]] = trunc{{( nuw)?( nsw)?}} i32 %b to i16
+    // CHECK-NOT: assume
+    // LLVM18-DAG: %[[TRUNC:.+]] = trunc i32 %b to i16
+    // LLVM19PLUS-DAG: %[[TRUNC:.+]] = trunc nuw i32 %b to i16
     // CHECK-DAG: shl i16 %a, %[[TRUNC]]
     a.unchecked_shl(b)
 }
@@ -31,7 +34,7 @@ pub unsafe fn unchecked_shl_unsigned_smaller(a: u16, b: u32) -> u16 {
 #[no_mangle]
 pub unsafe fn unchecked_shl_unsigned_bigger(a: u64, b: u32) -> u64 {
     // CHECK-NOT: assume
-    // CHECK: %[[EXT:.+]] = zext{{( nneg)?}} i32 %b to i64
+    // CHECK: %[[EXT:.+]] = zext i32 %b to i64
     // CHECK: shl i64 %a, %[[EXT]]
     a.unchecked_shl(b)
 }
@@ -49,13 +52,10 @@ pub unsafe fn unchecked_shr_signed_same(a: i32, b: u32) -> i32 {
 // CHECK-LABEL: @unchecked_shr_signed_smaller
 #[no_mangle]
 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, 16
-    // CHECK-DAG: tail call void @llvm.assume(i1 %[[INRANGE]])
-    // CHECK-DAG: %[[TRUNC:.+]] = trunc{{( nuw)?( nsw)?}}  i32 %b to i16
-    // CHECK-DAG: ashr i16 %a, %[[TRUNC]]
+    // CHECK-NOT: assume
+    // LLVM18: %[[TRUNC:.+]] = trunc i32 %b to i16
+    // LLVM19PLUS: %[[TRUNC:.+]] = trunc nuw i32 %b to i16
+    // CHECK: ashr i16 %a, %[[TRUNC]]
     a.unchecked_shr(b)
 }
 
@@ -63,7 +63,7 @@ pub unsafe fn unchecked_shr_signed_smaller(a: i16, b: u32) -> i16 {
 #[no_mangle]
 pub unsafe fn unchecked_shr_signed_bigger(a: i64, b: u32) -> i64 {
     // CHECK-NOT: assume
-    // CHECK: %[[EXT:.+]] = zext{{( nneg)?}} i32 %b to i64
+    // CHECK: %[[EXT:.+]] = zext i32 %b to i64
     // CHECK: ashr i64 %a, %[[EXT]]
     a.unchecked_shr(b)
 }
@@ -72,7 +72,7 @@ pub unsafe fn unchecked_shr_signed_bigger(a: i64, b: u32) -> i64 {
 #[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: %[[EXT:.+]] = zext i8 %b to i128
     // CHECK: lshr i128 %a, %[[EXT]]
     std::intrinsics::unchecked_shr(a, b)
 }
@@ -81,7 +81,7 @@ pub unsafe fn unchecked_shr_u128_i8(a: u128, b: i8) -> u128 {
 #[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: %[[EXT:.+]] = zext i8 %b to i128
     // CHECK: shl i128 %a, %[[EXT]]
     std::intrinsics::unchecked_shl(a, b)
 }
@@ -89,25 +89,19 @@ pub unsafe fn unchecked_shl_i128_u8(a: i128, b: u8) -> i128 {
 // 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{{( nuw)?( nsw)?}} i128 %b to i8
-    // CHECK-DAG: shl i8 %a, %[[TRUNC]]
+    // CHECK-NOT: assume
+    // LLVM18: %[[TRUNC:.+]] = trunc i128 %b to i8
+    // LLVM19PLUS: %[[TRUNC:.+]] = trunc nuw i128 %b to i8
+    // CHECK: 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{{( nuw)?( nsw)?}} i128 %b to i8
-    // CHECK-DAG: ashr i8 %a, %[[TRUNC]]
+    // CHECK-NOT: assume
+    // LLVM18: %[[TRUNC:.+]] = trunc i128 %b to i8
+    // LLVM19PLUS: %[[TRUNC:.+]] = trunc nuw i128 %b to i8
+    // CHECK: ashr i8 %a, %[[TRUNC]]
     std::intrinsics::unchecked_shr(a, b)
 }
diff --git a/tests/codegen/union-abi.rs b/tests/codegen/union-abi.rs
index 16022fad8af..28acc4de2f3 100644
--- a/tests/codegen/union-abi.rs
+++ b/tests/codegen/union-abi.rs
@@ -142,4 +142,4 @@ pub union UnionBool {
 pub fn test_UnionBool(b: UnionBool) -> bool {
     unsafe { b.b }
 }
-// CHECK: %_0 = trunc i8 %b to i1
+// CHECK: %_0 = trunc{{( nuw)?}} i8 %b to i1