From 7159aed51e7ef4d8459a10e449bf8e9cc09c9b98 Mon Sep 17 00:00:00 2001
From: clubby789 <jamie@hill-daniel.co.uk>
Date: Mon, 5 Feb 2024 16:15:47 +0000
Subject: [PATCH] Use `br` instead of conditional when branching on constant

---
 compiler/rustc_codegen_ssa/src/mir/block.rs | 18 ++++--
 tests/codegen/constant-branch.rs            | 67 +++++++++++++++++++++
 2 files changed, 80 insertions(+), 5 deletions(-)
 create mode 100644 tests/codegen/constant-branch.rs

diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
index 00007110938..9c7aadb81f8 100644
--- a/compiler/rustc_codegen_ssa/src/mir/block.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -319,7 +319,15 @@ fn codegen_switchint_terminator(
         targets: &SwitchTargets,
     ) {
         let discr = self.codegen_operand(bx, discr);
+        let discr_value = discr.immediate();
         let switch_ty = discr.layout.ty;
+        // If our discriminant is a constant we can branch directly
+        if let Some(const_discr) = bx.const_to_opt_u128(discr_value, false) {
+            let target = targets.target_for_value(const_discr);
+            bx.br(helper.llbb_with_cleanup(self, target));
+            return;
+        };
+
         let mut target_iter = targets.iter();
         if target_iter.len() == 1 {
             // If there are two targets (one conditional, one fallback), emit `br` instead of
@@ -330,14 +338,14 @@ fn codegen_switchint_terminator(
             if switch_ty == bx.tcx().types.bool {
                 // Don't generate trivial icmps when switching on bool.
                 match test_value {
-                    0 => bx.cond_br(discr.immediate(), llfalse, lltrue),
-                    1 => bx.cond_br(discr.immediate(), lltrue, llfalse),
+                    0 => bx.cond_br(discr_value, llfalse, lltrue),
+                    1 => bx.cond_br(discr_value, lltrue, llfalse),
                     _ => bug!(),
                 }
             } else {
                 let switch_llty = bx.immediate_backend_type(bx.layout_of(switch_ty));
                 let llval = bx.const_uint_big(switch_llty, test_value);
-                let cmp = bx.icmp(IntPredicate::IntEQ, discr.immediate(), llval);
+                let cmp = bx.icmp(IntPredicate::IntEQ, discr_value, llval);
                 bx.cond_br(cmp, lltrue, llfalse);
             }
         } else if self.cx.sess().opts.optimize == OptLevel::No
@@ -362,11 +370,11 @@ fn codegen_switchint_terminator(
             let ll2 = helper.llbb_with_cleanup(self, target2);
             let switch_llty = bx.immediate_backend_type(bx.layout_of(switch_ty));
             let llval = bx.const_uint_big(switch_llty, test_value1);
-            let cmp = bx.icmp(IntPredicate::IntEQ, discr.immediate(), llval);
+            let cmp = bx.icmp(IntPredicate::IntEQ, discr_value, llval);
             bx.cond_br(cmp, ll1, ll2);
         } else {
             bx.switch(
-                discr.immediate(),
+                discr_value,
                 helper.llbb_with_cleanup(self, targets.otherwise()),
                 target_iter.map(|(value, target)| (value, helper.llbb_with_cleanup(self, target))),
             );
diff --git a/tests/codegen/constant-branch.rs b/tests/codegen/constant-branch.rs
new file mode 100644
index 00000000000..3328b1eb4a8
--- /dev/null
+++ b/tests/codegen/constant-branch.rs
@@ -0,0 +1,67 @@
+//@ compile-flags: -Zmir-opt-level=0 -C no-prepopulate-passes -Copt-level=0
+// make sure that branching on a constant does not emit a conditional
+// branch or a switch
+
+#![crate_type = "lib"]
+
+// CHECK-LABEL: @if_bool
+#[no_mangle]
+pub fn if_bool() {
+    // CHECK: br label %{{.+}}
+    _ = if true {
+        0
+    } else {
+        1
+    };
+
+    // CHECK: br label %{{.+}}
+    _ = if false {
+        0
+    } else {
+        1
+    };
+}
+
+// CHECK-LABEL: @if_constant_int_eq
+#[no_mangle]
+pub fn if_constant_int_eq() {
+    let val = 0;
+    // CHECK: br label %{{.+}}
+    _ = if val == 0 {
+        0
+    } else {
+        1
+    };
+
+    // CHECK: br label %{{.+}}
+    _ = if val == 1 {
+        0
+    } else {
+        1
+    };
+}
+
+// CHECK-LABEL: @if_constant_match
+#[no_mangle]
+pub fn if_constant_match() {
+    // CHECK: br label %{{.+}}
+    _ = match 1 {
+        1 => 2,
+        2 => 3,
+        _ => 4
+    };
+
+    // CHECK: br label %{{.+}}
+    _ = match 1 {
+        2 => 3,
+        _ => 4
+    };
+
+    // CHECK: br label %[[MINUS1:.+]]
+    _ = match -1 {
+    // CHECK: [[MINUS1]]:
+    // CHECK: store i32 1
+        -1 => 1,
+        _ => 0,
+    }
+}