Put entry block at the end instead of shifting everything.

Author: cjgillot

compiler/rustc_mir_transform/src/coroutine.rs

@@ -1077,31 +1077,42 @@ fn compute_layout<'tcx>(
 /// Replaces the entry point of `body` with a block that switches on the coroutine discriminant and
 /// dispatches to blocks according to `cases`.
 ///
-/// After this function, the former entry point of the function will be bb1.
+/// After this function, the former entry point of the function will be the last block.
 fn insert_switch<'tcx>(
     body: &mut Body<'tcx>,
     cases: Vec<(usize, BasicBlock)>,
     transform: &TransformVisitor<'tcx>,
     default_block: BasicBlock,
 ) {
     let (assign, discr) = transform.get_discr(body);
-    let switch_targets =
-        SwitchTargets::new(cases.iter().map(|(i, bb)| ((*i) as u128, *bb)), default_block);
-    let switch = TerminatorKind::SwitchInt { discr: Operand::Move(discr), targets: switch_targets };
 
-    let source_info = SourceInfo::outermost(body.span);
-    body.basic_blocks_mut().raw.insert(
-        0,
-        BasicBlockData::new_stmts(
-            vec![assign],
-            Some(Terminator { source_info, kind: switch }),
-            false,
-        ),
+    // MIR validation ensures that no block targets `ENTRY_BLOCK`.
+    #[cfg(debug_assertions)]
+    for bb in body.basic_blocks.iter() {
+        for target in bb.terminator().successors() {
+            assert_ne!(target, START_BLOCK);
+        }
+    }
+
+    // Add the switch as entry block, and put the former entry block at the end.
+    let former_entry = std::mem::replace(
+        &mut body.basic_blocks_mut()[START_BLOCK],
+        BasicBlockData::new_stmts(vec![assign], None, false),
     );
+    let former_entry = body.basic_blocks_mut().push(former_entry);
 
-    for b in body.basic_blocks_mut().iter_mut() {
-        b.terminator_mut().successors_mut(|target| *target += 1);
+    // We may point to `START_BLOCK` in our `cases`, replace it with `former_entry`.
+    let mut switch_targets =
+        SwitchTargets::new(cases.iter().map(|(i, bb)| ((*i) as u128, *bb)), default_block);
+    for bb in switch_targets.all_targets_mut() {
+        if *bb == START_BLOCK {
+            *bb = former_entry;
+        }
     }
+
+    let switch = TerminatorKind::SwitchInt { discr: Operand::Move(discr), targets: switch_targets };
+    body.basic_blocks_mut()[START_BLOCK].terminator =
+        Some(Terminator { source_info: SourceInfo::outermost(body.span), kind: switch });
 }
 
 fn insert_term_block<'tcx>(body: &mut Body<'tcx>, kind: TerminatorKind<'tcx>) -> BasicBlock {

tests/mir-opt/coroutine/async_await.a-{closure#0}.StateTransform.diff

@@ -0,0 +1,56 @@
+- // MIR for `a::{closure#0}` before StateTransform
++ // MIR for `a::{closure#0}` after StateTransform
+  
+- fn a::{closure#0}(_1: {async fn body of a()}, _2: std::future::ResumeTy) -> ()
+- yields ()
+-  {
++ fn a::{closure#0}(_1: Pin<&mut {async fn body of a()}>, _2: &mut Context<'_>) -> Poll<()> {
++     coroutine layout {
++         variant_fields = {
++             Unresumed(0): [],
++             Returned (1): [],
++             Panicked (2): [],
++         }
++         storage_conflicts = BitMatrix(0x0) {}
++     }
+      debug _task_context => _2;
+-     let mut _0: ();
++     let mut _0: std::task::Poll<()>;
++     let mut _3: ();
++     let mut _4: u32;
++     let mut _5: &mut {async fn body of a()};
+  
+      bb0: {
+-         _0 = const ();
+-         drop(_1) -> [return: bb1, unwind: bb2];
++         _5 = copy (_1.0: &mut {async fn body of a()});
++         _4 = discriminant((*_5));
++         switchInt(move _4) -> [0: bb5, 1: bb3, otherwise: bb4];
+      }
+  
+      bb1: {
++         _0 = Poll::<()>::Ready(move _3);
++         discriminant((*_5)) = 1;
+          return;
+      }
+  
+-     bb2 (cleanup): {
+-         resume;
++     bb2: {
++         goto -> bb1;
++     }
++ 
++     bb3: {
++         assert(const false, "`async fn` resumed after completion") -> [success: bb3, unwind continue];
++     }
++ 
++     bb4: {
++         unreachable;
++     }
++ 
++     bb5: {
++         _3 = const ();
++         goto -> bb2;
+      }
+  }
+