Auto merge of #93387 - JakobDegen:improve_partialeq, r=tmiasko

Extend uninhabited enum variant branch elimination to also affect fallthrough

The `uninhabited_enum_branching` mir opt eliminates branches on variants where the data is uninhabited. This change extends this pass to also ensure that the `otherwise` case points to a trivially unreachable bb if all inhabited variants are present in the non-otherwise branches.

I believe it was `@scottmcm` who said that LLVM eliminates some of this information in its SimplifyCFG pass. This is unfortunate, but this change should still be at least a small improvement in principle (I don't think it will show up on any benchmarks)

7 files changed, 167 insertions, 20 deletions

diff --git a/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs b/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs
index cda9ba9dcc8..7133724d07d 100644
--- a/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs
+++ b/compiler/rustc_mir_transform/src/uninhabited_enum_branching.rs
@@ -3,7 +3,8 @@
 use crate::MirPass;
 use rustc_data_structures::stable_set::FxHashSet;
 use rustc_middle::mir::{
-    BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, SwitchTargets, TerminatorKind,
+    BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, SwitchTargets, Terminator,
+    TerminatorKind,
 };
 use rustc_middle::ty::layout::TyAndLayout;
 use rustc_middle::ty::{Ty, TyCtxt};
@@ -71,6 +72,28 @@ fn variant_discriminants<'tcx>(
     }
 }
 
+/// Ensures that the `otherwise` branch leads to an unreachable bb, returning `None` if so and a new
+/// bb to use as the new target if not.
+fn ensure_otherwise_unreachable<'tcx>(
+    body: &Body<'tcx>,
+    targets: &SwitchTargets,
+) -> Option<BasicBlockData<'tcx>> {
+    let otherwise = targets.otherwise();
+    let bb = &body.basic_blocks()[otherwise];
+    if bb.terminator().kind == TerminatorKind::Unreachable
+        && bb.statements.iter().all(|s| matches!(&s.kind, StatementKind::StorageDead(_)))
+    {
+        return None;
+    }
+
+    let mut new_block = BasicBlockData::new(Some(Terminator {
+        source_info: bb.terminator().source_info,
+        kind: TerminatorKind::Unreachable,
+    }));
+    new_block.is_cleanup = bb.is_cleanup;
+    Some(new_block)
+}
+
 impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
         sess.mir_opt_level() > 0
@@ -99,12 +122,25 @@ impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
             if let TerminatorKind::SwitchInt { targets, .. } =
                 &mut body.basic_blocks_mut()[bb].terminator_mut().kind
             {
-                let new_targets = SwitchTargets::new(
+                let mut new_targets = SwitchTargets::new(
                     targets.iter().filter(|(val, _)| allowed_variants.contains(val)),
                     targets.otherwise(),
                 );
 
-                *targets = new_targets;
+                if new_targets.iter().count() == allowed_variants.len() {
+                    if let Some(updated) = ensure_otherwise_unreachable(body, &new_targets) {
+                        let new_otherwise = body.basic_blocks_mut().push(updated);
+                        *new_targets.all_targets_mut().last_mut().unwrap() = new_otherwise;
+                    }
+                }
+
+                if let TerminatorKind::SwitchInt { targets, .. } =
+                    &mut body.basic_blocks_mut()[bb].terminator_mut().kind
+                {
+                    *targets = new_targets;
+                } else {
+                    unreachable!()
+                }
             } else {
                 unreachable!()
             }
diff --git a/src/test/mir-opt/early_otherwise_branch.opt2.EarlyOtherwiseBranch.diff b/src/test/mir-opt/early_otherwise_branch.opt2.EarlyOtherwiseBranch.diff
index 6d149b89edb..79b923a2889 100644
--- a/src/test/mir-opt/early_otherwise_branch.opt2.EarlyOtherwiseBranch.diff
+++ b/src/test/mir-opt/early_otherwise_branch.opt2.EarlyOtherwiseBranch.diff
@@ -31,7 +31,7 @@
           StorageDead(_5);                 // scope 0 at $DIR/early_otherwise_branch.rs:12:16: 12:17
           StorageDead(_4);                 // scope 0 at $DIR/early_otherwise_branch.rs:12:16: 12:17
           _8 = discriminant((_3.0: std::option::Option<u32>)); // scope 0 at $DIR/early_otherwise_branch.rs:12:11: 12:17
--         switchInt(move _8) -> [0_isize: bb1, 1_isize: bb3, otherwise: bb2]; // scope 0 at $DIR/early_otherwise_branch.rs:12:5: 12:17
+-         switchInt(move _8) -> [0_isize: bb1, 1_isize: bb3, otherwise: bb7]; // scope 0 at $DIR/early_otherwise_branch.rs:12:5: 12:17
 +         StorageLive(_11);                // scope 0 at $DIR/early_otherwise_branch.rs:12:5: 12:17
 +         _11 = discriminant((_3.1: std::option::Option<u32>)); // scope 0 at $DIR/early_otherwise_branch.rs:12:5: 12:17
 +         StorageLive(_12);                // scope 0 at $DIR/early_otherwise_branch.rs:12:5: 12:17
@@ -81,8 +81,10 @@
 +     bb4: {
           StorageDead(_3);                 // scope 0 at $DIR/early_otherwise_branch.rs:17:1: 17:2
           return;                          // scope 0 at $DIR/early_otherwise_branch.rs:17:2: 17:2
-+     }
-+ 
+      }
+  
+-     bb7: {
+-         unreachable;                     // scope 0 at $DIR/early_otherwise_branch.rs:15:14: 15:15
 +     bb5: {
 +         StorageDead(_12);                // scope 0 at $DIR/early_otherwise_branch.rs:12:5: 12:17
 +         switchInt(_8) -> [0_isize: bb3, 1_isize: bb2, otherwise: bb1]; // scope 0 at $DIR/early_otherwise_branch.rs:12:5: 12:17
diff --git a/src/test/mir-opt/early_otherwise_branch_68867.try_sum.EarlyOtherwiseBranch.before-SimplifyConstCondition-final.after.diff b/src/test/mir-opt/early_otherwise_branch_68867.try_sum.EarlyOtherwiseBranch.before-SimplifyConstCondition-final.after.diff
index fc2dcb25109..db6794db298 100644
--- a/src/test/mir-opt/early_otherwise_branch_68867.try_sum.EarlyOtherwiseBranch.before-SimplifyConstCondition-final.after.diff
+++ b/src/test/mir-opt/early_otherwise_branch_68867.try_sum.EarlyOtherwiseBranch.before-SimplifyConstCondition-final.after.diff
@@ -81,7 +81,7 @@
 -         StorageDead(_5);                 // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:23: 21:24
 +         nop;                             // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:23: 21:24
           _11 = discriminant((*(_4.0: &ViewportPercentageLength))); // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:14: 21:24
--         switchInt(move _11) -> [0_isize: bb1, 1_isize: bb3, 2_isize: bb4, 3_isize: bb5, otherwise: bb2]; // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
+-         switchInt(move _11) -> [0_isize: bb1, 1_isize: bb3, 2_isize: bb4, 3_isize: bb5, otherwise: bb11]; // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
 +         StorageLive(_34);                // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
 +         _34 = discriminant((*(_4.1: &ViewportPercentageLength))); // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
 +         StorageLive(_35);                // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
@@ -217,14 +217,9 @@
 -         StorageDead(_13);                // scope 0 at $DIR/early_otherwise_branch_68867.rs:22:49: 22:50
 -         StorageDead(_12);                // scope 0 at $DIR/early_otherwise_branch_68867.rs:22:49: 22:50
 -         goto -> bb10;                    // scope 0 at $DIR/early_otherwise_branch_68867.rs:22:49: 22:50
-+         nop;                             // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:5: 27:7
-+         discriminant(_0) = 0;            // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:5: 27:7
-+         nop;                             // scope 0 at $DIR/early_otherwise_branch_68867.rs:27:6: 27:7
-+         nop;                             // scope 0 at $DIR/early_otherwise_branch_68867.rs:28:1: 28:2
-+         return;                          // scope 0 at $DIR/early_otherwise_branch_68867.rs:28:2: 28:2
-      }
-  
-      bb7: {
+-     }
+- 
+-     bb7: {
 -         StorageLive(_17);                // scope 0 at $DIR/early_otherwise_branch_68867.rs:23:14: 23:17
 -         _17 = (((*(_4.0: &ViewportPercentageLength)) as Vh).0: f32); // scope 0 at $DIR/early_otherwise_branch_68867.rs:23:14: 23:17
 -         StorageLive(_18);                // scope 0 at $DIR/early_otherwise_branch_68867.rs:23:24: 23:29
@@ -289,10 +284,18 @@
 - 
 -     bb10: {
 -         ((_0 as Ok).0: ViewportPercentageLength) = move _3; // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:5: 27:7
--         discriminant(_0) = 0;            // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:5: 27:7
++         nop;                             // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:5: 27:7
+          discriminant(_0) = 0;            // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:5: 27:7
 -         StorageDead(_3);                 // scope 0 at $DIR/early_otherwise_branch_68867.rs:27:6: 27:7
 -         StorageDead(_4);                 // scope 0 at $DIR/early_otherwise_branch_68867.rs:28:1: 28:2
--         return;                          // scope 0 at $DIR/early_otherwise_branch_68867.rs:28:2: 28:2
++         nop;                             // scope 0 at $DIR/early_otherwise_branch_68867.rs:27:6: 27:7
++         nop;                             // scope 0 at $DIR/early_otherwise_branch_68867.rs:28:1: 28:2
+          return;                          // scope 0 at $DIR/early_otherwise_branch_68867.rs:28:2: 28:2
+      }
+  
+-     bb11: {
+-         unreachable;                     // scope 0 at $DIR/early_otherwise_branch_68867.rs:28:2: 28:2
++     bb7: {
 +         StorageDead(_35);                // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
 +         switchInt(_11) -> [0_isize: bb2, 1_isize: bb3, 2_isize: bb4, 3_isize: bb5, otherwise: bb1]; // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
       }
diff --git a/src/test/mir-opt/early_otherwise_branch_68867.try_sum.EarlyOtherwiseBranch.diff b/src/test/mir-opt/early_otherwise_branch_68867.try_sum.EarlyOtherwiseBranch.diff
index 28c650d72dc..c8d8ae7766d 100644
--- a/src/test/mir-opt/early_otherwise_branch_68867.try_sum.EarlyOtherwiseBranch.diff
+++ b/src/test/mir-opt/early_otherwise_branch_68867.try_sum.EarlyOtherwiseBranch.diff
@@ -67,7 +67,7 @@
           StorageDead(_6);                 // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:23: 21:24
           StorageDead(_5);                 // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:23: 21:24
           _11 = discriminant((*(_4.0: &ViewportPercentageLength))); // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:14: 21:24
--         switchInt(move _11) -> [0_isize: bb1, 1_isize: bb3, 2_isize: bb4, 3_isize: bb5, otherwise: bb2]; // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
+-         switchInt(move _11) -> [0_isize: bb1, 1_isize: bb3, 2_isize: bb4, 3_isize: bb5, otherwise: bb11]; // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
 +         StorageLive(_34);                // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
 +         _34 = discriminant((*(_4.1: &ViewportPercentageLength))); // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
 +         StorageLive(_35);                // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
@@ -206,8 +206,10 @@
           StorageDead(_3);                 // scope 0 at $DIR/early_otherwise_branch_68867.rs:27:6: 27:7
           StorageDead(_4);                 // scope 0 at $DIR/early_otherwise_branch_68867.rs:28:1: 28:2
           return;                          // scope 0 at $DIR/early_otherwise_branch_68867.rs:28:2: 28:2
-+     }
-+ 
+      }
+  
+-     bb11: {
+-         unreachable;                     // scope 0 at $DIR/early_otherwise_branch_68867.rs:28:2: 28:2
 +     bb7: {
 +         StorageDead(_35);                // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
 +         switchInt(_11) -> [0_isize: bb2, 1_isize: bb3, 2_isize: bb4, 3_isize: bb5, otherwise: bb1]; // scope 0 at $DIR/early_otherwise_branch_68867.rs:21:8: 21:24
diff --git a/src/test/mir-opt/uninhabited_fallthrough_elimination.eliminate_fallthrough.UninhabitedEnumBranching.diff b/src/test/mir-opt/uninhabited_fallthrough_elimination.eliminate_fallthrough.UninhabitedEnumBranching.diff
new file mode 100644
index 00000000000..7e843b65e88
--- /dev/null
+++ b/src/test/mir-opt/uninhabited_fallthrough_elimination.eliminate_fallthrough.UninhabitedEnumBranching.diff
@@ -0,0 +1,38 @@
+- // MIR for `eliminate_fallthrough` before UninhabitedEnumBranching
++ // MIR for `eliminate_fallthrough` after UninhabitedEnumBranching
+  
+  fn eliminate_fallthrough(_1: S) -> u32 {
+      debug s => _1;                       // in scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:21:26: 21:27
+      let mut _0: u32;                     // return place in scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:21:35: 21:38
+      let mut _2: isize;                   // in scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:23:9: 23:10
+  
+      bb0: {
+          _2 = discriminant(_1);           // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:22:11: 22:12
+-         switchInt(move _2) -> [1_isize: bb3, 2_isize: bb2, otherwise: bb1]; // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:22:5: 22:12
++         switchInt(move _2) -> [1_isize: bb3, 2_isize: bb2, otherwise: bb5]; // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:22:5: 22:12
+      }
+  
+      bb1: {
+          _0 = const 3_u32;                // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:25:14: 25:15
+          goto -> bb4;                     // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:25:14: 25:15
+      }
+  
+      bb2: {
+          _0 = const 1_u32;                // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:23:14: 23:15
+          goto -> bb4;                     // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:23:14: 23:15
+      }
+  
+      bb3: {
+          _0 = const 2_u32;                // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:24:14: 24:15
+          goto -> bb4;                     // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:24:14: 24:15
+      }
+  
+      bb4: {
+          return;                          // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:27:2: 27:2
++     }
++ 
++     bb5: {
++         unreachable;                     // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:25:14: 25:15
+      }
+  }
+  
diff --git a/src/test/mir-opt/uninhabited_fallthrough_elimination.keep_fallthrough.UninhabitedEnumBranching.diff b/src/test/mir-opt/uninhabited_fallthrough_elimination.keep_fallthrough.UninhabitedEnumBranching.diff
new file mode 100644
index 00000000000..5da011d427a
--- /dev/null
+++ b/src/test/mir-opt/uninhabited_fallthrough_elimination.keep_fallthrough.UninhabitedEnumBranching.diff
@@ -0,0 +1,34 @@
+- // MIR for `keep_fallthrough` before UninhabitedEnumBranching
++ // MIR for `keep_fallthrough` after UninhabitedEnumBranching
+  
+  fn keep_fallthrough(_1: S) -> u32 {
+      debug s => _1;                       // in scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:12:21: 12:22
+      let mut _0: u32;                     // return place in scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:12:30: 12:33
+      let mut _2: isize;                   // in scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:14:9: 14:13
+  
+      bb0: {
+          _2 = discriminant(_1);           // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:13:11: 13:12
+-         switchInt(move _2) -> [0_isize: bb2, 1_isize: bb3, otherwise: bb1]; // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:13:5: 13:12
++         switchInt(move _2) -> [1_isize: bb3, otherwise: bb1]; // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:13:5: 13:12
+      }
+  
+      bb1: {
+          _0 = const 3_u32;                // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:16:14: 16:15
+          goto -> bb4;                     // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:16:14: 16:15
+      }
+  
+      bb2: {
+          _0 = const 1_u32;                // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:14:17: 14:18
+          goto -> bb4;                     // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:14:17: 14:18
+      }
+  
+      bb3: {
+          _0 = const 2_u32;                // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:15:14: 15:15
+          goto -> bb4;                     // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:15:14: 15:15
+      }
+  
+      bb4: {
+          return;                          // scope 0 at $DIR/uninhabited_fallthrough_elimination.rs:18:2: 18:2
+      }
+  }
+  
diff --git a/src/test/mir-opt/uninhabited_fallthrough_elimination.rs b/src/test/mir-opt/uninhabited_fallthrough_elimination.rs
new file mode 100644
index 00000000000..0853883f8b8
--- /dev/null
+++ b/src/test/mir-opt/uninhabited_fallthrough_elimination.rs
@@ -0,0 +1,32 @@
+enum Empty {}
+
+enum S {
+    A(Empty),
+    B,
+    C,
+}
+
+use S::*;
+
+// EMIT_MIR uninhabited_fallthrough_elimination.keep_fallthrough.UninhabitedEnumBranching.diff
+fn keep_fallthrough(s: S) -> u32 {
+    match s {
+        A(_) => 1,
+        B => 2,
+        _ => 3,
+    }
+}
+
+// EMIT_MIR uninhabited_fallthrough_elimination.eliminate_fallthrough.UninhabitedEnumBranching.diff
+fn eliminate_fallthrough(s: S) -> u32 {
+    match s {
+        C => 1,
+        B => 2,
+        _ => 3,
+    }
+}
+
+fn main() {
+    keep_fallthrough(B);
+    eliminate_fallthrough(B);
+}