Use trunc nuw+br for 0/1 branches even in optimized builds

scottmcm · scottmcm · commit 5a982e7eccd4 · 2025-03-01T22:37:11.000-08:00
Rather than needing to use `switch` for them to include the `unreachable` arm
diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -3,6 +3,7 @@ use std::cmp;
 use rustc_abi::{BackendRepr, ExternAbi, HasDataLayout, Reg, WrappingRange};
 use rustc_ast as ast;
 use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
+use rustc_data_structures::packed::Pu128;
 use rustc_hir::lang_items::LangItem;
 use rustc_middle::mir::{self, AssertKind, InlineAsmMacro, SwitchTargets, UnwindTerminateReason};
 use rustc_middle::ty::layout::{HasTyCtxt, LayoutOf, ValidityRequirement};
@@ -406,6 +407,39 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 let cmp = bx.icmp(IntPredicate::IntEQ, discr_value, llval);
                 bx.cond_br_with_expect(cmp, lltarget, llotherwise, expect);
             }
+        } else if target_iter.len() == 2
+            && self.mir[targets.otherwise()].is_empty_unreachable()
+            && targets.all_values().contains(&Pu128(0))
+            && targets.all_values().contains(&Pu128(1))
+        {
+            // This is the really common case for `bool`, `Option`, etc.
+            // By using `trunc nuw` we communicate that other values are
+            // impossible without needing `switch` or `assume`s.
+            let true_bb = targets.target_for_value(1);
+            let false_bb = targets.target_for_value(0);
+            let true_ll = helper.llbb_with_cleanup(self, true_bb);
+            let false_ll = helper.llbb_with_cleanup(self, false_bb);
+
+            let expected_cond_value = if self.cx.sess().opts.optimize == OptLevel::No {
+                None
+            } else {
+                match (self.cold_blocks[true_bb], self.cold_blocks[false_bb]) {
+                    // Same coldness, no expectation
+                    (true, true) | (false, false) => None,
+                    // Different coldness, expect the non-cold one
+                    (true, false) => Some(false),
+                    (false, true) => Some(true),
+                }
+            };
+
+            let bool_ty = bx.tcx().types.bool;
+            let cond = if switch_ty == bool_ty {
+                discr_value
+            } else {
+                let bool_llty = bx.immediate_backend_type(bx.layout_of(bool_ty));
+                bx.unchecked_utrunc(discr_value, bool_llty)
+            };
+            bx.cond_br_with_expect(cond, true_ll, false_ll, expected_cond_value);
         } else if self.cx.sess().opts.optimize == OptLevel::No
             && target_iter.len() == 2
             && self.mir[targets.otherwise()].is_empty_unreachable()
diff --git a/tests/codegen/enum/enum-two-variants-match.rs b/tests/codegen/enum/enum-two-variants-match.rs
@@ -0,0 +1,51 @@
+//@ compile-flags: -Copt-level=3 -C no-prepopulate-passes
+//@ min-llvm-version: 19 (for trunc nuw)
+//@ only-x86_64 (because these discriminants are isize)
+
+#![crate_type = "lib"]
+
+// CHECK-LABEL: @option_match
+#[no_mangle]
+pub fn option_match(x: Option<i32>) -> u16 {
+    // CHECK: %x = alloca [8 x i8]
+    // CHECK: store i32 %0, ptr %x
+    // CHECK: %[[TAG:.+]] = load i32, ptr %x
+    // CHECK-SAME: !range ![[ZERO_ONE_32:[0-9]+]]
+    // CHECK: %[[DISCR:.+]] = zext i32 %[[TAG]] to i64
+    // CHECK: %[[COND:.+]] = trunc nuw i64 %[[DISCR]] to i1
+    // CHECK: br i1 %[[COND]], label %[[TRUE:[a-z0-9]+]], label %[[FALSE:[a-z0-9]+]]
+
+    // CHECK: [[TRUE]]:
+    // CHECK: store i16 13
+
+    // CHECK: [[FALSE]]:
+    // CHECK: store i16 42
+    match x {
+        Some(_) => 13,
+        None => 42,
+    }
+}
+
+// CHECK-LABEL: @result_match
+#[no_mangle]
+pub fn result_match(x: Result<u64, i64>) -> u16 {
+    // CHECK: %x = alloca [16 x i8]
+    // CHECK: store i64 %0, ptr %x
+    // CHECK: %[[DISCR:.+]] = load i64, ptr %x
+    // CHECK-SAME: !range ![[ZERO_ONE_64:[0-9]+]]
+    // CHECK: %[[COND:.+]] = trunc nuw i64 %[[DISCR]] to i1
+    // CHECK: br i1 %[[COND]], label %[[TRUE:[a-z0-9]+]], label %[[FALSE:[a-z0-9]+]]
+
+    // CHECK: [[TRUE]]:
+    // CHECK: store i16 13
+
+    // CHECK: [[FALSE]]:
+    // CHECK: store i16 42
+    match x {
+        Err(_) => 13,
+        Ok(_) => 42,
+    }
+}
+
+// CHECK: ![[ZERO_ONE_32]] = !{i32 0, i32 2}
+// CHECK: ![[ZERO_ONE_64]] = !{i64 0, i64 2}
diff --git a/tests/codegen/intrinsics/cold_path2.rs b/tests/codegen/intrinsics/cold_path2.rs
@@ -26,7 +26,8 @@ pub fn test(x: Option<bool>) {
     }
 
     // CHECK-LABEL: @test(
-    // CHECK: br i1 %1, label %bb2, label %bb1, !prof ![[NUM:[0-9]+]]
+    // CHECK: %[[IS_NONE:.+]] = icmp eq i8 %0, 2
+    // CHECK: br i1 %[[IS_NONE]], label %bb2, label %bb1, !prof ![[NUM:[0-9]+]]
     // CHECK: bb1:
     // CHECK: path_a
     // CHECK: bb2:
diff --git a/tests/codegen/issues/issue-101082.rs b/tests/codegen/issues/issue-101082.rs
@@ -1,5 +1,6 @@
 //@ compile-flags: -Copt-level=3
 //@ revisions: host x86-64-v3
+//@ min-llvm-version: 19
 
 // This particular CPU regressed in #131563
 //@[x86-64-v3] only-x86_64
diff --git a/tests/codegen/issues/issue-122600-ptr-discriminant-update.rs b/tests/codegen/issues/issue-122600-ptr-discriminant-update.rs
@@ -3,6 +3,11 @@
 
 #![crate_type = "lib"]
 
+// The bug here was that it was loading and storing the whole value.
+// It's ok for it to load the discriminant,
+// to preserve the UB from `unreachable_unchecked`,
+// but it better only store the constant discriminant of `B`.
+
 pub enum State {
     A([u8; 753]),
     B([u8; 753]),
@@ -11,9 +16,27 @@ pub enum State {
 // CHECK-LABEL: @update
 #[no_mangle]
 pub unsafe fn update(s: *mut State) {
-    // CHECK-NEXT: start:
-    // CHECK-NEXT: store i8
-    // CHECK-NEXT: ret
+    // CHECK-NOT: alloca
+
+    // CHECK-NOT: load
+    // CHECK-NOT: store
+    // CHECK-NOT: memcpy
+    // CHECK-NOT: 75{{3|4}}
+
+    // CHECK: %[[TAG:.+]] = load i8, ptr %s, align 1
+    // CHECK-NEXT: trunc nuw i8 %[[TAG]] to i1
+
+    // CHECK-NOT: load
+    // CHECK-NOT: store
+    // CHECK-NOT: memcpy
+    // CHECK-NOT: 75{{3|4}}
+
+    // CHECK: store i8 1, ptr %s, align 1
+
+    // CHECK-NOT: load
+    // CHECK-NOT: store
+    // CHECK-NOT: memcpy
+    // CHECK-NOT: 75{{3|4}}
     let State::A(v) = s.read() else { std::hint::unreachable_unchecked() };
     s.write(State::B(v));
 }
diff --git a/tests/codegen/try_question_mark_nop.rs b/tests/codegen/try_question_mark_nop.rs
@@ -16,8 +16,8 @@ use std::ptr::NonNull;
 #[no_mangle]
 pub fn option_nop_match_32(x: Option<u32>) -> Option<u32> {
     // CHECK: start:
-    // TWENTY-NEXT: %trunc = trunc nuw i32 %0 to i1
-    // TWENTY-NEXT: %.2 = select i1 %trunc, i32 %1, i32 undef
+    // TWENTY-NEXT: %[[IS_SOME:.+]] = trunc nuw i32 %0 to i1
+    // TWENTY-NEXT: %.2 = select i1 %[[IS_SOME]], i32 %1, i32 undef
     // CHECK-NEXT: [[REG1:%.*]] = insertvalue { i32, i32 } poison, i32 %0, 0
     // NINETEEN-NEXT: [[REG2:%.*]] = insertvalue { i32, i32 } [[REG1]], i32 %1, 1
     // TWENTY-NEXT: [[REG2:%.*]] = insertvalue { i32, i32 } [[REG1]], i32 %.2, 1
@@ -32,8 +32,8 @@ pub fn option_nop_match_32(x: Option<u32>) -> Option<u32> {
 #[no_mangle]
 pub fn option_nop_traits_32(x: Option<u32>) -> Option<u32> {
     // CHECK: start:
-    // TWENTY-NEXT: %trunc = trunc nuw i32 %0 to i1
-    // TWENTY-NEXT: %.1 = select i1 %trunc, i32 %1, i32 undef
+    // TWENTY-NEXT: %[[IS_SOME:.+]] = trunc nuw i32 %0 to i1
+    // TWENTY-NEXT: %.1 = select i1 %[[IS_SOME]], i32 %1, i32 undef
     // CHECK-NEXT: insertvalue { i32, i32 }
     // CHECK-NEXT: insertvalue { i32, i32 }
     // CHECK-NEXT: ret { i32, i32 }

Original file line number	Diff line number	Diff line change
`@@ -26,7 +26,8 @@ pub fn test(x: Option<bool>) {`
`26`	`26`	`}`
`27`	`27`
`28`	`28`	`// CHECK-LABEL: @test(`
`29`		`- // CHECK: br i1 %1, label %bb2, label %bb1, !prof ![[NUM:[0-9]+]]`
	`29`	`+ // CHECK: %[[IS_NONE:.+]] = icmp eq i8 %0, 2`
	`30`	`+ // CHECK: br i1 %[[IS_NONE]], label %bb2, label %bb1, !prof ![[NUM:[0-9]+]]`
`30`	`31`	`// CHECK: bb1:`
`31`	`32`	`// CHECK: path_a`
`32`	`33`	`// CHECK: bb2:`