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.
rust-lang · Dec 11, 2023 · 15de84f · 15de84f
1 parent 8a37655
commit 15de84f
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -103,6 +103,7 @@
 + 
 +     bb8: {
 +         _1 = CoroutineState::<i32, bool>::Yielded(move _8);
++         StorageDead(_8);
 +         discriminant((*_6)) = 3;
 +         goto -> bb4;
 +     }