Rollup merge of #146552 - cjgillot:resume-noremap, r=jackh726

StateTransform: Do not renumber resume local.

MIR parameters are not explicitly assigned-to when entering the MIR body. If we want to save their values inside the coroutine state, we need to do so explicitly.

This was done by renaming the `_2` local, and introducing an explicit assignment pre-transform. This particular trick confuses me.

This version makes explicit that we are assigning parameters to saved locals.

r? ``@dingxiangfei2009``

Author: Zalathar

compiler/rustc_middle/src/mir/mod.rs

@@ -471,7 +471,7 @@ impl<'tcx> Body<'tcx> {
 
     /// Returns an iterator over all function arguments.
     #[inline]
-    pub fn args_iter(&self) -> impl Iterator<Item = Local> + ExactSizeIterator {
+    pub fn args_iter(&self) -> impl Iterator<Item = Local> + ExactSizeIterator + use<> {
         (1..self.arg_count + 1).map(Local::new)
     }
 

compiler/rustc_middle/src/mir/query.rs

@@ -17,7 +17,7 @@ use crate::ty::{self, CoroutineArgsExt, OpaqueHiddenType, Ty};
 rustc_index::newtype_index! {
     #[derive(HashStable)]
     #[encodable]
-    #[debug_format = "_{}"]
+    #[debug_format = "_s{}"]
     pub struct CoroutineSavedLocal {}
 }
 

compiler/rustc_mir_transform/src/coroutine.rs

@@ -1340,14 +1340,13 @@ fn create_cases<'tcx>(
                     }
                 }
 
-                if operation == Operation::Resume {
+                if operation == Operation::Resume && point.resume_arg != CTX_ARG.into() {
                     // Move the resume argument to the destination place of the `Yield` terminator
-                    let resume_arg = CTX_ARG;
                     statements.push(Statement::new(
                         source_info,
                         StatementKind::Assign(Box::new((
                             point.resume_arg,
-                            Rvalue::Use(Operand::Move(resume_arg.into())),
+                            Rvalue::Use(Operand::Move(CTX_ARG.into())),
                         ))),
                     ));
                 }
@@ -1439,7 +1438,10 @@ fn check_field_tys_sized<'tcx>(
 }
 
 impl<'tcx> crate::MirPass<'tcx> for StateTransform {
+    #[instrument(level = "debug", skip(self, tcx, body), ret)]
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        debug!(def_id = ?body.source.def_id());
+
         let Some(old_yield_ty) = body.yield_ty() else {
             // This only applies to coroutines
             return;
@@ -1518,31 +1520,7 @@ impl<'tcx> crate::MirPass<'tcx> for StateTransform {
             cleanup_async_drops(body);
         }
 
-        // We also replace the resume argument and insert an `Assign`.
-        // This is needed because the resume argument `_2` might be live across a `yield`, in which
-        // case there is no `Assign` to it that the transform can turn into a store to the coroutine
-        // state. After the yield the slot in the coroutine state would then be uninitialized.
-        let resume_local = CTX_ARG;
-        let resume_ty = body.local_decls[resume_local].ty;
-        let old_resume_local = replace_local(resume_local, resume_ty, body, tcx);
-
-        // When first entering the coroutine, move the resume argument into its old local
-        // (which is now a generator interior).
-        let source_info = SourceInfo::outermost(body.span);
-        let stmts = &mut body.basic_blocks_mut()[START_BLOCK].statements;
-        stmts.insert(
-            0,
-            Statement::new(
-                source_info,
-                StatementKind::Assign(Box::new((
-                    old_resume_local.into(),
-                    Rvalue::Use(Operand::Move(resume_local.into())),
-                ))),
-            ),
-        );
-
         let always_live_locals = always_storage_live_locals(body);
-
         let movable = coroutine_kind.movability() == hir::Movability::Movable;
         let liveness_info =
             locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
@@ -1583,6 +1561,21 @@ impl<'tcx> crate::MirPass<'tcx> for StateTransform {
         };
         transform.visit_body(body);
 
+        // MIR parameters are not explicitly assigned-to when entering the MIR body.
+        // If we want to save their values inside the coroutine state, we need to do so explicitly.
+        let source_info = SourceInfo::outermost(body.span);
+        let args_iter = body.args_iter();
+        body.basic_blocks.as_mut()[START_BLOCK].statements.splice(
+            0..0,
+            args_iter.filter_map(|local| {
+                let (ty, variant_index, idx) = transform.remap[local]?;
+                let lhs = transform.make_field(variant_index, idx, ty);
+                let rhs = Rvalue::Use(Operand::Move(local.into()));
+                let assign = StatementKind::Assign(Box::new((lhs, rhs)));
+                Some(Statement::new(source_info, assign))
+            }),
+        );
+
         // Update our MIR struct to reflect the changes we've made
         body.arg_count = 2; // self, resume arg
         body.spread_arg = None;

tests/mir-opt/async_drop_live_dead.a-{closure#0}.coroutine_drop_async.0.panic-abort.mir

@@ -1,7 +1,7 @@
 // MIR for `a::{closure#0}` 0 coroutine_drop_async
 
 fn a::{closure#0}(_1: Pin<&mut {async fn body of a<T>()}>, _2: &mut Context<'_>) -> Poll<()> {
-    debug _task_context => _19;
+    debug _task_context => _2;
     debug x => ((*(_1.0: &mut {async fn body of a<T>()})).0: T);
     let mut _0: std::task::Poll<()>;
     let _3: T;
@@ -20,15 +20,14 @@ fn a::{closure#0}(_1: Pin<&mut {async fn body of a<T>()}>, _2: &mut Context<'_>)
     let mut _16: &mut impl std::future::Future<Output = ()>;
     let mut _17: std::pin::Pin<&mut impl std::future::Future<Output = ()>>;
     let mut _18: isize;
-    let mut _19: &mut std::task::Context<'_>;
-    let mut _20: u32;
+    let mut _19: u32;
     scope 1 {
         debug x => (((*(_1.0: &mut {async fn body of a<T>()})) as variant#4).0: T);
     }
 
     bb0: {
-        _20 = discriminant((*(_1.0: &mut {async fn body of a<T>()})));
-        switchInt(move _20) -> [0: bb9, 3: bb12, 4: bb13, otherwise: bb14];
+        _19 = discriminant((*(_1.0: &mut {async fn body of a<T>()})));
+        switchInt(move _19) -> [0: bb9, 3: bb12, 4: bb13, otherwise: bb14];
     }
 
     bb1: {

tests/mir-opt/async_drop_live_dead.a-{closure#0}.coroutine_drop_async.0.panic-unwind.mir

@@ -1,7 +1,7 @@
 // MIR for `a::{closure#0}` 0 coroutine_drop_async
 
 fn a::{closure#0}(_1: Pin<&mut {async fn body of a<T>()}>, _2: &mut Context<'_>) -> Poll<()> {
-    debug _task_context => _19;
+    debug _task_context => _2;
     debug x => ((*(_1.0: &mut {async fn body of a<T>()})).0: T);
     let mut _0: std::task::Poll<()>;
     let _3: T;
@@ -20,15 +20,14 @@ fn a::{closure#0}(_1: Pin<&mut {async fn body of a<T>()}>, _2: &mut Context<'_>)
     let mut _16: &mut impl std::future::Future<Output = ()>;
     let mut _17: std::pin::Pin<&mut impl std::future::Future<Output = ()>>;
     let mut _18: isize;
-    let mut _19: &mut std::task::Context<'_>;
-    let mut _20: u32;
+    let mut _19: u32;
     scope 1 {
         debug x => (((*(_1.0: &mut {async fn body of a<T>()})) as variant#4).0: T);
     }
 
     bb0: {
-        _20 = discriminant((*(_1.0: &mut {async fn body of a<T>()})));
-        switchInt(move _20) -> [0: bb12, 2: bb18, 3: bb16, 4: bb17, otherwise: bb19];
+        _19 = discriminant((*(_1.0: &mut {async fn body of a<T>()})));
+        switchInt(move _19) -> [0: bb12, 2: bb18, 3: bb16, 4: bb17, otherwise: bb19];
     }
 
     bb1: {

tests/mir-opt/building/async_await.a-{closure#0}.coroutine_resume.0.mir

@@ -10,19 +10,17 @@
 } */
 
 fn a::{closure#0}(_1: Pin<&mut {async fn body of a()}>, _2: &mut Context<'_>) -> Poll<()> {
-    debug _task_context => _4;
+    debug _task_context => _2;
     let mut _0: std::task::Poll<()>;
     let mut _3: ();
-    let mut _4: &mut std::task::Context<'_>;
-    let mut _5: u32;
+    let mut _4: u32;
 
     bb0: {
-        _5 = discriminant((*(_1.0: &mut {async fn body of a()})));
-        switchInt(move _5) -> [0: bb1, 1: bb4, otherwise: bb5];
+        _4 = discriminant((*(_1.0: &mut {async fn body of a()})));
+        switchInt(move _4) -> [0: bb1, 1: bb4, otherwise: bb5];
     }
 
     bb1: {
-        _4 = move _2;
         _3 = const ();
         goto -> bb3;
     }

tests/mir-opt/building/async_await.b-{closure#0}.coroutine_resume.0.mir

@@ -1,7 +1,7 @@
 // MIR for `b::{closure#0}` 0 coroutine_resume
 /* coroutine_layout = CoroutineLayout {
     field_tys: {
-        _0: CoroutineSavedTy {
+        _s0: CoroutineSavedTy {
             ty: Coroutine(
                 DefId(0:5 ~ async_await[ccf8]::a::{closure#0}),
                 [
@@ -18,7 +18,7 @@
             },
             ignore_for_traits: false,
         },
-        _1: CoroutineSavedTy {
+        _s1: CoroutineSavedTy {
             ty: Coroutine(
                 DefId(0:5 ~ async_await[ccf8]::a::{closure#0}),
                 [
@@ -40,17 +40,17 @@
         Unresumed(0): [],
         Returned (1): [],
         Panicked (2): [],
-        Suspend0 (3): [_0],
-        Suspend1 (4): [_1],
+        Suspend0 (3): [_s0],
+        Suspend1 (4): [_s1],
     },
     storage_conflicts: BitMatrix(2x2) {
-        (_0, _0),
-        (_1, _1),
+        (_s0, _s0),
+        (_s1, _s1),
     },
 } */
 
 fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) -> Poll<()> {
-    debug _task_context => _38;
+    debug _task_context => _2;
     let mut _0: std::task::Poll<()>;
     let _3: ();
     let mut _4: {async fn body of a()};
@@ -85,8 +85,7 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
     let mut _35: &mut std::task::Context<'_>;
     let mut _36: ();
     let mut _37: ();
-    let mut _38: &mut std::task::Context<'_>;
-    let mut _39: u32;
+    let mut _38: u32;
     scope 1 {
         debug __awaitee => (((*(_1.0: &mut {async fn body of b()})) as variant#3).0: {async fn body of a()});
         let _17: ();
@@ -103,12 +102,11 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
     }
 
     bb0: {
-        _39 = discriminant((*(_1.0: &mut {async fn body of b()})));
-        switchInt(move _39) -> [0: bb1, 1: bb29, 3: bb27, 4: bb28, otherwise: bb8];
+        _38 = discriminant((*(_1.0: &mut {async fn body of b()})));
+        switchInt(move _38) -> [0: bb1, 1: bb29, 3: bb27, 4: bb28, otherwise: bb8];
     }
 
     bb1: {
-        _38 = move _2;
         StorageLive(_3);
         StorageLive(_4);
         StorageLive(_5);
@@ -143,7 +141,7 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
         StorageLive(_13);
         StorageLive(_14);
         StorageLive(_15);
-        _15 = copy _38;
+        _15 = copy _2;
         _14 = move _15;
         goto -> bb6;
     }
@@ -198,7 +196,7 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
 
     bb11: {
         StorageDead(_20);
-        _38 = move _19;
+        _2 = move _19;
         StorageDead(_19);
         _7 = const ();
         goto -> bb4;
@@ -245,7 +243,7 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
         StorageLive(_29);
         StorageLive(_30);
         StorageLive(_31);
-        _31 = copy _38;
+        _31 = copy _2;
         _30 = move _31;
         goto -> bb18;
     }
@@ -295,7 +293,7 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body of b()}>, _2: &mut Context<'_>) ->
 
     bb22: {
         StorageDead(_36);
-        _38 = move _35;
+        _2 = move _35;
         StorageDead(_35);
         _7 = const ();
         goto -> bb16;