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``

3 files changed, 22 insertions, 29 deletions

diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index da2245b12d2..28142382b13 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/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)
     }
 
diff --git a/compiler/rustc_middle/src/mir/query.rs b/compiler/rustc_middle/src/mir/query.rs
index 466b9c7a3c2..a509c40c89c 100644
--- a/compiler/rustc_middle/src/mir/query.rs
+++ b/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 {}
 }
 
diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs
index 4603c695ded..08316aaed3b 100644
--- a/compiler/rustc_mir_transform/src/coroutine.rs
+++ b/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;