mir-opt: Create an indirect BB to add StorageDead

Author: dianqk

compiler/rustc_mir_transform/src/early_otherwise_branch.rs

@@ -193,8 +193,20 @@ impl<'tcx> crate::MirPass<'tcx> for EarlyOtherwiseBranch {
             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;
+            let mut true_case = opt_data.destination;
+            let mut false_case = eq_bb;
+            // Create an indirect BB to add `StorageDead` If the jump target is itself.
+            for bb in [&mut false_case, &mut true_case].into_iter() {
+                if *bb == parent {
+                    *bb = patch.new_block(BasicBlockData::new(
+                        Some(Terminator {
+                            kind: TerminatorKind::Goto { target: parent },
+                            source_info: bbs[parent].terminator().source_info,
+                        }),
+                        bbs[parent].is_cleanup,
+                    ));
+                }
+            }
             patch.patch_terminator(
                 parent,
                 TerminatorKind::if_(Operand::Move(Place::from(comp_temp)), true_case, false_case),
@@ -210,9 +222,9 @@ impl<'tcx> crate::MirPass<'tcx> for EarlyOtherwiseBranch {
 
             // 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() {
+            for bb in [false_case, true_case].into_iter() {
                 patch.add_statement(
-                    Location { block: *bb, statement_index: 0 },
+                    Location { block: bb, statement_index: 0 },
                     StatementKind::StorageDead(comp_temp),
                 );
             }

tests/mir-opt/early_otherwise_branch.rs

@@ -1,6 +1,10 @@
 //@ test-mir-pass: EarlyOtherwiseBranch
 //@ compile-flags: -Zmir-enable-passes=+UnreachableEnumBranching
 
+#![feature(custom_mir, core_intrinsics)]
+
+use std::intrinsics::mir::*;
+
 enum Option2<T> {
     Some(T),
     None,
@@ -124,11 +128,38 @@ fn opt5_failed_type(x: u32, y: u64) -> u32 {
     }
 }
 
+// EMIT_MIR early_otherwise_branch.target_self.EarlyOtherwiseBranch.diff
+#[custom_mir(dialect = "runtime")]
+fn target_self(val: i32) {
+    // CHECK-LABEL: fn target_self(
+    mir! {
+        {
+            Goto(bb1)
+        }
+        bb1 = {
+            match val {
+                0 => bb2,
+                _ => bb1,
+            }
+        }
+        bb2 = {
+            match val {
+                0 => bb3,
+                _ => bb1,
+            }
+        }
+        bb3 = {
+            Return()
+        }
+    }
+}
+
 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);
+    target_self(1);
 }

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

@@ -0,0 +1,35 @@
+- // MIR for `target_self` before EarlyOtherwiseBranch
++ // MIR for `target_self` after EarlyOtherwiseBranch
+  
+  fn target_self(_1: i32) -> () {
+      let mut _0: ();
++     let mut _2: bool;
+  
+      bb0: {
+          goto -> bb1;
+      }
+  
+      bb1: {
+-         switchInt(copy _1) -> [0: bb2, otherwise: bb1];
++         StorageLive(_2);
++         _2 = Ne(copy _1, copy _1);
++         switchInt(move _2) -> [0: bb3, otherwise: bb4];
+      }
+  
+      bb2: {
+-         switchInt(copy _1) -> [0: bb3, otherwise: bb1];
++         return;
+      }
+  
+      bb3: {
+-         return;
++         StorageDead(_2);
++         switchInt(copy _1) -> [0: bb2, otherwise: bb1];
++     }
++ 
++     bb4: {
++         StorageDead(_2);
++         goto -> bb1;
+      }
+  }
+