From 9cc67c530d288fcae72438e6960445e012ff563e Mon Sep 17 00:00:00 2001
From: Ben Kimock <kimockb@gmail.com>
Date: Sat, 16 Mar 2024 19:12:37 -0400
Subject: [PATCH] Lower assume(false) to an unreachable terminator

---
 compiler/rustc_codegen_ssa/src/mir/block.rs   | 15 +++++++++---
 .../rustc_codegen_ssa/src/mir/statement.rs    | 19 +++++++++++++--
 tests/codegen/discriminant-swap.rs            | 23 +++++++++++++++++++
 3 files changed, 52 insertions(+), 5 deletions(-)
 create mode 100644 tests/codegen/discriminant-swap.rs

diff --git a/compiler/rustc_codegen_ssa/src/mir/block.rs b/compiler/rustc_codegen_ssa/src/mir/block.rs
index 9bb2a52826585..88db95fd2571d 100644
--- a/compiler/rustc_codegen_ssa/src/mir/block.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/block.rs
@@ -23,6 +23,7 @@ use rustc_target::abi::{self, HasDataLayout, WrappingRange};
 use rustc_target::spec::abi::Abi;
 
 use std::cmp;
+use std::ops::ControlFlow;
 
 // Indicates if we are in the middle of merging a BB's successor into it. This
 // can happen when BB jumps directly to its successor and the successor has no
@@ -1213,13 +1214,21 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
 
             debug!("codegen_block({:?}={:?})", bb, data);
 
+            let mut replaced_terminator = false;
             for statement in &data.statements {
-                self.codegen_statement(bx, statement);
+                if let ControlFlow::Break(()) = self.codegen_statement(bx, statement) {
+                    replaced_terminator = true;
+                    break;
+                }
             }
 
-            let merging_succ = self.codegen_terminator(bx, bb, data.terminator());
-            if let MergingSucc::False = merging_succ {
+            if replaced_terminator {
                 break;
+            } else {
+                let merging_succ = self.codegen_terminator(bx, bb, data.terminator());
+                if let MergingSucc::False = merging_succ {
+                    break;
+                }
             }
 
             // We are merging the successor into the produced backend basic
diff --git a/compiler/rustc_codegen_ssa/src/mir/statement.rs b/compiler/rustc_codegen_ssa/src/mir/statement.rs
index ac7dfbb261dec..68fbc733cb094 100644
--- a/compiler/rustc_codegen_ssa/src/mir/statement.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/statement.rs
@@ -6,9 +6,15 @@ use super::FunctionCx;
 use super::LocalRef;
 use crate::traits::*;
 
+use std::ops::ControlFlow;
+
 impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
     #[instrument(level = "debug", skip(self, bx))]
-    pub fn codegen_statement(&mut self, bx: &mut Bx, statement: &mir::Statement<'tcx>) {
+    pub fn codegen_statement(
+        &mut self,
+        bx: &mut Bx,
+        statement: &mir::Statement<'tcx>,
+    ) -> ControlFlow<()> {
         self.set_debug_loc(bx, statement.source_info);
         match statement.kind {
             mir::StatementKind::Assign(box (ref place, ref rvalue)) => {
@@ -70,7 +76,15 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             mir::StatementKind::Intrinsic(box NonDivergingIntrinsic::Assume(ref op)) => {
                 if !matches!(bx.tcx().sess.opts.optimize, OptLevel::No | OptLevel::Less) {
                     let op_val = self.codegen_operand(bx, op);
-                    bx.assume(op_val.immediate());
+                    let imm = op_val.immediate();
+                    if let Some(value) = bx.const_to_opt_uint(imm) {
+                        if value == 0 {
+                            bx.unreachable();
+                            return ControlFlow::Break(());
+                        }
+                    } else {
+                        bx.assume(imm);
+                    }
                 }
             }
             mir::StatementKind::Intrinsic(box NonDivergingIntrinsic::CopyNonOverlapping(
@@ -97,5 +111,6 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
             | mir::StatementKind::PlaceMention(..)
             | mir::StatementKind::Nop => {}
         }
+        ControlFlow::Continue(())
     }
 }
diff --git a/tests/codegen/discriminant-swap.rs b/tests/codegen/discriminant-swap.rs
new file mode 100644
index 0000000000000..db3a93f6b238a
--- /dev/null
+++ b/tests/codegen/discriminant-swap.rs
@@ -0,0 +1,23 @@
+//@ compile-flags: -O
+
+#![crate_type = "lib"]
+
+use std::hint::unreachable_unchecked;
+use std::ptr::{read, write};
+
+type T = [u8; 753];
+
+pub enum State {
+    A(T),
+    B(T),
+}
+
+// CHECK-LABEL: @init(ptr {{.*}}s)
+// CHECK-NEXT: start
+// CHECK-NEXT: store i8 1, ptr %s, align 1
+// CHECK-NEXT: ret void
+#[no_mangle]
+unsafe fn init(s: *mut State) {
+    let State::A(v) = read(s) else { unreachable_unchecked() };
+    write(s, State::B(v));
+}