From ef1831a21fede46bd7bd2d895fb82c194c0d7304 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tomasz=20Mi=C4=85sko?= <tomasz.miasko@gmail.com>
Date: Sat, 9 Dec 2023 00:00:00 +0000
Subject: [PATCH] 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.
---
 compiler/rustc_mir_transform/src/coroutine.rs    | 16 +++++++++++++++-
 ...nc_await.b-{closure#0}.coroutine_resume.0.mir |  7 +++++++
 ..._tiny.main-{closure#0}.coroutine_resume.0.mir |  3 +++
 ...inline_coroutine.main.Inline.panic-abort.diff |  1 +
 ...nline_coroutine.main.Inline.panic-unwind.diff |  1 +
 5 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs
index ed8b4ef3ba43..c4f295d69f8c 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 111dd8e97f9a..59e275124648 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 17b99c87c398..165aa3a05cb7 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 42dd7ba55cf4..68c612314f6d 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 7b8958c13fc5..1bf14e8c3b33 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;
 +     }