mir-opt: a sub-BB of a cleanup BB must also be a cleanup BB

Author: dianqk

compiler/rustc_mir_transform/src/early_otherwise_branch.rs

@@ -183,14 +183,15 @@ impl<'tcx> crate::MirPass<'tcx> for EarlyOtherwiseBranch {
             let eq_targets = SwitchTargets::new(eq_new_targets, parent_targets.otherwise());
 
             // Create `bbEq` in example above
-            let eq_switch = BasicBlockData::new(Some(Terminator {
+            let mut eq_switch = BasicBlockData::new(Some(Terminator {
                 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,
                     targets: eq_targets,
                 },
             }));
+            eq_switch.is_cleanup = bbs[parent].is_cleanup;
 
             let eq_bb = patch.new_block(eq_switch);
 

tests/mir-opt/early_otherwise_cleanup_branch.cleanup.EarlyOtherwiseBranch.diff

@@ -0,0 +1,135 @@
+- // MIR for `cleanup` before EarlyOtherwiseBranch
++ // MIR for `cleanup` after EarlyOtherwiseBranch
+  
+  fn cleanup(_1: Option<Option<Option<T>>>) -> () {
+      debug val => _1;
+      let mut _0: ();
+      let mut _2: isize;
+      let mut _3: isize;
+      let mut _4: isize;
+      let _5: T;
+      let mut _6: bool;
+      let mut _7: bool;
+      let mut _8: isize;
++     let mut _9: bool;
+      scope 1 {
+          debug _v => _5;
+      }
+  
+      bb0: {
+          _6 = const false;
+          _7 = const false;
+          _6 = const true;
+          _7 = const true;
+          _4 = discriminant(_1);
+          switchInt(copy _4) -> [0: bb4, 1: bb2, otherwise: bb1];
+      }
+  
+      bb1: {
+          unreachable;
+      }
+  
+      bb2: {
+          _3 = discriminant(((_1 as Some).0: std::option::Option<std::option::Option<T>>));
+          switchInt(copy _3) -> [0: bb5, 1: bb3, otherwise: bb1];
+      }
+  
+      bb3: {
+          _2 = discriminant(((((_1 as Some).0: std::option::Option<std::option::Option<T>>) as Some).0: std::option::Option<T>));
+          switchInt(copy _2) -> [0: bb6, 1: bb7, otherwise: bb1];
+      }
+  
+      bb4: {
+          _0 = const ();
+-         goto -> bb17;
++         goto -> bb15;
+      }
+  
+      bb5: {
+          _0 = const ();
+-         goto -> bb17;
++         goto -> bb15;
+      }
+  
+      bb6: {
+          _0 = const ();
+-         goto -> bb17;
++         goto -> bb15;
+      }
+  
+      bb7: {
+          StorageLive(_5);
+          _5 = move ((((((_1 as Some).0: std::option::Option<std::option::Option<T>>) as Some).0: std::option::Option<T>) as Some).0: T);
+          _0 = const ();
+-         drop(_5) -> [return: bb8, unwind: bb20];
++         drop(_5) -> [return: bb8, unwind: bb16];
+      }
+  
+      bb8: {
+          StorageDead(_5);
+-         goto -> bb17;
++         goto -> bb15;
+      }
+  
+      bb9 (cleanup): {
++         StorageDead(_9);
+          resume;
+      }
+  
+      bb10: {
+          return;
+      }
+  
+      bb11: {
+-         switchInt(copy _6) -> [0: bb12, otherwise: bb16];
++         switchInt(copy _6) -> [0: bb12, otherwise: bb14];
+      }
+  
+      bb12: {
+          _6 = const false;
+          goto -> bb10;
+      }
+  
+      bb13: {
+-         switchInt(copy _7) -> [0: bb14, otherwise: bb15];
+-     }
+- 
+-     bb14: {
+          _7 = const false;
+          goto -> bb12;
+      }
+  
+-     bb15: {
+-         goto -> bb14;
+-     }
+- 
+-     bb16: {
++     bb14: {
+          _8 = discriminant(((_1 as Some).0: std::option::Option<std::option::Option<T>>));
+          switchInt(move _8) -> [1: bb13, otherwise: bb12];
+      }
+  
+-     bb17: {
++     bb15: {
+          switchInt(copy _4) -> [1: bb11, otherwise: bb10];
+      }
+  
+-     bb18 (cleanup): {
+-         switchInt(copy _3) -> [1: bb19, otherwise: bb9];
++     bb16 (cleanup): {
++         StorageLive(_9);
++         _9 = Ne(copy _4, copy _3);
++         switchInt(move _9) -> [0: bb17, otherwise: bb9];
+      }
+  
+-     bb19 (cleanup): {
++     bb17 (cleanup): {
++         StorageDead(_9);
+          goto -> bb9;
+-     }
+- 
+-     bb20 (cleanup): {
+-         switchInt(copy _4) -> [1: bb18, otherwise: bb9];
+      }
+  }
+  

tests/mir-opt/early_otherwise_cleanup_branch.rs

@@ -0,0 +1,21 @@
+//@ test-mir-pass: EarlyOtherwiseBranch
+//@ compile-flags: -Zmir-enable-passes=+GVN,+SimplifyLocals-after-value-numbering
+
+// Ensure that the new sub-BB we create is also a cleanup BB.
+
+// EMIT_MIR early_otherwise_cleanup_branch.cleanup.EarlyOtherwiseBranch.diff
+fn cleanup<T>(val: Option<Option<Option<T>>>) {
+    // CHECK-LABEL: fn cleanup(
+    // CHECK: drop({{.*}}) -> [return: bb{{.*}}, unwind: [[PARENT_UNWIND_BB:bb.*]]];
+    // CHECK: [[PARENT_UNWIND_BB]] (cleanup): {
+    // CHECK-NEXT: StorageLive(_9);
+    // CHECK-NEXT: [[CMP_LOCAL:_.*]] = Ne
+    // CHECK-NEXT: switchInt(move [[CMP_LOCAL]]) -> [0: [[NEW_UNWIND_BB:bb.*]], otherwise
+    // CHECK: [[NEW_UNWIND_BB]] (cleanup): {
+    match val {
+        Some(Some(Some(_v))) => {}
+        Some(Some(None)) => {}
+        Some(None) => {}
+        None => {}
+    }
+}