Do not optimize out SwitchInt before borrowck, or if Zmir-preserve-ub

6 files changed, 32 insertions, 15 deletions

diff --git a/compiler/rustc_mir_transform/src/early_otherwise_branch.rs b/compiler/rustc_mir_transform/src/early_otherwise_branch.rs
index 57f7893be1b..ac97b08fbed 100644
--- a/compiler/rustc_mir_transform/src/early_otherwise_branch.rs
+++ b/compiler/rustc_mir_transform/src/early_otherwise_branch.rs
@@ -224,7 +224,7 @@ impl<'tcx> crate::MirPass<'tcx> for EarlyOtherwiseBranch {
         // Since this optimization adds new basic blocks and invalidates others,
         // clean up the cfg to make it nicer for other passes
         if should_cleanup {
-            simplify_cfg(body);
+            simplify_cfg(tcx, body);
         }
     }
 
diff --git a/compiler/rustc_mir_transform/src/inline.rs b/compiler/rustc_mir_transform/src/inline.rs
index 0ab24e48d44..babee81a2ad 100644
--- a/compiler/rustc_mir_transform/src/inline.rs
+++ b/compiler/rustc_mir_transform/src/inline.rs
@@ -63,7 +63,7 @@ impl<'tcx> crate::MirPass<'tcx> for Inline {
         let _guard = span.enter();
         if inline::<NormalInliner<'tcx>>(tcx, body) {
             debug!("running simplify cfg on {:?}", body.source);
-            simplify_cfg(body);
+            simplify_cfg(tcx, body);
             deref_finder(tcx, body);
         }
     }
@@ -99,7 +99,7 @@ impl<'tcx> crate::MirPass<'tcx> for ForceInline {
         let _guard = span.enter();
         if inline::<ForceInliner<'tcx>>(tcx, body) {
             debug!("running simplify cfg on {:?}", body.source);
-            simplify_cfg(body);
+            simplify_cfg(tcx, body);
             deref_finder(tcx, body);
         }
     }
diff --git a/compiler/rustc_mir_transform/src/match_branches.rs b/compiler/rustc_mir_transform/src/match_branches.rs
index 0d9d0368d37..d43a35124c2 100644
--- a/compiler/rustc_mir_transform/src/match_branches.rs
+++ b/compiler/rustc_mir_transform/src/match_branches.rs
@@ -45,7 +45,7 @@ impl<'tcx> crate::MirPass<'tcx> for MatchBranchSimplification {
         }
 
         if should_cleanup {
-            simplify_cfg(body);
+            simplify_cfg(tcx, body);
         }
     }
 
diff --git a/compiler/rustc_mir_transform/src/remove_place_mention.rs b/compiler/rustc_mir_transform/src/remove_place_mention.rs
index 15fe77d5319..cb598ceb4df 100644
--- a/compiler/rustc_mir_transform/src/remove_place_mention.rs
+++ b/compiler/rustc_mir_transform/src/remove_place_mention.rs
@@ -8,7 +8,7 @@ pub(super) struct RemovePlaceMention;
 
 impl<'tcx> crate::MirPass<'tcx> for RemovePlaceMention {
     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
-        !sess.opts.unstable_opts.mir_keep_place_mention
+        !sess.opts.unstable_opts.mir_preserve_ub
     }
 
     fn run_pass(&self, _: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
diff --git a/compiler/rustc_mir_transform/src/remove_unneeded_drops.rs b/compiler/rustc_mir_transform/src/remove_unneeded_drops.rs
index 8a8cdafc690..43f80508e4a 100644
--- a/compiler/rustc_mir_transform/src/remove_unneeded_drops.rs
+++ b/compiler/rustc_mir_transform/src/remove_unneeded_drops.rs
@@ -35,7 +35,7 @@ impl<'tcx> crate::MirPass<'tcx> for RemoveUnneededDrops {
         // if we applied optimizations, we potentially have some cfg to cleanup to
         // make it easier for further passes
         if should_simplify {
-            simplify_cfg(body);
+            simplify_cfg(tcx, body);
         }
     }
 
diff --git a/compiler/rustc_mir_transform/src/simplify.rs b/compiler/rustc_mir_transform/src/simplify.rs
index 84905f4a400..6c8fa1fe9b3 100644
--- a/compiler/rustc_mir_transform/src/simplify.rs
+++ b/compiler/rustc_mir_transform/src/simplify.rs
@@ -26,6 +26,13 @@
 //! Here the block (`{ return; }`) has the return type `char`, rather than `()`, but the MIR we
 //! naively generate still contains the `_a = ()` write in the unreachable block "after" the
 //! return.
+//!
+//! **WARNING**: This is one of the few optimizations that runs on built and analysis MIR, and
+//! so its effects may affect the type-checking, borrow-checking, and other analysis of MIR.
+//! We must be extremely careful to only apply optimizations that preserve UB and all
+//! non-determinism, since changes here can affect which programs compile in an insta-stable way.
+//! The normal logic that a program with UB can be changed to do anything does not apply to
+//! pre-"runtime" MIR!
 
 use rustc_index::{Idx, IndexSlice, IndexVec};
 use rustc_middle::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor};
@@ -66,8 +73,8 @@ impl SimplifyCfg {
     }
 }
 
-pub(super) fn simplify_cfg(body: &mut Body<'_>) {
-    CfgSimplifier::new(body).simplify();
+pub(super) fn simplify_cfg<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+    CfgSimplifier::new(tcx, body).simplify();
     remove_dead_blocks(body);
 
     // FIXME: Should probably be moved into some kind of pass manager
@@ -79,9 +86,9 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyCfg {
         self.name()
     }
 
-    fn run_pass(&self, _: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         debug!("SimplifyCfg({:?}) - simplifying {:?}", self.name(), body.source);
-        simplify_cfg(body);
+        simplify_cfg(tcx, body);
     }
 
     fn is_required(&self) -> bool {
@@ -90,12 +97,13 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyCfg {
 }
 
 struct CfgSimplifier<'a, 'tcx> {
+    preserve_switch_reads: bool,
     basic_blocks: &'a mut IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
     pred_count: IndexVec<BasicBlock, u32>,
 }
 
 impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
-    fn new(body: &'a mut Body<'tcx>) -> Self {
+    fn new(tcx: TyCtxt<'tcx>, body: &'a mut Body<'tcx>) -> Self {
         let mut pred_count = IndexVec::from_elem(0u32, &body.basic_blocks);
 
         // we can't use mir.predecessors() here because that counts
@@ -110,9 +118,12 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
             }
         }
 
+        // Preserve `SwitchInt` reads on built and analysis MIR, or if `-Zmir-preserve-ub`.
+        let preserve_switch_reads = matches!(body.phase, MirPhase::Built | MirPhase::Analysis(_))
+            || tcx.sess.opts.unstable_opts.mir_preserve_ub;
         let basic_blocks = body.basic_blocks_mut();
 
-        CfgSimplifier { basic_blocks, pred_count }
+        CfgSimplifier { preserve_switch_reads, basic_blocks, pred_count }
     }
 
     fn simplify(mut self) {
@@ -253,9 +264,15 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
 
     // turn a branch with all successors identical to a goto
     fn simplify_branch(&mut self, terminator: &mut Terminator<'tcx>) -> bool {
-        match terminator.kind {
-            TerminatorKind::SwitchInt { .. } => {}
-            _ => return false,
+        // Removing a `SwitchInt` terminator may remove reads that result in UB,
+        // so we must not apply this optimization before borrowck or when
+        // `-Zmir-preserve-ub` is set.
+        if self.preserve_switch_reads {
+            return false;
+        }
+
+        let TerminatorKind::SwitchInt { .. } = terminator.kind else {
+            return false;
         };
 
         let first_succ = {