End locals' live range before suspending coroutine

State transforms retains storage statements for locals that are not
stored inside a coroutine. It ensures those locals are live when
resuming by inserting StorageLive as appropriate. It forgot to end the
storage of those locals when suspending, which is fixed here.

While the end of live range is implicit when executing return, it is
nevertheless useful for inliner which would otherwise extend the live
range beyond return.

5 files changed, 27 insertions, 1 deletions

diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs
index ed8b4ef3ba4..c4f295d69f8 100644
--- a/compiler/rustc_mir_transform/src/coroutine.rs
+++ b/compiler/rustc_mir_transform/src/coroutine.rs
@@ -527,12 +527,26 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> {
                         resume_arg
                     };
 
+                let storage_liveness: GrowableBitSet<Local> =
+                    self.storage_liveness[block].clone().unwrap().into();
+
+                for i in 0..self.always_live_locals.domain_size() {
+                    let l = Local::new(i);
+                    let needs_storage_dead = storage_liveness.contains(l)
+                        && !self.remap.contains_key(&l)
+                        && !self.always_live_locals.contains(l);
+                    if needs_storage_dead {
+                        data.statements
+                            .push(Statement { source_info, kind: StatementKind::StorageDead(l) });
+                    }
+                }
+
                 self.suspension_points.push(SuspensionPoint {
                     state,
                     resume,
                     resume_arg,
                     drop,
-                    storage_liveness: self.storage_liveness[block].clone().unwrap().into(),
+                    storage_liveness,
                 });
 
                 VariantIdx::new(state)
diff --git a/tests/mir-opt/building/async_await.b-{closure#0}.coroutine_resume.0.mir b/tests/mir-opt/building/async_await.b-{closure#0}.coroutine_resume.0.mir
index 111dd8e97f9..59e27512464 100644
--- a/tests/mir-opt/building/async_await.b-{closure#0}.coroutine_resume.0.mir
+++ b/tests/mir-opt/building/async_await.b-{closure#0}.coroutine_resume.0.mir
@@ -178,6 +178,10 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body@$DIR/async_await.rs:15:18: 18:2}>,
         StorageLive(_20);
         _20 = ();
         _0 = Poll::<()>::Pending;
+        StorageDead(_3);
+        StorageDead(_4);
+        StorageDead(_19);
+        StorageDead(_20);
         discriminant((*(_1.0: &mut {async fn body@$DIR/async_await.rs:15:18: 18:2}))) = 3;
         return;
     }
@@ -276,6 +280,9 @@ fn b::{closure#0}(_1: Pin<&mut {async fn body@$DIR/async_await.rs:15:18: 18:2}>,
         StorageLive(_36);
         _36 = ();
         _0 = Poll::<()>::Pending;
+        StorageDead(_21);
+        StorageDead(_35);
+        StorageDead(_36);
         discriminant((*(_1.0: &mut {async fn body@$DIR/async_await.rs:15:18: 18:2}))) = 4;
         return;
     }
diff --git a/tests/mir-opt/coroutine_tiny.main-{closure#0}.coroutine_resume.0.mir b/tests/mir-opt/coroutine_tiny.main-{closure#0}.coroutine_resume.0.mir
index 17b99c87c39..165aa3a05cb 100644
--- a/tests/mir-opt/coroutine_tiny.main-{closure#0}.coroutine_resume.0.mir
+++ b/tests/mir-opt/coroutine_tiny.main-{closure#0}.coroutine_resume.0.mir
@@ -55,6 +55,9 @@ fn main::{closure#0}(_1: Pin<&mut {coroutine@$DIR/coroutine_tiny.rs:20:16: 20:24
         StorageLive(_7);
         _7 = ();
         _0 = CoroutineState::<(), ()>::Yielded(move _7);
+        StorageDead(_4);
+        StorageDead(_6);
+        StorageDead(_7);
         discriminant((*(_1.0: &mut {coroutine@$DIR/coroutine_tiny.rs:20:16: 20:24}))) = 3;
         return;
     }
diff --git a/tests/mir-opt/inline/inline_coroutine.main.Inline.panic-abort.diff b/tests/mir-opt/inline/inline_coroutine.main.Inline.panic-abort.diff
index 42dd7ba55cf..68c612314f6 100644
--- a/tests/mir-opt/inline/inline_coroutine.main.Inline.panic-abort.diff
+++ b/tests/mir-opt/inline/inline_coroutine.main.Inline.panic-abort.diff
@@ -89,6 +89,7 @@
 + 
 +     bb6: {
 +         _1 = CoroutineState::<i32, bool>::Yielded(move _8);
++         StorageDead(_8);
 +         discriminant((*_6)) = 3;
 +         goto -> bb2;
 +     }
diff --git a/tests/mir-opt/inline/inline_coroutine.main.Inline.panic-unwind.diff b/tests/mir-opt/inline/inline_coroutine.main.Inline.panic-unwind.diff
index 7b8958c13fc..1bf14e8c3b3 100644
--- a/tests/mir-opt/inline/inline_coroutine.main.Inline.panic-unwind.diff
+++ b/tests/mir-opt/inline/inline_coroutine.main.Inline.panic-unwind.diff
@@ -103,6 +103,7 @@
 + 
 +     bb8: {
 +         _1 = CoroutineState::<i32, bool>::Yielded(move _8);
++         StorageDead(_8);
 +         discriminant((*_6)) = 3;
 +         goto -> bb4;
 +     }