From fd7a2351d9020753fc58efa1a9a05bc754ba384a 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   | 10 ++++++-
 .../rustc_codegen_ssa/src/mir/statement.rs    | 28 +++++++++++++++++--
 tests/codegen/discriminant-swap.rs            | 28 +++++++++++++++++++
 3 files changed, 63 insertions(+), 3 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..e4984649f788c 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,10 +1214,17 @@ 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;
+                }
             }
 
+            if replaced_terminator {
+                break;
+            }
             let merging_succ = self.codegen_terminator(bx, bb, data.terminator());
             if let MergingSucc::False = merging_succ {
                 break;
diff --git a/compiler/rustc_codegen_ssa/src/mir/statement.rs b/compiler/rustc_codegen_ssa/src/mir/statement.rs
index ac7dfbb261dec..9e1fba23329f3 100644
--- a/compiler/rustc_codegen_ssa/src/mir/statement.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/statement.rs
@@ -6,9 +6,22 @@ use super::FunctionCx;
 use super::LocalRef;
 use crate::traits::*;
 
+use std::ops::ControlFlow;
+
 impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
+    /// Lower a single MIR statement, returning [`ControlFlow::Break`] if we must not continue lowering
+    /// the rest of the statements in the block.
+    ///
+    /// Since we are lowering a polymorphic MIR body, we can discover only at this point that we
+    /// are lowering `assume(false)`. If we encounter such a statement, there is no reason to lower
+    /// the rest of the block; we just emit an unreachable terminator and return
+    /// [`ControlFlow::Break`].
     #[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 +83,17 @@ 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 we are lowering assume(false), just produce an unreachable
+                        // terminator. We don't emit anything for assume(true).
+                        if value == 0 {
+                            bx.unreachable();
+                            return ControlFlow::Break(());
+                        }
+                    } else {
+                        bx.assume(imm);
+                    }
                 }
             }
             mir::StatementKind::Intrinsic(box NonDivergingIntrinsic::CopyNonOverlapping(
@@ -97,5 +120,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..9a7f36828c9ff
--- /dev/null
+++ b/tests/codegen/discriminant-swap.rs
@@ -0,0 +1,28 @@
+// The lowering of the function below initially reads and writes to the entire pointee, even
+// though it only needs to do a store to the discriminant.
+// This test ensures that std::hint::unreachable_unchecked does not prevent the desired
+// optimization.
+
+//@ 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));
+}