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.

Author: tmiasko

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)

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;
     }

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;
     }

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;
 +     }

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;
 +     }