Transforms a match containing negative numbers into an assignment statement as well

dianqk · dianqk · commit 7a476356b92b · 2024-02-04T08:06:42.000+08:00
diff --git a/compiler/rustc_mir_transform/src/match_branches.rs b/compiler/rustc_mir_transform/src/match_branches.rs
@@ -65,13 +65,13 @@ trait SimplifyMatch<'tcx> {
             _ => unreachable!(),
         };
 
-        if !self.can_simplify(tcx, targets, param_env, bbs) {
+        let discr_ty = discr.ty(local_decls, tcx);
+        if !self.can_simplify(tcx, targets, param_env, bbs, discr_ty) {
             return false;
         }
 
         // Take ownership of items now that we know we can optimize.
         let discr = discr.clone();
-        let discr_ty = discr.ty(local_decls, tcx);
 
         // Introduce a temporary for the discriminant value.
         let source_info = bbs[switch_bb_idx].terminator().source_info;
@@ -101,6 +101,7 @@ trait SimplifyMatch<'tcx> {
         targets: &SwitchTargets,
         param_env: ParamEnv<'tcx>,
         bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+        discr_ty: Ty<'tcx>,
     ) -> bool;
 
     fn new_stmts(
@@ -154,6 +155,7 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToIf {
         targets: &SwitchTargets,
         param_env: ParamEnv<'tcx>,
         bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+        _discr_ty: Ty<'tcx>,
     ) -> bool {
         if targets.iter().len() != 1 {
             return false;
@@ -265,7 +267,7 @@ struct SimplifyToExp {
 enum CompareType<'tcx, 'a> {
     Same(&'a StatementKind<'tcx>),
     Eq(&'a Place<'tcx>, Ty<'tcx>, ScalarInt),
-    Discr(&'a Place<'tcx>, Ty<'tcx>),
+    Discr(&'a Place<'tcx>, Ty<'tcx>, bool),
 }
 
 enum TransfromType {
@@ -279,7 +281,7 @@ impl From<CompareType<'_, '_>> for TransfromType {
         match compare_type {
             CompareType::Same(_) => TransfromType::Same,
             CompareType::Eq(_, _, _) => TransfromType::Eq,
-            CompareType::Discr(_, _) => TransfromType::Discr,
+            CompareType::Discr(_, _, _) => TransfromType::Discr,
         }
     }
 }
@@ -330,6 +332,7 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
         targets: &SwitchTargets,
         param_env: ParamEnv<'tcx>,
         bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+        discr_ty: Ty<'tcx>,
     ) -> bool {
         if targets.iter().len() < 2 || targets.iter().len() > 64 {
             return false;
@@ -352,6 +355,7 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
             return false;
         }
 
+        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) = iter.next().unwrap();
         let second_stmts = &bbs[second_target].statements;
@@ -379,12 +383,30 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
                     ) {
                         (Some(f), Some(s)) if f == s => CompareType::Eq(lhs_f, f_c.const_.ty(), f),
                         (Some(f), Some(s))
-                            if Some(f) == ScalarInt::try_from_uint(first_val, f.size())
-                                && Some(s) == ScalarInt::try_from_uint(second_val, s.size()) =>
+                            if ((f_c.const_.ty().is_signed() || discr_ty.is_signed())
+                                && f.try_to_int(f.size()).unwrap()
+                                    == ScalarInt::try_from_uint(first_val, discr_size)
+                                        .unwrap()
+                                        .try_to_int(discr_size)
+                                        .unwrap()
+                                && s.try_to_int(s.size()).unwrap()
+                                    == ScalarInt::try_from_uint(second_val, discr_size)
+                                        .unwrap()
+                                        .try_to_int(discr_size)
+                                        .unwrap())
+                                || (Some(f) == ScalarInt::try_from_uint(first_val, f.size())
+                                    && Some(s)
+                                        == ScalarInt::try_from_uint(second_val, s.size())) =>
                         {
-                            CompareType::Discr(lhs_f, f_c.const_.ty())
+                            CompareType::Discr(
+                                lhs_f,
+                                f_c.const_.ty(),
+                                f_c.const_.ty().is_signed() || discr_ty.is_signed(),
+                            )
+                        }
+                        _ => {
+                            return false;
                         }
-                        _ => return false,
                     }
                 }
 
@@ -409,15 +431,26 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
                         && s_c.const_.ty() == f_ty
                         && s_c.const_.try_eval_scalar_int(tcx, param_env) == Some(val) => {}
                     (
-                        CompareType::Discr(lhs_f, f_ty),
+                        CompareType::Discr(lhs_f, f_ty, is_signed),
                         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 false;
                         };
-                        if Some(f) != ScalarInt::try_from_uint(other_val, f.size()) {
-                            return false;
+                        if is_signed
+                            && s_c.const_.ty().is_signed()
+                            && f.try_to_int(f.size()).unwrap()
+                                == ScalarInt::try_from_uint(other_val, discr_size)
+                                    .unwrap()
+                                    .try_to_int(discr_size)
+                                    .unwrap()
+                        {
+                            continue;
+                        }
+                        if Some(f) == ScalarInt::try_from_uint(other_val, f.size()) {
+                            continue;
                         }
+                        return false;
                     }
                     _ => return false,
                 }
diff --git a/tests/mir-opt/matches_reduce_branches.match_i16_i8.MatchBranchSimplification.diff b/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: bb1, otherwise: bb2];
-      }
-  
-      bb1: {
-          _0 = const -3_i8;
-          goto -> bb5;
-      }
-  
-      bb2: {
-          unreachable;
-      }
-  
-      bb3: {
-          _0 = const -1_i8;
-          goto -> bb5;
-      }
-  
-      bb4: {
-          _0 = const 2_i8;
-          goto -> bb5;
-      }
-  
-      bb5: {
+-         switchInt(move _2) -> [65535: bb3, 2: bb4, 65533: bb1, otherwise: bb2];
+-     }
+- 
+-     bb1: {
+-         _0 = const -3_i8;
+-         goto -> bb5;
+-     }
+- 
+-     bb2: {
+-         unreachable;
+-     }
+- 
+-     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;
       }
   }
diff --git a/tests/mir-opt/matches_reduce_branches.match_i8_i16.MatchBranchSimplification.diff b/tests/mir-opt/matches_reduce_branches.match_i8_i16.MatchBranchSimplification.diff
@@ -5,32 +5,37 @@
       debug i => _1;
       let mut _0: i16;
       let mut _2: i8;
++     let mut _3: i8;
   
       bb0: {
           _2 = discriminant(_1);
-          switchInt(move _2) -> [255: bb3, 2: bb4, 253: bb1, otherwise: bb2];
-      }
-  
-      bb1: {
-          _0 = const -3_i16;
-          goto -> bb5;
-      }
-  
-      bb2: {
-          unreachable;
-      }
-  
-      bb3: {
-          _0 = const -1_i16;
-          goto -> bb5;
-      }
-  
-      bb4: {
-          _0 = const 2_i16;
-          goto -> bb5;
-      }
-  
-      bb5: {
+-         switchInt(move _2) -> [255: bb3, 2: bb4, 253: bb1, otherwise: bb2];
+-     }
+- 
+-     bb1: {
+-         _0 = const -3_i16;
+-         goto -> bb5;
+-     }
+- 
+-     bb2: {
+-         unreachable;
+-     }
+- 
+-     bb3: {
+-         _0 = const -1_i16;
+-         goto -> bb5;
+-     }
+- 
+-     bb4: {
+-         _0 = const 2_i16;
+-         goto -> bb5;
+-     }
+- 
+-     bb5: {
++         StorageLive(_3);
++         _3 = move _2;
++         _0 = _3 as i16 (IntToInt);
++         StorageDead(_3);
           return;
       }
   }
diff --git a/tests/mir-opt/matches_reduce_branches.rs b/tests/mir-opt/matches_reduce_branches.rs
@@ -144,7 +144,9 @@ enum EnumAi8 {
 // EMIT_MIR matches_reduce_branches.match_i8_i16.MatchBranchSimplification.diff
 fn match_i8_i16(i: EnumAi8) -> i16 {
     // CHECK-LABEL: fn match_i8_i16(
-    // CHECK: switchInt
+    // CHECK-NOT: switchInt
+    // CHECK: _0 = _3 as i16 (IntToInt);
+    // CHECH: return
     match i {
         EnumAi8::A => -1,
         EnumAi8::B => 2,
@@ -173,7 +175,9 @@ enum EnumAi16 {
 // EMIT_MIR matches_reduce_branches.match_i16_i8.MatchBranchSimplification.diff
 fn match_i16_i8(i: EnumAi16) -> i8 {
     // CHECK-LABEL: fn match_i16_i8(
-    // CHECK: switchInt
+    // CHECK-NOT: switchInt
+    // CHECK: _0 = _3 as i8 (IntToInt);
+    // CHECH: return
     match i {
         EnumAi16::A => -1,
         EnumAi16::B => 2,
