Rollup merge of #89321 - tmiasko:rebase-resume-arg, r=estebank

Rebase resume argument projections during state transform

When remapping a resume argument with projections rebase them on top of
the new base.

The case where resume argument has projections is unusual, but might
arise with box syntax where the assignment is performed directly into
the box without an intermediate temporary.

Fixes #85635.

3 files changed, 12 insertions, 5 deletions

diff --git a/compiler/rustc_mir_transform/src/generator.rs b/compiler/rustc_mir_transform/src/generator.rs
index 06366b6fc31..bc72e9d94a9 100644
--- a/compiler/rustc_mir_transform/src/generator.rs
+++ b/compiler/rustc_mir_transform/src/generator.rs
@@ -342,7 +342,7 @@ impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
             let source_info = data.terminator().source_info;
             // We must assign the value first in case it gets declared dead below
             data.statements.extend(self.make_state(state_idx, v, source_info));
-            let state = if let Some((resume, resume_arg)) = resume {
+            let state = if let Some((resume, mut resume_arg)) = resume {
                 // Yield
                 let state = 3 + self.suspension_points.len();
 
@@ -350,7 +350,8 @@ impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
                 // live across a yield.
                 let resume_arg =
                     if let Some(&(ty, variant, idx)) = self.remap.get(&resume_arg.local) {
-                        self.make_field(variant, idx, ty)
+                        replace_base(&mut resume_arg, self.make_field(variant, idx, ty), self.tcx);
+                        resume_arg
                     } else {
                         resume_arg
                     };
diff --git a/src/test/ui/generator/yield-in-box.rs b/src/test/ui/generator/yield-in-box.rs
index 65f368df9cb..dd6fa7c151a 100644
--- a/src/test/ui/generator/yield-in-box.rs
+++ b/src/test/ui/generator/yield-in-box.rs
@@ -1,8 +1,10 @@
 // run-pass
-
 // Test that box-statements with yields in them work.
 
-#![feature(generators, box_syntax)]
+#![feature(generators, box_syntax, generator_trait)]
+use std::pin::Pin;
+use std::ops::Generator;
+use std::ops::GeneratorState;
 
 fn main() {
     let x = 0i32;
@@ -15,4 +17,8 @@ fn main() {
             _t => {}
         }
     };
+
+    let mut g = |_| box yield;
+    assert_eq!(Pin::new(&mut g).resume(1), GeneratorState::Yielded(()));
+    assert_eq!(Pin::new(&mut g).resume(2), GeneratorState::Complete(box 2));
 }
diff --git a/src/test/ui/generator/yield-in-box.stderr b/src/test/ui/generator/yield-in-box.stderr
index 24de18edb0f..7602e803945 100644
--- a/src/test/ui/generator/yield-in-box.stderr
+++ b/src/test/ui/generator/yield-in-box.stderr
@@ -1,5 +1,5 @@
 warning: unused generator that must be used
-  --> $DIR/yield-in-box.rs:9:5
+  --> $DIR/yield-in-box.rs:11:5
    |
 LL | /     || {
 LL | |         let y = 2u32;