Auto merge of rust-lang#127324 - DianQK:match-br, r=<try>

[WIP] Simplify match based on the cast result of `IntToInt`

Even though the current method has become simpler and clearer, I should still add more comprehensive tests.

r? ghost

Author: bors

compiler/rustc_mir_transform/src/match_branches.rs

@@ -41,10 +41,7 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
                 should_cleanup = true;
                 continue;
             }
-            // unsound: https://github.com/rust-lang/rust/issues/124150
-            if tcx.sess.opts.unstable_opts.unsound_mir_opts
-                && SimplifyToExp::default().simplify(tcx, body, bb_idx, param_env).is_some()
-            {
+            if SimplifyToExp::default().simplify(tcx, body, bb_idx, param_env).is_some() {
                 should_cleanup = true;
                 continue;
             }
@@ -263,33 +260,49 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToIf {
     }
 }
 
+/// Check if the cast constant using `IntToInt` is equal to the target constant.
+fn can_cast(
+    from_val: impl Into<u128>,
+    from_size: Size,
+    target_scalar: ScalarInt,
+    target_is_signed: bool,
+) -> bool {
+    let from_scalar = ScalarInt::try_from_uint(from_val.into(), from_size).unwrap();
+    let to_size = target_scalar.size();
+    let cast_scalar = if target_is_signed {
+        ScalarInt::try_from_int(from_scalar.to_int(from_size), to_size).unwrap()
+    } else {
+        ScalarInt::try_from_uint(from_scalar.to_uint(from_size), to_size).unwrap()
+    };
+    cast_scalar == target_scalar
+}
+
 #[derive(Default)]
 struct SimplifyToExp {
-    transfrom_types: Vec<TransfromType>,
+    transfrom_kinds: Vec<TransfromKind>,
 }
 
 #[derive(Clone, Copy)]
-enum CompareType<'tcx, 'a> {
+enum ExpectedTransformKind<'tcx, 'a> {
     /// Identical statements.
     Same(&'a StatementKind<'tcx>),
     /// Assignment statements have the same value.
-    Eq(&'a Place<'tcx>, Ty<'tcx>, ScalarInt),
+    SameByEq { place: &'a Place<'tcx>, ty: Ty<'tcx>, scalar: ScalarInt },
     /// Enum variant comparison type.
-    Discr { place: &'a Place<'tcx>, ty: Ty<'tcx>, is_signed: bool },
+    Cast { place: &'a Place<'tcx>, ty: Ty<'tcx> },
 }
 
-enum TransfromType {
+enum TransfromKind {
     Same,
-    Eq,
-    Discr,
+    Cast,
 }
 
-impl From<CompareType<'_, '_>> for TransfromType {
-    fn from(compare_type: CompareType<'_, '_>) -> Self {
+impl From<ExpectedTransformKind<'_, '_>> for TransfromKind {
+    fn from(compare_type: ExpectedTransformKind<'_, '_>) -> Self {
         match compare_type {
-            CompareType::Same(_) => TransfromType::Same,
-            CompareType::Eq(_, _, _) => TransfromType::Eq,
-            CompareType::Discr { .. } => TransfromType::Discr,
+            ExpectedTransformKind::Same(_) => TransfromKind::Same,
+            ExpectedTransformKind::SameByEq { .. } => TransfromKind::Same,
+            ExpectedTransformKind::Cast { .. } => TransfromKind::Cast,
         }
     }
 }
@@ -353,7 +366,7 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
             return None;
         }
         let mut target_iter = targets.iter();
-        let (first_val, first_target) = target_iter.next().unwrap();
+        let (first_case_val, first_target) = target_iter.next().unwrap();
         let first_terminator_kind = &bbs[first_target].terminator().kind;
         // Check that destinations are identical, and if not, then don't optimize this block
         if !targets
@@ -365,22 +378,18 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
 
         let discr_size = tcx.layout_of(param_env.and(discr_ty)).unwrap().size;
         let first_stmts = &bbs[first_target].statements;
-        let (second_val, second_target) = target_iter.next().unwrap();
+        let (second_case_val, second_target) = target_iter.next().unwrap();
         let second_stmts = &bbs[second_target].statements;
         if first_stmts.len() != second_stmts.len() {
             return None;
         }
 
-        fn int_equal(l: ScalarInt, r: impl Into<u128>, size: Size) -> bool {
-            l.to_bits_unchecked() == ScalarInt::try_from_uint(r, size).unwrap().to_bits_unchecked()
-        }
-
         // We first compare the two branches, and then the other branches need to fulfill the same conditions.
-        let mut compare_types = Vec::new();
+        let mut expected_transform_kinds = Vec::new();
         for (f, s) in iter::zip(first_stmts, second_stmts) {
             let compare_type = match (&f.kind, &s.kind) {
                 // If two statements are exactly the same, we can optimize.
-                (f_s, s_s) if f_s == s_s => CompareType::Same(f_s),
+                (f_s, s_s) if f_s == s_s => ExpectedTransformKind::Same(f_s),
 
                 // If two statements are assignments with the match values to the same place, we can optimize.
                 (
@@ -394,22 +403,27 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
                         f_c.const_.try_eval_scalar_int(tcx, param_env),
                         s_c.const_.try_eval_scalar_int(tcx, param_env),
                     ) {
-                        (Some(f), Some(s)) if f == s => CompareType::Eq(lhs_f, f_c.const_.ty(), f),
-                        // Enum variants can also be simplified to an assignment statement if their values are equal.
-                        // We need to consider both unsigned and signed scenarios here.
+                        (Some(f), Some(s)) if f == s => ExpectedTransformKind::SameByEq {
+                            place: lhs_f,
+                            ty: f_c.const_.ty(),
+                            scalar: f,
+                        },
+                        // Enum variants can also be simplified to an assignment statement,
+                        // if we can use `IntToInt` cast to get an equal value.
                         (Some(f), Some(s))
-                            if ((f_c.const_.ty().is_signed() || discr_ty.is_signed())
-                                && int_equal(f, first_val, discr_size)
-                                && int_equal(s, second_val, discr_size))
-                                || (Some(f) == ScalarInt::try_from_uint(first_val, f.size())
-                                    && Some(s)
-                                        == ScalarInt::try_from_uint(second_val, s.size())) =>
+                            if (can_cast(
+                                first_case_val,
+                                discr_size,
+                                f,
+                                f_c.const_.ty().is_signed(),
+                            ) && can_cast(
+                                second_case_val,
+                                discr_size,
+                                s,
+                                s_c.const_.ty().is_signed(),
+                            )) =>
                         {
-                            CompareType::Discr {
-                                place: lhs_f,
-                                ty: f_c.const_.ty(),
-                                is_signed: f_c.const_.ty().is_signed() || discr_ty.is_signed(),
-                            }
+                            ExpectedTransformKind::Cast { place: lhs_f, ty: f_c.const_.ty() }
                         }
                         _ => {
                             return None;
@@ -420,47 +434,36 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
                 // Otherwise we cannot optimize. Try another block.
                 _ => return None,
             };
-            compare_types.push(compare_type);
+            expected_transform_kinds.push(compare_type);
         }
 
         // All remaining BBs need to fulfill the same pattern as the two BBs from the previous step.
         for (other_val, other_target) in target_iter {
             let other_stmts = &bbs[other_target].statements;
-            if compare_types.len() != other_stmts.len() {
+            if expected_transform_kinds.len() != other_stmts.len() {
                 return None;
             }
-            for (f, s) in iter::zip(&compare_types, other_stmts) {
+            for (f, s) in iter::zip(&expected_transform_kinds, other_stmts) {
                 match (*f, &s.kind) {
-                    (CompareType::Same(f_s), s_s) if f_s == s_s => {}
+                    (ExpectedTransformKind::Same(f_s), s_s) if f_s == s_s => {}
                     (
-                        CompareType::Eq(lhs_f, f_ty, val),
+                        ExpectedTransformKind::SameByEq { place: lhs_f, ty: f_ty, scalar },
                         StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))),
                     ) if lhs_f == lhs_s
                         && s_c.const_.ty() == f_ty
-                        && s_c.const_.try_eval_scalar_int(tcx, param_env) == Some(val) => {}
+                        && s_c.const_.try_eval_scalar_int(tcx, param_env) == Some(scalar) => {}
                     (
-                        CompareType::Discr { place: lhs_f, ty: f_ty, is_signed },
+                        ExpectedTransformKind::Cast { place: lhs_f, ty: f_ty },
                         StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))),
-                    ) if lhs_f == lhs_s && s_c.const_.ty() == f_ty => {
-                        let Some(f) = s_c.const_.try_eval_scalar_int(tcx, param_env) else {
-                            return None;
-                        };
-                        if is_signed
-                            && s_c.const_.ty().is_signed()
-                            && int_equal(f, other_val, discr_size)
-                        {
-                            continue;
-                        }
-                        if Some(f) == ScalarInt::try_from_uint(other_val, f.size()) {
-                            continue;
-                        }
-                        return None;
-                    }
+                    ) if let Some(f) = s_c.const_.try_eval_scalar_int(tcx, param_env)
+                        && lhs_f == lhs_s
+                        && s_c.const_.ty() == f_ty
+                        && can_cast(other_val, discr_size, f, f_ty.is_signed()) => {}
                     _ => return None,
                 }
             }
         }
-        self.transfrom_types = compare_types.into_iter().map(|c| c.into()).collect();
+        self.transfrom_kinds = expected_transform_kinds.into_iter().map(|c| c.into()).collect();
         Some(())
     }
 
@@ -478,13 +481,13 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
         let (_, first) = targets.iter().next().unwrap();
         let first = &bbs[first];
 
-        for (t, s) in iter::zip(&self.transfrom_types, &first.statements) {
+        for (t, s) in iter::zip(&self.transfrom_kinds, &first.statements) {
             match (t, &s.kind) {
-                (TransfromType::Same, _) | (TransfromType::Eq, _) => {
+                (TransfromKind::Same, _) => {
                     patch.add_statement(parent_end, s.kind.clone());
                 }
                 (
-                    TransfromType::Discr,
+                    TransfromKind::Cast,
                     StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))),
                 ) => {
                     let operand = Operand::Copy(Place::from(discr_local));

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

@@ -5,37 +5,42 @@
       debug i => _1;
       let mut _0: u128;
       let mut _2: i128;
++     let mut _3: i128;
 
       bb0: {
           _2 = discriminant(_1);
-          switchInt(move _2) -> [1: bb3, 2: bb4, 3: bb5, 340282366920938463463374607431768211455: bb2, otherwise: bb1];
-      }
-  
-      bb1: {
-          unreachable;
-      }
-  
-      bb2: {
-          _0 = const core::num::<impl u128>::MAX;
-          goto -> bb6;
-      }
-  
-      bb3: {
-          _0 = const 1_u128;
-          goto -> bb6;
-      }
-  
-      bb4: {
-          _0 = const 2_u128;
-          goto -> bb6;
-      }
-  
-      bb5: {
-          _0 = const 3_u128;
-          goto -> bb6;
-      }
-  
-      bb6: {
+-         switchInt(move _2) -> [1: bb3, 2: bb4, 3: bb5, 340282366920938463463374607431768211455: bb2, otherwise: bb1];
+-     }
+- 
+-     bb1: {
+-         unreachable;
+-     }
+- 
+-     bb2: {
+-         _0 = const core::num::<impl u128>::MAX;
+-         goto -> bb6;
+-     }
+- 
+-     bb3: {
+-         _0 = const 1_u128;
+-         goto -> bb6;
+-     }
+- 
+-     bb4: {
+-         _0 = const 2_u128;
+-         goto -> bb6;
+-     }
+- 
+-     bb5: {
+-         _0 = const 3_u128;
+-         goto -> bb6;
+-     }
+- 
+-     bb6: {
++         StorageLive(_3);
++         _3 = move _2;
++         _0 = _3 as u128 (IntToInt);
++         StorageDead(_3);
           return;
       }
   }

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

@@ -5,32 +5,37 @@
       debug i => _1;
       let mut _0: i8;
       let mut _2: i16;
++     let mut _3: i16;
 
       bb0: {
           _2 = discriminant(_1);
-          switchInt(move _2) -> [65535: bb3, 2: bb4, 65533: bb2, otherwise: bb1];
-      }
-  
-      bb1: {
-          unreachable;
-      }
-  
-      bb2: {
-          _0 = const -3_i8;
-          goto -> bb5;
-      }
-  
-      bb3: {
-          _0 = const -1_i8;
-          goto -> bb5;
-      }
-  
-      bb4: {
-          _0 = const 2_i8;
-          goto -> bb5;
-      }
-  
-      bb5: {
+-         switchInt(move _2) -> [65535: bb3, 2: bb4, 65533: bb2, otherwise: bb1];
+-     }
+- 
+-     bb1: {
+-         unreachable;
+-     }
+- 
+-     bb2: {
+-         _0 = const -3_i8;
+-         goto -> bb5;
+-     }
+- 
+-     bb3: {
+-         _0 = const -1_i8;
+-         goto -> bb5;
+-     }
+- 
+-     bb4: {
+-         _0 = const 2_i8;
+-         goto -> bb5;
+-     }
+- 
+-     bb5: {
++         StorageLive(_3);
++         _3 = move _2;
++         _0 = _3 as i8 (IntToInt);
++         StorageDead(_3);
           return;
       }
   }