rust-lang · dianqk · Mar 13, 2024
diff --git a/compiler/rustc_middle/src/mir/patch.rs b/compiler/rustc_middle/src/mir/patch.rs
@@ -11,8 +11,6 @@ pub struct MirPatch<'tcx> {
     resume_block: Option<BasicBlock>,
     // Only for unreachable in cleanup path.
     unreachable_cleanup_block: Option<BasicBlock>,
-    // Only for unreachable not in cleanup path.
-    unreachable_no_cleanup_block: Option<BasicBlock>,
     // Cached block for UnwindTerminate (with reason)
     terminate_block: Option<(BasicBlock, UnwindTerminateReason)>,
     body_span: Span,
@@ -29,7 +27,6 @@ impl<'tcx> MirPatch<'tcx> {
             next_local: body.local_decls.len(),
             resume_block: None,
             unreachable_cleanup_block: None,
-            unreachable_no_cleanup_block: None,
             terminate_block: None,
             body_span: body.span,
         };
@@ -46,12 +43,9 @@ impl<'tcx> MirPatch<'tcx> {
             // Check if we already have an unreachable block
             if matches!(block.terminator().kind, TerminatorKind::Unreachable)
                 && block.statements.is_empty()
+                && block.is_cleanup
             {
-                if block.is_cleanup {
-                    result.unreachable_cleanup_block = Some(bb);
-                } else {
-                    result.unreachable_no_cleanup_block = Some(bb);
-                }
+                result.unreachable_cleanup_block = Some(bb);
                 continue;
             }
 
@@ -101,23 +95,6 @@ impl<'tcx> MirPatch<'tcx> {
         bb
     }
 
-    pub fn unreachable_no_cleanup_block(&mut self) -> BasicBlock {
-        if let Some(bb) = self.unreachable_no_cleanup_block {
-            return bb;
-        }
-
-        let bb = self.new_block(BasicBlockData {
-            statements: vec![],
-            terminator: Some(Terminator {
-                source_info: SourceInfo::outermost(self.body_span),
-                kind: TerminatorKind::Unreachable,
-            }),
-            is_cleanup: false,
-        });
-        self.unreachable_no_cleanup_block = Some(bb);
-        bb
-    }
-
     pub fn terminate_block(&mut self, reason: UnwindTerminateReason) -> BasicBlock {
         if let Some((cached_bb, cached_reason)) = self.terminate_block
             && reason == cached_reason

diff --git a/compiler/rustc_middle/src/mir/terminator.rs b/compiler/rustc_middle/src/mir/terminator.rs
@@ -74,17 +74,6 @@ impl SwitchTargets {
     pub fn target_for_value(&self, value: u128) -> BasicBlock {
         self.iter().find_map(|(v, t)| (v == value).then_some(t)).unwrap_or_else(|| self.otherwise())
     }
-
-    /// Adds a new target to the switch. But You cannot add an already present value.
-    #[inline]
-    pub fn add_target(&mut self, value: u128, bb: BasicBlock) {
-        let value = Pu128(value);
-        if self.values.contains(&value) {
-            bug!("target value {:?} already present", value);
-        }
-        self.values.push(value);
-        self.targets.insert(self.targets.len() - 1, bb);
-    }
 }
 
 pub struct SwitchTargetsIter<'a> {

diff --git a/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs b/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs
@@ -2,10 +2,8 @@
 
 use crate::MirPass;
 use rustc_data_structures::fx::FxHashSet;
-use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::{
-    BasicBlock, BasicBlockData, BasicBlocks, Body, Local, Operand, Rvalue, StatementKind,
-    TerminatorKind,
+    BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, Terminator, TerminatorKind,
 };
 use rustc_middle::ty::layout::TyAndLayout;
 use rustc_middle::ty::{Ty, TyCtxt};
@@ -79,8 +77,7 @@ impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         trace!("UninhabitedEnumBranching starting for {:?}", body.source);
 
-        let mut unreachable_targets = Vec::new();
-        let mut patch = MirPatch::new(body);
+        let mut removable_switchs = Vec::new();
 
         for (bb, bb_data) in body.basic_blocks.iter_enumerated() {
             trace!("processing block {:?}", bb);
@@ -95,73 +92,46 @@ impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
                 tcx.param_env_reveal_all_normalized(body.source.def_id()).and(discriminant_ty),
             );
 
-            let mut allowed_variants = if let Ok(layout) = layout {
+            let allowed_variants = if let Ok(layout) = layout {
                 variant_discriminants(&layout, discriminant_ty, tcx)
-            } else if let Some(variant_range) = discriminant_ty.variant_range(tcx) {
-                variant_range
-                    .map(|variant| {
-                        discriminant_ty.discriminant_for_variant(tcx, variant).unwrap().val
-                    })
-                    .collect()
             } else {
                 continue;
             };
 
             trace!("allowed_variants = {:?}", allowed_variants);
 
-            unreachable_targets.clear();
-            let TerminatorKind::SwitchInt { targets, discr } = &bb_data.terminator().kind else {
-                bug!()
-            };
+            let terminator = bb_data.terminator();
+            let TerminatorKind::SwitchInt { targets, .. } = &terminator.kind else { bug!() };
 
+            let mut reachable_count = 0;
             for (index, (val, _)) in targets.iter().enumerate() {
-                if !allowed_variants.remove(&val) {
-                    unreachable_targets.push(index);
-                }
-            }
-            let otherwise_is_empty_unreachable =
-                body.basic_blocks[targets.otherwise()].is_empty_unreachable();
-            // After resolving https://github.com/llvm/llvm-project/issues/78578,
-            // we can remove the limit on the number of successors.
-            fn check_successors(basic_blocks: &BasicBlocks<'_>, bb: BasicBlock) -> bool {
-                let mut successors = basic_blocks[bb].terminator().successors();
-                let Some(first_successor) = successors.next() else { return true };
-                if successors.next().is_some() {
-                    return true;
+                if allowed_variants.contains(&val) {
+                    reachable_count += 1;
+                } else {
+                    removable_switchs.push((bb, index));
                 }
-                if let TerminatorKind::SwitchInt { .. } =
-                    &basic_blocks[first_successor].terminator().kind
-                {
-                    return false;
-                };
-                true
             }
-            let otherwise_is_last_variant = !otherwise_is_empty_unreachable
-                && allowed_variants.len() == 1
-                && check_successors(&body.basic_blocks, targets.otherwise());
-            let replace_otherwise_to_unreachable = otherwise_is_last_variant
-                || !otherwise_is_empty_unreachable && allowed_variants.is_empty();
 
-            if unreachable_targets.is_empty() && !replace_otherwise_to_unreachable {
-                continue;
+            if reachable_count == allowed_variants.len() {
+                removable_switchs.push((bb, targets.iter().count()));
             }
+        }
 
-            let unreachable_block = patch.unreachable_no_cleanup_block();
-            let mut targets = targets.clone();
-            if replace_otherwise_to_unreachable {
-                if otherwise_is_last_variant {
-                    #[allow(rustc::potential_query_instability)]
-                    let last_variant = *allowed_variants.iter().next().unwrap();
-                    targets.add_target(last_variant, targets.otherwise());
-                }
-                unreachable_targets.push(targets.iter().count());
-            }
-            for index in unreachable_targets.iter() {
-                targets.all_targets_mut()[*index] = unreachable_block;
-            }
-            patch.patch_terminator(bb, TerminatorKind::SwitchInt { targets, discr: discr.clone() });
+        if removable_switchs.is_empty() {
+            return;
         }
 
-        patch.apply(body);
+        let new_block = BasicBlockData::new(Some(Terminator {
+            source_info: body.basic_blocks[removable_switchs[0].0].terminator().source_info,
+            kind: TerminatorKind::Unreachable,
+        }));
+        let unreachable_block = body.basic_blocks.as_mut().push(new_block);
+
+        for (bb, index) in removable_switchs {
+            let bb = &mut body.basic_blocks.as_mut()[bb];
+            let terminator = bb.terminator_mut();
+            let TerminatorKind::SwitchInt { targets, .. } = &mut terminator.kind else { bug!() };
+            targets.all_targets_mut()[index] = unreachable_block;
+        }
     }
 }
diff --git a/tests/codegen/enum/uninhabited_enum_default_branch.rs b/tests/codegen/enum/uninhabited_enum_default_branch.rs
diff --git a/tests/mir-opt/funky_arms.float_to_exponential_common.GVN.panic-abort.diff b/tests/mir-opt/funky_arms.float_to_exponential_common.GVN.panic-abort.diff
@@ -69,7 +69,7 @@
           StorageLive(_6);
           _6 = ((*_1).4: std::option::Option<usize>);
           _7 = discriminant(_6);
-          switchInt(move _7) -> [1: bb4, 0: bb6, otherwise: bb9];
+          switchInt(move _7) -> [1: bb4, otherwise: bb6];
       }
 
       bb4: {
@@ -135,9 +135,5 @@
           StorageDead(_6);
           return;
       }
-
-      bb9: {
-          unreachable;
-      }
   }
 
diff --git a/tests/mir-opt/funky_arms.float_to_exponential_common.GVN.panic-unwind.diff b/tests/mir-opt/funky_arms.float_to_exponential_common.GVN.panic-unwind.diff
@@ -69,7 +69,7 @@
           StorageLive(_6);
           _6 = ((*_1).4: std::option::Option<usize>);
           _7 = discriminant(_6);
-          switchInt(move _7) -> [1: bb4, 0: bb6, otherwise: bb9];
+          switchInt(move _7) -> [1: bb4, otherwise: bb6];
       }
 
       bb4: {
@@ -135,9 +135,5 @@
           StorageDead(_6);
           return;
       }
-
-      bb9: {
-          unreachable;
-      }
   }
 
diff --git a/tests/mir-opt/issues/issue_59352.num_to_digit.PreCodegen.after.panic-abort.mir b/tests/mir-opt/issues/issue_59352.num_to_digit.PreCodegen.after.panic-abort.mir
@@ -33,21 +33,21 @@ fn num_to_digit(_1: char) -> u32 {
         _3 = &_2;
         StorageLive(_4);
         _4 = discriminant(_2);
-        switchInt(move _4) -> [1: bb2, 0: bb6, otherwise: bb8];
+        StorageDead(_3);
+        StorageDead(_2);
+        switchInt(move _4) -> [1: bb2, otherwise: bb7];
     }
 
     bb2: {
         StorageDead(_4);
-        StorageDead(_3);
-        StorageDead(_2);
         StorageLive(_5);
         _5 = char::methods::<impl char>::to_digit(move _1, const 8_u32) -> [return: bb3, unwind unreachable];
     }
 
     bb3: {
         StorageLive(_6);
         _6 = discriminant(_5);
-        switchInt(move _6) -> [0: bb4, 1: bb5, otherwise: bb8];
+        switchInt(move _6) -> [0: bb4, 1: bb5, otherwise: bb6];
     }
 
     bb4: {
@@ -58,22 +58,20 @@ fn num_to_digit(_1: char) -> u32 {
         _0 = move ((_5 as Some).0: u32);
         StorageDead(_6);
         StorageDead(_5);
-        goto -> bb7;
+        goto -> bb8;
     }
 
     bb6: {
-        StorageDead(_4);
-        StorageDead(_3);
-        StorageDead(_2);
-        _0 = const 0_u32;
-        goto -> bb7;
+        unreachable;
     }
 
     bb7: {
-        return;
+        StorageDead(_4);
+        _0 = const 0_u32;
+        goto -> bb8;
     }
 
     bb8: {
-        unreachable;
+        return;
     }
 }
diff --git a/tests/mir-opt/issues/issue_59352.num_to_digit.PreCodegen.after.panic-unwind.mir b/tests/mir-opt/issues/issue_59352.num_to_digit.PreCodegen.after.panic-unwind.mir
@@ -33,21 +33,21 @@ fn num_to_digit(_1: char) -> u32 {
         _3 = &_2;
         StorageLive(_4);
         _4 = discriminant(_2);
-        switchInt(move _4) -> [1: bb2, 0: bb6, otherwise: bb8];
+        StorageDead(_3);
+        StorageDead(_2);
+        switchInt(move _4) -> [1: bb2, otherwise: bb7];
     }
 
     bb2: {
         StorageDead(_4);
-        StorageDead(_3);
-        StorageDead(_2);
         StorageLive(_5);
         _5 = char::methods::<impl char>::to_digit(move _1, const 8_u32) -> [return: bb3, unwind continue];
     }
 
     bb3: {
         StorageLive(_6);
         _6 = discriminant(_5);
-        switchInt(move _6) -> [0: bb4, 1: bb5, otherwise: bb8];
+        switchInt(move _6) -> [0: bb4, 1: bb5, otherwise: bb6];
     }
 
     bb4: {
@@ -58,22 +58,20 @@ fn num_to_digit(_1: char) -> u32 {
         _0 = move ((_5 as Some).0: u32);
         StorageDead(_6);
         StorageDead(_5);
-        goto -> bb7;
+        goto -> bb8;
     }
 
     bb6: {
-        StorageDead(_4);
-        StorageDead(_3);
-        StorageDead(_2);
-        _0 = const 0_u32;
-        goto -> bb7;
+        unreachable;
     }
 
     bb7: {
-        return;
+        StorageDead(_4);
+        _0 = const 0_u32;
+        goto -> bb8;
     }
 
     bb8: {
-        unreachable;
+        return;
     }
 }