Lower assume(false) to an unreachable terminator

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;

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(())
     }
 }

tests/codegen/discriminant-swap.rs

@@ -0,0 +1,27 @@
+// 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));
+}