Auto merge of #129047 - DianQK:early_otherwise_branch_scalar, r=cjgillot

Apply `EarlyOtherwiseBranch` to scalar value

In the future, I'm thinking of hoisting discriminant via GVN so that we only need to write very little code here.

r? `@cjgillot`

compiler/rustc_mir_transform/src/early_otherwise_branch.rs

@@ -133,18 +133,29 @@ fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
 
             let mut patch = MirPatch::new(body);
 
-            // create temp to store second discriminant in, `_s` in example above
+            let (second_discriminant_temp, second_operand) = if opt_data.need_hoist_discriminant {
-            let second_discriminant_temp =
+                // create temp to store second discriminant in, `_s` in example above
-                patch.new_temp(opt_data.child_ty, opt_data.child_source.span);
+                let second_discriminant_temp =
+                    patch.new_temp(opt_data.child_ty, opt_data.child_source.span);
 
-            patch.add_statement(parent_end, StatementKind::StorageLive(second_discriminant_temp));
+                patch.add_statement(
+                    parent_end,
+                    StatementKind::StorageLive(second_discriminant_temp),
+                );
 
-            // create assignment of discriminant
+                // create assignment of discriminant
-            patch.add_assign(
+                patch.add_assign(
-                parent_end,
+                    parent_end,
-                Place::from(second_discriminant_temp),
+                    Place::from(second_discriminant_temp),
-                Rvalue::Discriminant(opt_data.child_place),
+                    Rvalue::Discriminant(opt_data.child_place),
-            );
+                );
+                (
+                    Some(second_discriminant_temp),
+                    Operand::Move(Place::from(second_discriminant_temp)),
+                )
+            } else {
+                (None, Operand::Copy(opt_data.child_place))
+            };
 
             // create temp to store inequality comparison between the two discriminants, `_t` in
             // example above
@@ -153,11 +164,9 @@ fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             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
+            // create inequality comparison
-            let comp_rvalue = Rvalue::BinaryOp(
+            let comp_rvalue =
-                nequal,
+                Rvalue::BinaryOp(nequal, Box::new((parent_op.clone(), second_operand)));
-                Box::new((parent_op.clone(), Operand::Move(Place::from(second_discriminant_temp)))),
-            );
             patch.add_statement(
                 parent_end,
                 StatementKind::Assign(Box::new((Place::from(comp_temp), comp_rvalue))),
@@ -193,8 +202,13 @@ fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
                 TerminatorKind::if_(Operand::Move(Place::from(comp_temp)), true_case, false_case),
             );
 
-            // generate StorageDead for the second_discriminant_temp not in use anymore
+            if let Some(second_discriminant_temp) = second_discriminant_temp {
-            patch.add_statement(parent_end, StatementKind::StorageDead(second_discriminant_temp));
+                // generate StorageDead for the second_discriminant_temp not in use anymore
+                patch.add_statement(
+                    parent_end,
+                    StatementKind::StorageDead(second_discriminant_temp),
+                );
+            }
 
             // Generate a StorageDead for comp_temp in each of the targets, since we moved it into
             // the switch
@@ -222,6 +236,7 @@ struct OptimizationData<'tcx> {
     child_place: Place<'tcx>,
     child_ty: Ty<'tcx>,
     child_source: SourceInfo,
+    need_hoist_discriminant: bool,
 }
 
 fn evaluate_candidate<'tcx>(
@@ -235,44 +250,12 @@ fn evaluate_candidate<'tcx>(
         return None;
     };
     let parent_ty = parent_discr.ty(body.local_decls(), tcx);
-    if !bbs[targets.otherwise()].is_empty_unreachable() {
-        // Someone could write code like this:
-        // ```rust
-        // let Q = val;
-        // if discriminant(P) == otherwise {
-        //     let ptr = &mut Q as *mut _ as *mut u8;
-        //     // It may be difficult for us to effectively determine whether values are valid.
-        //     // Invalid values can come from all sorts of corners.
-        //     unsafe { *ptr = 10; }
-        // }
-        //
-        // match P {
-        //    A => match Q {
-        //        A => {
-        //            // code
-        //        }
-        //        _ => {
-        //            // don't use Q
-        //        }
-        //    }
-        //    _ => {
-        //        // don't use Q
-        //    }
-        // };
-        // ```
-        //
-        // Hoisting the `discriminant(Q)` out of the `A` arm causes us to compute the discriminant
-        // of an invalid value, which is UB.
-        // In order to fix this, **we would either need to show that the discriminant computation of
-        // `place` is computed in all branches**.
-        // FIXME(#95162) For the moment, we adopt a conservative approach and
-        // consider only the `otherwise` branch has no statements and an unreachable terminator.
-        return None;
-    }
     let (_, child) = targets.iter().next()?;
-    let child_terminator = &bbs[child].terminator();
+
-    let TerminatorKind::SwitchInt { targets: child_targets, discr: child_discr } =
+    let Terminator {
-        &child_terminator.kind
+        kind: TerminatorKind::SwitchInt { targets: child_targets, discr: child_discr },
+        source_info,
+    } = bbs[child].terminator()
     else {
         return None;
     };
@@ -280,25 +263,115 @@ fn evaluate_candidate<'tcx>(
     if child_ty != parent_ty {
         return None;
     }
-    let Some(StatementKind::Assign(boxed)) = &bbs[child].statements.first().map(|x| &x.kind) else {
+
+    // We only handle:
+    // ```
+    // bb4: {
+    //     _8 = discriminant((_3.1: Enum1));
+    //    switchInt(move _8) -> [2: bb7, otherwise: bb1];
+    // }
+    // ```
+    // and
+    // ```
+    // bb2: {
+    //     switchInt((_3.1: u64)) -> [1: bb5, otherwise: bb1];
+    // }
+    // ```
+    if bbs[child].statements.len() > 1 {
         return None;
+    }
+
+    // When thie BB has exactly one statement, this statement should be discriminant.
+    let need_hoist_discriminant = bbs[child].statements.len() == 1;
+    let child_place = if need_hoist_discriminant {
+        if !bbs[targets.otherwise()].is_empty_unreachable() {
+            // Someone could write code like this:
+            // ```rust
+            // let Q = val;
+            // if discriminant(P) == otherwise {
+            //     let ptr = &mut Q as *mut _ as *mut u8;
+            //     // It may be difficult for us to effectively determine whether values are valid.
+            //     // Invalid values can come from all sorts of corners.
+            //     unsafe { *ptr = 10; }
+            // }
+            //
+            // match P {
+            //    A => match Q {
+            //        A => {
+            //            // code
+            //        }
+            //        _ => {
+            //            // don't use Q
+            //        }
+            //    }
+            //    _ => {
+            //        // don't use Q
+            //    }
+            // };
+            // ```
+            //
+            // Hoisting the `discriminant(Q)` out of the `A` arm causes us to compute the discriminant of an
+            // invalid value, which is UB.
+            // In order to fix this, **we would either need to show that the discriminant computation of
+            // `place` is computed in all branches**.
+            // FIXME(#95162) For the moment, we adopt a conservative approach and
+            // consider only the `otherwise` branch has no statements and an unreachable terminator.
+            return None;
+        }
+        // Handle:
+        // ```
+        // bb4: {
+        //     _8 = discriminant((_3.1: Enum1));
+        //    switchInt(move _8) -> [2: bb7, otherwise: bb1];
+        // }
+        // ```
+        let [
+            Statement {
+                kind: StatementKind::Assign(box (_, Rvalue::Discriminant(child_place))),
+                ..
+            },
+        ] = bbs[child].statements.as_slice()
+        else {
+            return None;
+        };
+        *child_place
+    } else {
+        // Handle:
+        // ```
+        // bb2: {
+        //     switchInt((_3.1: u64)) -> [1: bb5, otherwise: bb1];
+        // }
+        // ```
+        let Operand::Copy(child_place) = child_discr else {
+            return None;
+        };
+        *child_place
     };
-    let (_, Rvalue::Discriminant(child_place)) = &**boxed else {
+    let destination = if need_hoist_discriminant || bbs[targets.otherwise()].is_empty_unreachable()
-        return None;
+    {
+        child_targets.otherwise()
+    } else {
+        targets.otherwise()
     };
-    let destination = child_targets.otherwise();
 
     // Verify that the optimization is legal for each branch
     for (value, child) in targets.iter() {
-        if !verify_candidate_branch(&bbs[child], value, *child_place, destination) {
+        if !verify_candidate_branch(
+            &bbs[child],
+            value,
+            child_place,
+            destination,
+            need_hoist_discriminant,
+        ) {
             return None;
         }
     }
     Some(OptimizationData {
         destination,
-        child_place: *child_place,
+        child_place,
         child_ty,
-        child_source: child_terminator.source_info,
+        child_source: *source_info,
+        need_hoist_discriminant,
     })
 }
 
@@ -307,31 +380,48 @@ fn verify_candidate_branch<'tcx>(
     value: u128,
     place: Place<'tcx>,
     destination: BasicBlock,
+    need_hoist_discriminant: bool,
 ) -> bool {
-    // In order for the optimization to be correct, the branch must...
+    // In order for the optimization to be correct, the terminator must be a `SwitchInt`.
-    // ...have exactly one statement
+    let TerminatorKind::SwitchInt { discr: switch_op, targets } = &branch.terminator().kind else {
-    if let [statement] = branch.statements.as_slice()
+        return false;
-        // ...assign the discriminant of `place` in that statement
+    };
-        && let StatementKind::Assign(boxed) = &statement.kind
+    if need_hoist_discriminant {
-        && let (discr_place, Rvalue::Discriminant(from_place)) = &**boxed
+        // If we need hoist discriminant, the branch must have exactly one statement.
-        && *from_place == place
+        let [statement] = branch.statements.as_slice() else {
-        // ...make that assignment to a local
+            return false;
-        && discr_place.projection.is_empty()
+        };
-        // ...terminate on a `SwitchInt` that invalidates that local
+        // The statement must assign the discriminant of `place`.
-        && let TerminatorKind::SwitchInt { discr: switch_op, targets, .. } =
+        let StatementKind::Assign(box (discr_place, Rvalue::Discriminant(from_place))) =
-            &branch.terminator().kind
+            statement.kind
-        && *switch_op == Operand::Move(*discr_place)
+        else {
-        // ...fall through to `destination` if the switch misses
+            return false;
-        && destination == targets.otherwise()
+        };
-        // ...have a branch for value `value`
+        if from_place != place {
-        && let mut iter = targets.iter()
+            return false;
-        && let Some((target_value, _)) = iter.next()
+        }
-        && target_value == value
+        // The assignment must invalidate a local that terminate on a `SwitchInt`.
-        // ...and have no more branches
+        if !discr_place.projection.is_empty() || *switch_op != Operand::Move(discr_place) {
-        && iter.next().is_none()
+            return false;
-    {
+        }
-        true
     } else {
-        false
+        // If we don't need hoist discriminant, the branch must not have any statements.
+        if !branch.statements.is_empty() {
+            return false;
+        }
+        // The place on `SwitchInt` must be the same.
+        if *switch_op != Operand::Copy(place) {
+            return false;
+        }
     }
+    // It must fall through to `destination` if the switch misses.
+    if destination != targets.otherwise() {
+        return false;
+    }
+    // It must have exactly one branch for value `value` and have no more branches.
+    let mut iter = targets.iter();
+    let (Some((target_value, _)), None) = (iter.next(), iter.next()) else {
+        return false;
+    };
+    target_value == value
 }

tests/mir-opt/early_otherwise_branch.opt5.EarlyOtherwiseBranch.diff

@@ -0,0 +1,77 @@
+- // MIR for `opt5` before EarlyOtherwiseBranch
++ // MIR for `opt5` after EarlyOtherwiseBranch
+  
+  fn opt5(_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;
++     let mut _6: bool;
+  
+      bb0: {
+          StorageLive(_3);
+          StorageLive(_4);
+          _4 = copy _1;
+          StorageLive(_5);
+          _5 = copy _2;
+          _3 = (move _4, move _5);
+          StorageDead(_5);
+          StorageDead(_4);
+-         switchInt(copy (_3.0: u32)) -> [1: bb2, 2: bb3, 3: bb4, otherwise: bb1];
++         StorageLive(_6);
++         _6 = Ne(copy (_3.0: u32), copy (_3.1: u32));
++         switchInt(move _6) -> [0: bb6, otherwise: bb1];
+      }
+  
+      bb1: {
++         StorageDead(_6);
+          _0 = const 0_u32;
+-         goto -> bb8;
++         goto -> bb5;
+      }
+  
+      bb2: {
+-         switchInt(copy (_3.1: u32)) -> [1: bb7, otherwise: bb1];
++         _0 = const 6_u32;
++         goto -> bb5;
+      }
+  
+      bb3: {
+-         switchInt(copy (_3.1: u32)) -> [2: bb6, otherwise: bb1];
++         _0 = const 5_u32;
++         goto -> bb5;
+      }
+  
+      bb4: {
+-         switchInt(copy (_3.1: u32)) -> [3: bb5, otherwise: bb1];
++         _0 = const 4_u32;
++         goto -> bb5;
+      }
+  
+      bb5: {
+-         _0 = const 6_u32;
+-         goto -> bb8;
++         StorageDead(_3);
++         return;
+      }
+  
+      bb6: {
+-         _0 = const 5_u32;
+-         goto -> bb8;
+-     }
+- 
+-     bb7: {
+-         _0 = const 4_u32;
+-         goto -> bb8;
+-     }
+- 
+-     bb8: {
+-         StorageDead(_3);
+-         return;
++         StorageDead(_6);
++         switchInt(copy (_3.0: u32)) -> [1: bb4, 2: bb3, 3: bb2, otherwise: bb1];
+      }
+  }
+  

tests/mir-opt/early_otherwise_branch.opt5_failed.EarlyOtherwiseBranch.diff

@@ -0,0 +1,61 @@
+- // MIR for `opt5_failed` before EarlyOtherwiseBranch
++ // MIR for `opt5_failed` after EarlyOtherwiseBranch
+  
+  fn opt5_failed(_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 = copy _1;
+          StorageLive(_5);
+          _5 = copy _2;
+          _3 = (move _4, move _5);
+          StorageDead(_5);
+          StorageDead(_4);
+          switchInt(copy (_3.0: u32)) -> [1: bb2, 2: bb3, 3: bb4, otherwise: bb1];
+      }
+  
+      bb1: {
+          _0 = const 0_u32;
+          goto -> bb8;
+      }
+  
+      bb2: {
+          switchInt(copy (_3.1: u32)) -> [1: bb7, otherwise: bb1];
+      }
+  
+      bb3: {
+          switchInt(copy (_3.1: u32)) -> [2: bb6, otherwise: bb1];
+      }
+  
+      bb4: {
+          switchInt(copy (_3.1: u32)) -> [2: bb5, otherwise: bb1];
+      }
+  
+      bb5: {
+          _0 = const 6_u32;
+          goto -> bb8;
+      }
+  
+      bb6: {
+          _0 = const 5_u32;
+          goto -> bb8;
+      }
+  
+      bb7: {
+          _0 = const 4_u32;
+          goto -> bb8;
+      }
+  
+      bb8: {
+          StorageDead(_3);
+          return;
+      }
+  }
+  

tests/mir-opt/early_otherwise_branch.opt5_failed_type.EarlyOtherwiseBranch.diff

@@ -0,0 +1,61 @@
+- // MIR for `opt5_failed_type` before EarlyOtherwiseBranch
++ // MIR for `opt5_failed_type` after EarlyOtherwiseBranch
+  
+  fn opt5_failed_type(_1: u32, _2: u64) -> u32 {
+      debug x => _1;
+      debug y => _2;
+      let mut _0: u32;
+      let mut _3: (u32, u64);
+      let mut _4: u32;
+      let mut _5: u64;
+  
+      bb0: {
+          StorageLive(_3);
+          StorageLive(_4);
+          _4 = copy _1;
+          StorageLive(_5);
+          _5 = copy _2;
+          _3 = (move _4, move _5);
+          StorageDead(_5);
+          StorageDead(_4);
+          switchInt(copy (_3.0: u32)) -> [1: bb2, 2: bb3, 3: bb4, otherwise: bb1];
+      }
+  
+      bb1: {
+          _0 = const 0_u32;
+          goto -> bb8;
+      }
+  
+      bb2: {
+          switchInt(copy (_3.1: u64)) -> [1: bb7, otherwise: bb1];
+      }
+  
+      bb3: {
+          switchInt(copy (_3.1: u64)) -> [2: bb6, otherwise: bb1];
+      }
+  
+      bb4: {
+          switchInt(copy (_3.1: u64)) -> [3: bb5, otherwise: bb1];
+      }
+  
+      bb5: {
+          _0 = const 6_u32;
+          goto -> bb8;
+      }
+  
+      bb6: {
+          _0 = const 5_u32;
+          goto -> bb8;
+      }
+  
+      bb7: {
+          _0 = const 4_u32;
+          goto -> bb8;
+      }
+  
+      bb8: {
+          StorageDead(_3);
+          return;
+      }
+  }
+  

tests/mir-opt/early_otherwise_branch.rs

@@ -78,9 +78,57 @@ fn opt4(x: Option2<u32>, y: Option2<u32>) -> u32 {
     }
 }
 
+// EMIT_MIR early_otherwise_branch.opt5.EarlyOtherwiseBranch.diff
+fn opt5(x: u32, y: u32) -> u32 {
+    // CHECK-LABEL: fn opt5(
+    // CHECK: let mut [[CMP_LOCAL:_.*]]: bool;
+    // CHECK: bb0: {
+    // CHECK: [[CMP_LOCAL]] = Ne(
+    // CHECK: switchInt(move [[CMP_LOCAL]]) -> [
+    // CHECK-NEXT: }
+    match (x, y) {
+        (1, 1) => 4,
+        (2, 2) => 5,
+        (3, 3) => 6,
+        _ => 0,
+    }
+}
+
+// EMIT_MIR early_otherwise_branch.opt5_failed.EarlyOtherwiseBranch.diff
+fn opt5_failed(x: u32, y: u32) -> u32 {
+    // CHECK-LABEL: fn opt5_failed(
+    // CHECK: bb0: {
+    // CHECK-NOT: Ne(
+    // CHECK: switchInt(
+    // CHECK-NEXT: }
+    match (x, y) {
+        (1, 1) => 4,
+        (2, 2) => 5,
+        (3, 2) => 6,
+        _ => 0,
+    }
+}
+
+// EMIT_MIR early_otherwise_branch.opt5_failed_type.EarlyOtherwiseBranch.diff
+fn opt5_failed_type(x: u32, y: u64) -> u32 {
+    // CHECK-LABEL: fn opt5_failed_type(
+    // CHECK: bb0: {
+    // CHECK-NOT: Ne(
+    // CHECK: switchInt(
+    // CHECK-NEXT: }
+    match (x, y) {
+        (1, 1) => 4,
+        (2, 2) => 5,
+        (3, 3) => 6,
+        _ => 0,
+    }
+}
+
 fn main() {
     opt1(None, Some(0));
     opt2(None, Some(0));
     opt3(Option2::None, Option2::Some(false));
     opt4(Option2::None, Option2::Some(0));
+    opt5(0, 0);
+    opt5_failed(0, 0);
 }