Apply EarlyOtherwiseBranch for all targets with the same value

dianqk · dianqk · commit 6658b16e9523 · 2024-02-23T23:45:01.000+08:00
diff --git a/compiler/rustc_mir_transform/src/early_otherwise_branch.rs b/compiler/rustc_mir_transform/src/early_otherwise_branch.rs
@@ -159,27 +159,12 @@ impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch {
                 (None, Operand::Copy(opt_data.child_place))
             };
 
-            // create temp to store inequality comparison between the two discriminants, `_t` in
-            // example above
-            let nequal = BinOp::Ne;
-            let comp_res_type = nequal.ty(tcx, parent_ty, opt_data.child_ty);
-            let comp_temp = patch.new_temp(comp_res_type, opt_data.child_source.span);
-            patch.add_statement(parent_end, StatementKind::StorageLive(comp_temp));
-
-            // create inequality comparison between the two discriminants
-            let comp_rvalue =
-                Rvalue::BinaryOp(nequal, Box::new((parent_op.clone(), second_operand)));
-            patch.add_statement(
-                parent_end,
-                StatementKind::Assign(Box::new((Place::from(comp_temp), comp_rvalue))),
-            );
-
             let eq_new_targets = parent_targets.iter().map(|(value, child)| {
                 let TerminatorKind::SwitchInt { targets, .. } = &bbs[child].terminator().kind
                 else {
                     unreachable!()
                 };
-                (value, targets.target_for_value(value))
+                (value, targets.all_targets()[0])
             });
             let eq_targets = SwitchTargets::new(eq_new_targets, opt_data.destination);
 
@@ -188,36 +173,77 @@ impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch {
                 source_info: bbs[parent].terminator().source_info,
                 kind: TerminatorKind::SwitchInt {
                     // switch on the first discriminant, so we can mark the second one as dead
-                    discr: parent_op,
+                    discr: parent_op.clone(),
                     targets: eq_targets,
                 },
             }));
 
             let eq_bb = patch.new_block(eq_switch);
 
-            // Jump to it on the basis of the inequality comparison
-            let true_case = opt_data.destination;
-            let false_case = eq_bb;
-            patch.patch_terminator(
-                parent,
-                TerminatorKind::if_(Operand::Move(Place::from(comp_temp)), true_case, false_case),
-            );
+            if let Some(same_target_value) = opt_data.same_target_value {
+                let t = TerminatorKind::SwitchInt {
+                    discr: second_operand,
+                    targets: SwitchTargets::static_if(
+                        same_target_value,
+                        eq_bb,
+                        opt_data.destination,
+                    ),
+                };
+                patch.patch_terminator(parent, t);
+
+                if let Some(second_discriminant_temp) = second_discriminant_temp {
+                    // Generate a StorageDead for second_discriminant_temp in each of the targets, since we moved it into
+                    // the switch
+                    for bb in [eq_bb, opt_data.destination].iter() {
+                        patch.add_statement(
+                            Location { block: *bb, statement_index: 0 },
+                            StatementKind::StorageDead(second_discriminant_temp),
+                        );
+                    }
+                }
+            } else {
+                // create temp to store inequality comparison between the two discriminants, `_t` in
+                // example above
+                let nequal = BinOp::Ne;
+                let comp_res_type = nequal.ty(tcx, parent_ty, opt_data.child_ty);
+                let comp_temp = patch.new_temp(comp_res_type, opt_data.child_source.span);
+                patch.add_statement(parent_end, StatementKind::StorageLive(comp_temp));
 
-            if let Some(second_discriminant_temp) = second_discriminant_temp {
-                // generate StorageDead for the second_discriminant_temp not in use anymore
+                // create inequality comparison between the two discriminants
+                let comp_rvalue = Rvalue::BinaryOp(nequal, Box::new((parent_op, second_operand)));
                 patch.add_statement(
                     parent_end,
-                    StatementKind::StorageDead(second_discriminant_temp),
+                    StatementKind::Assign(Box::new((Place::from(comp_temp), comp_rvalue))),
                 );
-            }
 
-            // Generate a StorageDead for comp_temp in each of the targets, since we moved it into
-            // the switch
-            for bb in [false_case, true_case].iter() {
-                patch.add_statement(
-                    Location { block: *bb, statement_index: 0 },
-                    StatementKind::StorageDead(comp_temp),
+                // Jump to it on the basis of the inequality comparison
+                let true_case = opt_data.destination;
+                let false_case = eq_bb;
+                patch.patch_terminator(
+                    parent,
+                    TerminatorKind::if_(
+                        Operand::Move(Place::from(comp_temp)),
+                        true_case,
+                        false_case,
+                    ),
                 );
+
+                // Generate a StorageDead for comp_temp in each of the targets, since we moved it into
+                // the switch
+                for bb in [false_case, true_case].iter() {
+                    patch.add_statement(
+                        Location { block: *bb, statement_index: 0 },
+                        StatementKind::StorageDead(comp_temp),
+                    );
+                }
+
+                if let Some(second_discriminant_temp) = second_discriminant_temp {
+                    // generate StorageDead for the second_discriminant_temp not in use anymore
+                    patch.add_statement(
+                        parent_end,
+                        StatementKind::StorageDead(second_discriminant_temp),
+                    );
+                }
             }
 
             patch.apply(body);
@@ -286,6 +312,7 @@ struct OptimizationData<'tcx> {
     child_ty: Ty<'tcx>,
     child_source: SourceInfo,
     hoist_discriminant: bool,
+    same_target_value: Option<u128>,
 }
 
 fn evaluate_candidate<'tcx>(
@@ -375,16 +402,38 @@ fn evaluate_candidate<'tcx>(
         targets.otherwise()
     };
 
+    let TerminatorKind::SwitchInt { targets: child_targets, .. } = &bbs[child].terminator().kind
+    else {
+        return None;
+    };
     // Verify that the optimization is legal for each branch
-    for (value, child) in targets.iter() {
+    let Some((may_same_target_value, _)) = child_targets.iter().next() else {
+        return None;
+    };
+    let mut same_target_value = Some(may_same_target_value);
+    for (_, child) in targets.iter() {
         if !verify_candidate_branch(
             &bbs[child],
-            value,
+            may_same_target_value,
             child_place,
             destination,
             hoist_discriminant,
         ) {
-            return None;
+            same_target_value = None;
+            break;
+        }
+    }
+    if same_target_value.is_none() {
+        for (value, child) in targets.iter() {
+            if !verify_candidate_branch(
+                &bbs[child],
+                value,
+                child_place,
+                destination,
+                hoist_discriminant,
+            ) {
+                return None;
+            }
         }
     }
     Some(OptimizationData {
@@ -393,6 +442,7 @@ fn evaluate_candidate<'tcx>(
         child_ty,
         child_source: child_terminator.source_info,
         hoist_discriminant,
+        same_target_value,
     })
 }
 
diff --git a/tests/mir-opt/early_otherwise_branch.opt6.EarlyOtherwiseBranch.diff b/tests/mir-opt/early_otherwise_branch.opt6.EarlyOtherwiseBranch.diff
@@ -0,0 +1,72 @@
+- // MIR for `opt6` before EarlyOtherwiseBranch
++ // MIR for `opt6` after EarlyOtherwiseBranch
+  
+  fn opt6(_1: u32, _2: u32) -> u32 {
+      debug x => _1;
+      debug y => _2;
+      let mut _0: u32;
+      let mut _3: (u32, u32);
+      let mut _4: u32;
+      let mut _5: u32;
+  
+      bb0: {
+          StorageLive(_3);
+          StorageLive(_4);
+          _4 = _1;
+          StorageLive(_5);
+          _5 = _2;
+          _3 = (move _4, move _5);
+          StorageDead(_5);
+          StorageDead(_4);
+-         switchInt((_3.0: u32)) -> [1: bb2, 2: bb3, 3: bb4, otherwise: bb1];
++         switchInt((_3.1: u32)) -> [10: bb6, otherwise: bb1];
+      }
+  
+      bb1: {
+          _0 = const 0_u32;
+-         goto -> bb8;
++         goto -> bb5;
+      }
+  
+      bb2: {
+-         switchInt((_3.1: u32)) -> [10: bb5, otherwise: bb1];
++         _0 = const 4_u32;
++         goto -> bb5;
+      }
+  
+      bb3: {
+-         switchInt((_3.1: u32)) -> [10: bb6, otherwise: bb1];
++         _0 = const 5_u32;
++         goto -> bb5;
+      }
+  
+      bb4: {
+-         switchInt((_3.1: u32)) -> [10: bb7, otherwise: bb1];
++         _0 = const 6_u32;
++         goto -> bb5;
+      }
+  
+      bb5: {
+-         _0 = const 4_u32;
+-         goto -> bb8;
++         StorageDead(_3);
++         return;
+      }
+  
+      bb6: {
+-         _0 = const 5_u32;
+-         goto -> bb8;
+-     }
+- 
+-     bb7: {
+-         _0 = const 6_u32;
+-         goto -> bb8;
+-     }
+- 
+-     bb8: {
+-         StorageDead(_3);
+-         return;
++         switchInt((_3.0: u32)) -> [1: bb2, 2: bb3, 3: bb4, otherwise: bb1];
+      }
+  }
+  
diff --git a/tests/mir-opt/early_otherwise_branch.opt7.EarlyOtherwiseBranch.diff b/tests/mir-opt/early_otherwise_branch.opt7.EarlyOtherwiseBranch.diff
@@ -0,0 +1,94 @@
+- // MIR for `opt7` before EarlyOtherwiseBranch
++ // MIR for `opt7` after EarlyOtherwiseBranch
+  
+  fn opt7(_1: Option<u32>, _2: Option<u32>) -> u32 {
+      debug x => _1;
+      debug y => _2;
+      let mut _0: u32;
+      let mut _3: (std::option::Option<u32>, std::option::Option<u32>);
+      let mut _4: std::option::Option<u32>;
+      let mut _5: std::option::Option<u32>;
+      let mut _6: isize;
+      let mut _7: isize;
+      let mut _8: isize;
+      let _9: u32;
+      let _10: u32;
+      let _11: u32;
++     let mut _12: isize;
+      scope 1 {
+          debug a => _9;
+          debug b => _10;
+      }
+      scope 2 {
+          debug b => _11;
+      }
+  
+      bb0: {
+          StorageLive(_3);
+          StorageLive(_4);
+          _4 = _1;
+          StorageLive(_5);
+          _5 = _2;
+          _3 = (move _4, move _5);
+          StorageDead(_5);
+          StorageDead(_4);
+          _8 = discriminant((_3.0: std::option::Option<u32>));
+-         switchInt(move _8) -> [0: bb2, 1: bb3, otherwise: bb7];
++         StorageLive(_12);
++         _12 = discriminant((_3.1: std::option::Option<u32>));
++         switchInt(move _12) -> [1: bb5, otherwise: bb1];
+      }
+  
+      bb1: {
++         StorageDead(_12);
+          _0 = const 1_u32;
+-         goto -> bb6;
++         goto -> bb4;
+      }
+  
+      bb2: {
+-         _6 = discriminant((_3.1: std::option::Option<u32>));
+-         switchInt(move _6) -> [1: bb5, otherwise: bb1];
+-     }
+- 
+-     bb3: {
+-         _7 = discriminant((_3.1: std::option::Option<u32>));
+-         switchInt(move _7) -> [1: bb4, otherwise: bb1];
+-     }
+- 
+-     bb4: {
+          StorageLive(_10);
+          _10 = (((_3.1: std::option::Option<u32>) as Some).0: u32);
+          StorageLive(_9);
+          _9 = (((_3.0: std::option::Option<u32>) as Some).0: u32);
+          _0 = const 0_u32;
+          StorageDead(_9);
+          StorageDead(_10);
+-         goto -> bb6;
++         goto -> bb4;
+      }
+  
+-     bb5: {
++     bb3: {
+          StorageLive(_11);
+          _11 = (((_3.1: std::option::Option<u32>) as Some).0: u32);
+          _0 = const 2_u32;
+          StorageDead(_11);
+-         goto -> bb6;
++         goto -> bb4;
+      }
+  
+-     bb6: {
++     bb4: {
+          StorageDead(_3);
+          return;
+      }
+  
+-     bb7: {
+-         unreachable;
++     bb5: {
++         StorageDead(_12);
++         switchInt(_8) -> [0: bb3, 1: bb2, otherwise: bb1];
+      }
+  }
+  
diff --git a/tests/mir-opt/early_otherwise_branch.rs b/tests/mir-opt/early_otherwise_branch.rs
@@ -83,10 +83,44 @@ fn opt5(x: u32, y: u32) -> u32 {
     }
 }
 
+// EMIT_MIR early_otherwise_branch.opt6.EarlyOtherwiseBranch.diff
+fn opt6(x: u32, y: u32) -> u32 {
+    // CHECK-LABEL: fn opt6(
+    // CHECK: bb0: {
+    // CHECK: switchInt((_{{.*}}: u32)) -> [10: [[SWITCH_BB:bb.*]], otherwise: [[OTHERWISE:bb.*]]];
+    // CHECK-NEXT: }
+    // CHECK: [[SWITCH_BB]]:
+    // CHECK:  switchInt((_{{.*}}: u32)) -> [1: bb{{.*}}, 2: bb{{.*}}, 3: bb{{.*}}, otherwise: [[OTHERWISE]]];
+    // CHECK-NEXT: }
+    match (x, y) {
+        (1, 10) => 4,
+        (2, 10) => 5,
+        (3, 10) => 6,
+        _ => 0,
+    }
+}
+
+// EMIT_MIR early_otherwise_branch.opt7.EarlyOtherwiseBranch.diff
+fn opt7(x: Option<u32>, y: Option<u32>) -> u32 {
+    // CHECK-LABEL: fn opt7(
+    // CHECK: bb0: {
+    // CHECK: [[LOCAL1:_.*]] = discriminant({{.*}});
+    // CHECK: [[LOCAL2:_.*]] = discriminant({{.*}});
+    // CHECK: switchInt(move [[LOCAL2]]) -> [
+    // CHECK-NEXT: }
+    match (x, y) {
+        (Some(a), Some(b)) => 0,
+        (None, Some(b)) => 2,
+        _ => 1,
+    }
+}
+
 fn main() {
     opt1(None, Some(0));
     opt2(None, Some(0));
     opt3(None, Some(false));
     opt4(0, 0);
     opt5(0, 0);
+    opt6(0, 0);
+    opt7(None, Some(0));
 }
