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

4 files changed, 322 insertions, 1 deletions

diff --git a/compiler/rustc_mir_transform/src/early_otherwise_branch.rs b/compiler/rustc_mir_transform/src/early_otherwise_branch.rs
index 2f2d07c739c..be6056250fd 100644
--- a/compiler/rustc_mir_transform/src/early_otherwise_branch.rs
+++ b/compiler/rustc_mir_transform/src/early_otherwise_branch.rs
@@ -179,7 +179,7 @@ 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
@@ -187,6 +187,7 @@ impl<'tcx> crate::MirPass<'tcx> for EarlyOtherwiseBranch {
                     targets: eq_targets,
                 },
             }));
+            eq_switch.is_cleanup = bbs[parent].is_cleanup;
 
             let eq_bb = patch.new_block(eq_switch);
 
diff --git a/tests/mir-opt/early_otherwise_branch_unwind.poll.EarlyOtherwiseBranch.diff b/tests/mir-opt/early_otherwise_branch_unwind.poll.EarlyOtherwiseBranch.diff
new file mode 100644
index 00000000000..5f03b94ccb8
--- /dev/null
+++ b/tests/mir-opt/early_otherwise_branch_unwind.poll.EarlyOtherwiseBranch.diff
@@ -0,0 +1,142 @@
+- // MIR for `poll` before EarlyOtherwiseBranch
++ // MIR for `poll` after EarlyOtherwiseBranch
+  
+  fn poll(_1: Poll<Result<Option<Vec<u8>>, u8>>) -> () {
+      debug val => _1;
+      let mut _0: ();
+      let mut _2: isize;
+      let mut _3: isize;
+      let mut _4: isize;
+      let _5: std::vec::Vec<u8>;
+      let _6: u8;
+      let mut _7: bool;
+      let mut _8: bool;
+      let mut _9: isize;
++     let mut _10: bool;
+      scope 1 {
+          debug _trailers => _5;
+      }
+      scope 2 {
+          debug _err => _6;
+      }
+  
+      bb0: {
+          _7 = const false;
+          _8 = const false;
+          _7 = const true;
+          _8 = const true;
+          _4 = discriminant(_1);
+          switchInt(copy _4) -> [0: bb2, 1: bb4, otherwise: bb1];
+      }
+  
+      bb1: {
+          unreachable;
+      }
+  
+      bb2: {
+          _3 = discriminant(((_1 as Ready).0: std::result::Result<std::option::Option<std::vec::Vec<u8>>, u8>));
+          switchInt(copy _3) -> [0: bb3, 1: bb6, otherwise: bb1];
+      }
+  
+      bb3: {
+          _2 = discriminant(((((_1 as Ready).0: std::result::Result<std::option::Option<std::vec::Vec<u8>>, u8>) as Ok).0: std::option::Option<std::vec::Vec<u8>>));
+          switchInt(copy _2) -> [0: bb5, 1: bb7, otherwise: bb1];
+      }
+  
+      bb4: {
+          _0 = const ();
+-         goto -> bb17;
++         goto -> bb15;
+      }
+  
+      bb5: {
+          _0 = const ();
+-         goto -> bb17;
++         goto -> bb15;
+      }
+  
+      bb6: {
+          StorageLive(_6);
+          _6 = copy ((((_1 as Ready).0: std::result::Result<std::option::Option<std::vec::Vec<u8>>, u8>) as Err).0: u8);
+          _0 = const ();
+          StorageDead(_6);
+-         goto -> bb17;
++         goto -> bb15;
+      }
+  
+      bb7: {
+          StorageLive(_5);
+          _5 = move ((((((_1 as Ready).0: std::result::Result<std::option::Option<std::vec::Vec<u8>>, u8>) as Ok).0: std::option::Option<std::vec::Vec<u8>>) as Some).0: std::vec::Vec<u8>);
+          _0 = const ();
+-         drop(_5) -> [return: bb8, unwind: bb20];
++         drop(_5) -> [return: bb8, unwind: bb16];
+      }
+  
+      bb8: {
+          StorageDead(_5);
+-         goto -> bb17;
++         goto -> bb15;
+      }
+  
+      bb9 (cleanup): {
++         StorageDead(_10);
+          resume;
+      }
+  
+      bb10: {
+          return;
+      }
+  
+      bb11: {
+-         switchInt(copy _7) -> [0: bb12, otherwise: bb16];
++         switchInt(copy _7) -> [0: bb12, otherwise: bb14];
+      }
+  
+      bb12: {
+          _7 = const false;
+          goto -> bb10;
+      }
+  
+      bb13: {
+-         switchInt(copy _8) -> [0: bb14, otherwise: bb15];
+-     }
+- 
+-     bb14: {
+          _8 = const false;
+          goto -> bb12;
+      }
+  
+-     bb15: {
+-         goto -> bb14;
+-     }
+- 
+-     bb16: {
++     bb14: {
+          _9 = discriminant(((_1 as Ready).0: std::result::Result<std::option::Option<std::vec::Vec<u8>>, u8>));
+          switchInt(move _9) -> [0: bb13, otherwise: bb12];
+      }
+  
+-     bb17: {
++     bb15: {
+          switchInt(copy _4) -> [0: bb11, otherwise: bb10];
+      }
+  
+-     bb18 (cleanup): {
+-         switchInt(copy _3) -> [0: bb19, otherwise: bb9];
++     bb16 (cleanup): {
++         StorageLive(_10);
++         _10 = Ne(copy _4, copy _3);
++         switchInt(move _10) -> [0: bb17, otherwise: bb9];
+      }
+  
+-     bb19 (cleanup): {
++     bb17 (cleanup): {
++         StorageDead(_10);
+          goto -> bb9;
+-     }
+- 
+-     bb20 (cleanup): {
+-         switchInt(copy _4) -> [0: bb18, otherwise: bb9];
+      }
+  }
+  
diff --git a/tests/mir-opt/early_otherwise_branch_unwind.rs b/tests/mir-opt/early_otherwise_branch_unwind.rs
new file mode 100644
index 00000000000..7df58c7a64f
--- /dev/null
+++ b/tests/mir-opt/early_otherwise_branch_unwind.rs
@@ -0,0 +1,43 @@
+//@ test-mir-pass: EarlyOtherwiseBranch
+//@ compile-flags: -Zmir-enable-passes=+GVN,+SimplifyLocals-after-value-numbering
+//@ needs-unwind
+
+use std::task::Poll;
+
+// We find a matching pattern in the unwind path,
+// and we need to create a cleanup BB for this case to meet the unwind invariants rule.
+
+// EMIT_MIR early_otherwise_branch_unwind.unwind.EarlyOtherwiseBranch.diff
+fn unwind<T>(val: Option<Option<Option<T>>>) {
+    // CHECK-LABEL: fn unwind(
+    // CHECK: drop({{.*}}) -> [return: bb{{.*}}, unwind: [[PARENT_UNWIND_BB:bb.*]]];
+    // CHECK: [[PARENT_UNWIND_BB]] (cleanup): {
+    // CHECK-NEXT: StorageLive
+    // 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 => {}
+    }
+}
+
+// From https://github.com/rust-lang/rust/issues/130769#issuecomment-2370443086.
+// EMIT_MIR early_otherwise_branch_unwind.poll.EarlyOtherwiseBranch.diff
+pub fn poll(val: Poll<Result<Option<Vec<u8>>, u8>>) {
+    // CHECK-LABEL: fn poll(
+    // CHECK: drop({{.*}}) -> [return: bb{{.*}}, unwind: [[PARENT_UNWIND_BB:bb.*]]];
+    // CHECK: [[PARENT_UNWIND_BB]] (cleanup): {
+    // CHECK-NEXT: StorageLive
+    // CHECK-NEXT: [[CMP_LOCAL:_.*]] = Ne
+    // CHECK-NEXT: switchInt(move [[CMP_LOCAL]]) -> [0: [[NEW_UNWIND_BB:bb.*]], otherwise
+    // CHECK: [[NEW_UNWIND_BB]] (cleanup): {
+    match val {
+        Poll::Ready(Ok(Some(_trailers))) => {}
+        Poll::Ready(Err(_err)) => {}
+        Poll::Ready(Ok(None)) => {}
+        Poll::Pending => {}
+    }
+}
diff --git a/tests/mir-opt/early_otherwise_branch_unwind.unwind.EarlyOtherwiseBranch.diff b/tests/mir-opt/early_otherwise_branch_unwind.unwind.EarlyOtherwiseBranch.diff
new file mode 100644
index 00000000000..1405b9d314a
--- /dev/null
+++ b/tests/mir-opt/early_otherwise_branch_unwind.unwind.EarlyOtherwiseBranch.diff
@@ -0,0 +1,135 @@
+- // MIR for `unwind` before EarlyOtherwiseBranch
++ // MIR for `unwind` after EarlyOtherwiseBranch
+  
+  fn unwind(_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];
+      }
+  }
+