MatchBranchSimplification: Consider empty-unreachable otherwise branch

Author: clubby789

compiler/rustc_data_structures/src/packed.rs

@@ -18,6 +18,13 @@ impl Pu128 {
     }
 }
 
+impl From<Pu128> for u128 {
+    #[inline]
+    fn from(value: Pu128) -> Self {
+        value.get()
+    }
+}
+
 impl From<u128> for Pu128 {
     #[inline]
     fn from(value: u128) -> Self {

compiler/rustc_middle/src/mir/terminator.rs

@@ -67,6 +67,17 @@ impl SwitchTargets {
         &mut self.targets
     }
 
+    /// Returns a slice with all considered values (not including the fallback).
+    #[inline]
+    pub fn all_values(&self) -> &[Pu128] {
+        &self.values
+    }
+
+    #[inline]
+    pub fn all_values_mut(&mut self) -> &mut [Pu128] {
+        &mut self.values
+    }
+
     /// Finds the `BasicBlock` to which this `SwitchInt` will branch given the
     /// specific value. This cannot fail, as it'll return the `otherwise`
     /// branch if there's not a specific match for the value.

compiler/rustc_mir_transform/src/match_branches.rs

@@ -7,6 +7,7 @@ use rustc_middle::ty::layout::{IntegerExt, TyAndLayout};
 use rustc_middle::ty::{ParamEnv, ScalarInt, Ty, TyCtxt};
 use rustc_target::abi::Integer;
 use rustc_type_ir::TyKind::*;
+use tracing::instrument;
 
 use super::simplify::simplify_cfg;
 
@@ -57,7 +58,7 @@ impl<'tcx> crate::MirPass<'tcx> for MatchBranchSimplification {
 }
 
 trait SimplifyMatch<'tcx> {
-    /// Simplifies a match statement, returning true if the simplification succeeds, false
+    /// Simplifies a match statement, returning `Some` if the simplification succeeds, `None`
     /// otherwise. Generic code is written here, and we generally don't need a custom
     /// implementation.
     fn simplify(
@@ -156,6 +157,7 @@ struct SimplifyToIf;
 /// }
 /// ```
 impl<'tcx> SimplifyMatch<'tcx> for SimplifyToIf {
+    #[instrument(level = "debug", skip(self, tcx), ret)]
     fn can_simplify(
         &mut self,
         tcx: TyCtxt<'tcx>,
@@ -164,12 +166,15 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToIf {
         bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
         _discr_ty: Ty<'tcx>,
     ) -> Option<()> {
-        if targets.iter().len() != 1 {
-            return None;
-        }
+        let (first, second) = match targets.all_targets() {
+            &[first, otherwise] => (first, otherwise),
+            &[first, second, otherwise] if bbs[otherwise].is_empty_unreachable() => (first, second),
+            _ => {
+                return None;
+            }
+        };
+
         // We require that the possible target blocks all be distinct.
-        let (_, first) = targets.iter().next().unwrap();
-        let second = targets.otherwise();
         if first == second {
             return None;
         }
@@ -218,8 +223,14 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToIf {
         discr_local: Local,
         discr_ty: Ty<'tcx>,
     ) {
-        let (val, first) = targets.iter().next().unwrap();
-        let second = targets.otherwise();
+        let ((val, first), second) = match (targets.all_targets(), targets.all_values()) {
+            (&[first, otherwise], &[val]) => ((val, first), otherwise),
+            (&[first, second, otherwise], &[val, _]) if bbs[otherwise].is_empty_unreachable() => {
+                ((val, first), second)
+            }
+            _ => unreachable!(),
+        };
+
         // We already checked that first and second are different blocks,
         // and bb_idx has a different terminator from both of them.
         let first = &bbs[first];
@@ -294,7 +305,7 @@ struct SimplifyToExp {
     transform_kinds: Vec<TransformKind>,
 }
 
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Debug)]
 enum ExpectedTransformKind<'a, 'tcx> {
     /// Identical statements.
     Same(&'a StatementKind<'tcx>),
@@ -359,6 +370,7 @@ impl From<ExpectedTransformKind<'_, '_>> for TransformKind {
 /// }
 /// ```
 impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
+    #[instrument(level = "debug", skip(self, tcx), ret)]
     fn can_simplify(
         &mut self,
         tcx: TyCtxt<'tcx>,

tests/mir-opt/matches_reduce_branches.my_is_some.MatchBranchSimplification.diff

@@ -4,27 +4,32 @@
   fn my_is_some(_1: Option<T>) -> bool {
       let mut _0: bool;
       let mut _2: isize;
++     let mut _3: isize;
 
       bb0: {
           _2 = discriminant(_1);
-          switchInt(copy _2) -> [0: bb1, 1: bb2, otherwise: bb3];
-      }
-  
-      bb1: {
-          _0 = const false;
-          goto -> bb4;
-      }
-  
-      bb2: {
-          _0 = const true;
-          goto -> bb4;
-      }
-  
-      bb3: {
-          unreachable;
-      }
-  
-      bb4: {
+-         switchInt(copy _2) -> [0: bb1, 1: bb2, otherwise: bb3];
+-     }
+- 
+-     bb1: {
+-         _0 = const false;
+-         goto -> bb4;
+-     }
+- 
+-     bb2: {
+-         _0 = const true;
+-         goto -> bb4;
+-     }
+- 
+-     bb3: {
+-         unreachable;
+-     }
+- 
+-     bb4: {
++         StorageLive(_3);
++         _3 = copy _2;
++         _0 = Ne(copy _3, const 0_isize);
++         StorageDead(_3);
           return;
       }
   }

tests/mir-opt/matches_reduce_branches.rs

@@ -25,7 +25,7 @@ fn foo(bar: Option<()>) {
 #[custom_mir(dialect = "built")]
 fn my_is_some<T>(bar: Option<T>) -> bool {
     // CHECK-LABEL: fn my_is_some(
-    // CHECK: switchInt
+    // CHECK: = Ne
     // CHECK: return
     mir! {
         {