Fix coroutine MIR saved local remapping

compiler/rustc_mir_dataflow/src/impls/storage_liveness.rs

@@ -155,7 +155,6 @@ impl<'tcx> Analysis<'tcx> for MaybeRequiresStorage<'_, 'tcx> {
         match &stmt.kind {
             StatementKind::StorageDead(l) => state.kill(*l),
 
-            // If a place is assigned to in a statement, it needs storage for that statement.
             StatementKind::Assign((place, _)) => {
                 state.gen_(place.local);
             }
@@ -180,12 +179,35 @@ impl<'tcx> Analysis<'tcx> for MaybeRequiresStorage<'_, 'tcx> {
     fn apply_primary_statement_effect(
         &self,
         state: &mut Self::Domain,
-        _: &Statement<'tcx>,
+        stmt: &Statement<'tcx>,
         loc: Location,
     ) {
         // If we move from a place then it only stops needing storage *after*
         // that statement.
         self.check_for_move(state, loc);
+
+        match &stmt.kind {
+            // If a place is assigned to in a statement, it needs storage after that statement.
+            // Even if the place was moved from in the rvalue (e.g. `x = x + 1` or `x = f(move x)`),
+            // the assignment restores a valid value into the place.
+            StatementKind::Assign((place, _)) => {
+                state.gen_(place.local);
+            }
+            StatementKind::SetDiscriminant { place, .. } => {
+                state.gen_(place.local);
+            }
+
+            StatementKind::StorageDead(_)
+            | StatementKind::AscribeUserType(..)
+            | StatementKind::PlaceMention(..)
+            | StatementKind::Coverage(..)
+            | StatementKind::FakeRead(..)
+            | StatementKind::ConstEvalCounter
+            | StatementKind::Nop
+            | StatementKind::Intrinsic(..)
+            | StatementKind::BackwardIncompatibleDropHint { .. }
+            | StatementKind::StorageLive(..) => {}
+        }
     }
 
     fn apply_early_terminator_effect(

compiler/rustc_mir_transform/src/coroutine/mod.rs

@@ -79,6 +79,7 @@ use rustc_span::def_id::DefId;
 use tracing::{debug, instrument};
 
 use crate::deref_separator::deref_finder;
+use crate::patch::MirPatch;
 use crate::{abort_unwinding_calls, pass_manager as pm, simplify};
 
 pub(super) struct StateTransform;
@@ -199,6 +200,8 @@ struct TransformVisitor<'tcx> {
     old_yield_ty: Ty<'tcx>,
 
     old_ret_ty: Ty<'tcx>,
+
+    patch: Option<MirPatch<'tcx>>,
 }
 
 impl<'tcx> TransformVisitor<'tcx> {
@@ -408,11 +411,51 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> {
     }
 
     #[tracing::instrument(level = "trace", skip(self), ret)]
-    fn visit_place(&mut self, place: &mut Place<'tcx>, _: PlaceContext, _location: Location) {
+    fn visit_place(&mut self, place: &mut Place<'tcx>, _: PlaceContext, location: Location) {
         // Replace an Local in the remap with a coroutine struct access
         if let Some(&Some((ty, variant_index, idx))) = self.remap.get(place.local) {
             replace_base(place, self.make_field(variant_index, idx, ty), self.tcx);
         }
+        if let Some(new_projection) = self.process_projection(&place.projection, location) {
+            place.projection = self.tcx.mk_place_elems(&new_projection);
+        }
+    }
+
+    fn process_projection_elem(
+        &mut self,
+        elem: PlaceElem<'tcx>,
+        location: Location,
+    ) -> Option<PlaceElem<'tcx>> {
+        match elem {
+            PlaceElem::Index(local) => {
+                if let Some(&Some((ty, variant, idx))) = self.remap.get(local) {
+                    // `PlaceElem::Index` only accepts a `Local`, not an arbitrary `Place`.
+                    // If the local in indexing was saved across a yield point and remapped to a
+                    // coroutine struct field, we cannot inline the struct field access into
+                    // the index projection.
+                    // For example, an local storing the counter to track which element to drop in
+                    // an array is one such case.
+                    //
+                    // Instead, we inject an assignment before this location to restore the
+                    // saved local from the coroutine struct (`local = copy $projection`),
+                    // and leave the `PlaceElem::Index(local)` projection unchanged.
+                    let field = self.make_field(variant, idx, ty);
+                    self.patch.as_mut().unwrap().add_assign(
+                        location,
+                        Place::from(local),
+                        Rvalue::Use(Operand::Copy(field), WithRetag::No),
+                    );
+                }
+                None
+            }
+            PlaceElem::Field(..)
+            | PlaceElem::OpaqueCast(..)
+            | PlaceElem::UnwrapUnsafeBinder(..)
+            | PlaceElem::Deref
+            | PlaceElem::ConstantIndex { .. }
+            | PlaceElem::Subslice { .. }
+            | PlaceElem::Downcast(..) => None,
+        }
     }
 
     #[tracing::instrument(level = "trace", skip(self, stmt), ret)]
@@ -1096,6 +1139,7 @@ impl<'tcx> crate::MirPass<'tcx> for StateTransform {
             new_ret_local,
             old_ret_ty,
             old_yield_ty,
+            patch: Some(MirPatch::new(body)),
         };
         transform.visit_body(body);
 
@@ -1116,6 +1160,7 @@ impl<'tcx> crate::MirPass<'tcx> for StateTransform {
                 Some(Statement::new(source_info, assign))
             }),
         );
+        transform.patch.take().unwrap().apply(body);
 
         // Remove the context argument within generator bodies.
         if matches!(coroutine_kind, CoroutineKind::Desugared(CoroutineDesugaring::Gen, _)) {

tests/mir-opt/coroutine/async_drop.array-{closure#0}.ElaborateDrops.diff

@@ -0,0 +1,220 @@
+- // MIR for `array::{closure#0}` before ElaborateDrops
++ // MIR for `array::{closure#0}` after ElaborateDrops
+
+  fn array::{closure#0}(_1: {async fn body of array()}, _2: std::future::ResumeTy) -> ()
+  yields ()
+   {
+      debug _task_context => _2;
+      let mut _0: ();
+      let _3: [AsyncInt; 2];
+      let mut _4: AsyncInt;
+      let mut _5: AsyncInt;
++     let mut _6: impl std::future::Future<Output = ()>;
++     let mut _7: std::future::ResumeTy;
++     let mut _8: std::task::Poll<()>;
++     let mut _9: isize;
++     let mut _10: std::pin::Pin<&mut impl std::future::Future<Output = ()>>;
++     let mut _11: &mut std::task::Context<'_>;
++     let mut _12: std::future::ResumeTy;
++     let mut _13: &mut impl std::future::Future<Output = ()>;
++     let mut _14: std::future::ResumeTy;
++     let mut _15: std::task::Poll<()>;
++     let mut _16: isize;
++     let mut _17: std::pin::Pin<&mut impl std::future::Future<Output = ()>>;
++     let mut _18: &mut std::task::Context<'_>;
++     let mut _19: std::future::ResumeTy;
++     let mut _20: &mut impl std::future::Future<Output = ()>;
++     let mut _21: std::pin::Pin<&mut [AsyncInt; 2]>;
++     let mut _22: &mut [AsyncInt; 2];
+      scope 1 {
+          debug array => _3;
+      }
+
+      bb0: {
+          StorageLive(_3);
+          StorageLive(_4);
+          _4 = AsyncInt(const 1_i32);
+          StorageLive(_5);
+          _5 = AsyncInt(const 2_i32);
+          _3 = [move _4, move _5];
+-         drop(_5) -> [return: bb1, unwind: bb9, drop: bb5];
++         goto -> bb1;
+      }
+
+      bb1: {
+          StorageDead(_5);
+-         drop(_4) -> [return: bb2, unwind: bb10, drop: bb6];
++         goto -> bb2;
+      }
+
+      bb2: {
+          StorageDead(_4);
+          _0 = const ();
+-         drop(_3) -> [return: bb3, unwind: bb11, drop: bb7];
++         goto -> bb34;
+      }
+
+      bb3: {
+          StorageDead(_3);
+-         drop(_1) -> [return: bb4, drop: bb8, unwind continue];
++         drop(_1) -> [return: bb4, unwind: bb12];
+      }
+
+      bb4: {
+          return;
+      }
+
+      bb5: {
+          StorageDead(_5);
+-         drop(_4) -> [return: bb6, unwind: bb13];
++         goto -> bb6;
+      }
+
+      bb6: {
+          StorageDead(_4);
+          goto -> bb7;
+      }
+
+      bb7: {
+          StorageDead(_3);
+-         drop(_1) -> [return: bb8, unwind continue];
++         goto -> bb8;
+      }
+
+      bb8: {
+          coroutine_drop;
+      }
+
+      bb9 (cleanup): {
+          StorageDead(_5);
+-         drop(_4) -> [return: bb10, unwind terminate(cleanup)];
++         goto -> bb10;
+      }
+
+      bb10 (cleanup): {
+          StorageDead(_4);
+          goto -> bb11;
+      }
+
+      bb11 (cleanup): {
+          StorageDead(_3);
+          drop(_1) -> [return: bb12, unwind terminate(cleanup)];
+      }
+
+      bb12 (cleanup): {
+          resume;
+      }
+
+      bb13 (cleanup): {
+          StorageDead(_4);
+          StorageDead(_3);
+-         drop(_1) -> [return: bb12, unwind terminate(cleanup)];
++         goto -> bb12;
++     }
++ 
++     bb14: {
++         StorageDead(_6);
++         goto -> bb3;
++     }
++ 
++     bb15: {
++         StorageDead(_6);
++         goto -> bb7;
++     }
++ 
++     bb16 (cleanup): {
++         StorageDead(_6);
++         goto -> bb11;
++     }
++ 
++     bb17: {
++         assert(const false, "`async fn` resumed after async drop") -> [success: bb17, unwind: bb16];
++     }
++ 
++     bb18: {
++         _2 = move _7;
++         StorageDead(_7);
++         goto -> bb17;
++     }
++ 
++     bb19: {
++         _2 = move _7;
++         StorageDead(_7);
++         goto -> bb25;
++     }
++ 
++     bb20: {
++         StorageLive(_7);
++         _7 = yield(const ()) -> [resume: bb18, drop: bb19];
++     }
++ 
++     bb21: {
++         unreachable;
++     }
++ 
++     bb22: {
++         _9 = discriminant(_8);
++         switchInt(move _9) -> [0: bb15, 1: bb20, otherwise: bb21];
++     }
++ 
++     bb23: {
++         _8 = <impl Future<Output = ()> as Future>::poll(move _10, move _11) -> [return: bb22, unwind: bb16];
++     }
++ 
++     bb24: {
++         _12 = move _2;
++         _11 = std::future::get_context::<'_, '_>(move _12) -> [return: bb23, unwind: bb16];
++     }
++ 
++     bb25: {
++         _13 = &mut _6;
++         _10 = Pin::<&mut impl Future<Output = ()>>::new_unchecked(move _13) -> [return: bb24, unwind: bb16];
++     }
++ 
++     bb26: {
++         _2 = move _14;
++         StorageDead(_14);
++         goto -> bb32;
++     }
++ 
++     bb27: {
++         _2 = move _14;
++         StorageDead(_14);
++         goto -> bb25;
++     }
++ 
++     bb28: {
++         StorageLive(_14);
++         _14 = yield(const ()) -> [resume: bb26, drop: bb27];
++     }
++ 
++     bb29: {
++         _16 = discriminant(_15);
++         switchInt(move _16) -> [0: bb14, 1: bb28, otherwise: bb21];
++     }
++ 
++     bb30: {
++         _15 = <impl Future<Output = ()> as Future>::poll(move _17, move _18) -> [return: bb29, unwind: bb16];
++     }
++ 
++     bb31: {
++         _19 = move _2;
++         _18 = std::future::get_context::<'_, '_>(move _19) -> [return: bb30, unwind: bb16];
++     }
++ 
++     bb32: {
++         _20 = &mut _6;
++         _17 = Pin::<&mut impl Future<Output = ()>>::new_unchecked(move _20) -> [return: bb31, unwind: bb16];
++     }
++ 
++     bb33: {
++         StorageLive(_6);
++         _6 = async_drop_in_place::<[AsyncInt; 2]>(copy (_21.0: &mut [AsyncInt; 2])) -> [return: bb32, unwind: bb16];
++     }
++ 
++     bb34: {
++         _22 = &mut _3;
++         _21 = Pin::<&mut [AsyncInt; 2]>::new_unchecked(move _22) -> [return: bb33, unwind: bb11];
+      }
+  }
+