Prereq1 for async drop - added drop & async_fut Option fields in Drop terminator

Author: azhogin

compiler/rustc_borrowck/src/lib.rs

@@ -690,7 +690,14 @@ impl<'a, 'mir, 'tcx, R> rustc_mir_dataflow::ResultsVisitor<'mir, 'tcx, R>
             TerminatorKind::SwitchInt { discr, targets: _ } => {
                 self.consume_operand(loc, (discr, span), flow_state);
             }
-            TerminatorKind::Drop { place, target: _, unwind: _, replace } => {
+            TerminatorKind::Drop {
+                place,
+                target: _,
+                unwind: _,
+                replace,
+                drop: _,
+                async_fut: _,
+            } => {
                 debug!(
                     "visit_terminator_drop \
                      loc: {:?} term: {:?} place: {:?} span: {:?}",

compiler/rustc_borrowck/src/polonius/loan_invalidations.rs

@@ -103,7 +103,14 @@ impl<'cx, 'tcx> Visitor<'tcx> for LoanInvalidationsGenerator<'cx, 'tcx> {
             TerminatorKind::SwitchInt { discr, targets: _ } => {
                 self.consume_operand(location, discr);
             }
-            TerminatorKind::Drop { place: drop_place, target: _, unwind: _, replace } => {
+            TerminatorKind::Drop {
+                place: drop_place,
+                target: _,
+                unwind: _,
+                replace,
+                drop: _,
+                async_fut: _,
+            } => {
                 let write_kind =
                     if *replace { WriteKind::Replace } else { WriteKind::StorageDeadOrDrop };
                 self.access_place(

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -1700,8 +1700,14 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 }
             }
             TerminatorKind::Unreachable => {}
-            TerminatorKind::Drop { target, unwind, .. }
-            | TerminatorKind::Assert { target, unwind, .. } => {
+            TerminatorKind::Drop { target, unwind, drop, .. } => {
+                self.assert_iscleanup(body, block_data, target, is_cleanup);
+                self.assert_iscleanup_unwind(body, block_data, unwind, is_cleanup);
+                if let Some(drop) = drop {
+                    self.assert_iscleanup(body, block_data, drop, is_cleanup);
+                }
+            }
+            TerminatorKind::Assert { target, unwind, .. } => {
                 self.assert_iscleanup(body, block_data, target, is_cleanup);
                 self.assert_iscleanup_unwind(body, block_data, unwind, is_cleanup);
             }

compiler/rustc_codegen_cranelift/src/base.rs

@@ -549,7 +549,11 @@ fn codegen_fn_body(fx: &mut FunctionCx<'_, '_, '_>, start_block: Block) {
             | TerminatorKind::CoroutineDrop => {
                 bug!("shouldn't exist at codegen {:?}", bb_data.terminator());
             }
-            TerminatorKind::Drop { place, target, unwind: _, replace: _ } => {
+            TerminatorKind::Drop { place, target, unwind: _, replace: _, drop, async_fut } => {
+                assert!(
+                    async_fut.is_none() && drop.is_none(),
+                    "Async Drop must be expanded or reset to sync before codegen"
+                );
                 let drop_place = codegen_place(fx, *place);
                 crate::abi::codegen_drop(fx, source_info, drop_place, *target);
             }

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -1347,16 +1347,21 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 MergingSucc::False
             }
 
-            mir::TerminatorKind::Drop { place, target, unwind, replace: _ } => self
-                .codegen_drop_terminator(
+            mir::TerminatorKind::Drop { place, target, unwind, replace: _, drop, async_fut } => {
+                assert!(
+                    async_fut.is_none() && drop.is_none(),
+                    "Async Drop must be expanded or reset to sync before codegen"
+                );
+                self.codegen_drop_terminator(
                     helper,
                     bx,
                     &terminator.source_info,
                     place,
                     target,
                     unwind,
                     mergeable_succ(),
-                ),
+                )
+            }
 
             mir::TerminatorKind::Assert { ref cond, expected, ref msg, target, unwind } => self
                 .codegen_assert_terminator(

compiler/rustc_const_eval/src/interpret/step.rs

@@ -528,7 +528,11 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 }
             }
 
-            Drop { place, target, unwind, replace: _ } => {
+            Drop { place, target, unwind, replace: _, drop, async_fut } => {
+                assert!(
+                    async_fut.is_none() && drop.is_none(),
+                    "Async Drop must be expanded or reset to sync in runtime MIR"
+                );
                 let place = self.eval_place(place)?;
                 let instance = Instance::resolve_drop_in_place(*self.tcx, place.layout.ty);
                 if let ty::InstanceKind::DropGlue(_, None) = instance.def {

compiler/rustc_middle/src/mir/pretty.rs

@@ -1015,7 +1015,13 @@ impl<'tcx> TerminatorKind<'tcx> {
             Call { target: None, unwind: _, .. } => vec![],
             Yield { drop: Some(_), .. } => vec!["resume".into(), "drop".into()],
             Yield { drop: None, .. } => vec!["resume".into()],
-            Drop { unwind: UnwindAction::Cleanup(_), .. } => vec!["return".into(), "unwind".into()],
+            Drop { unwind: UnwindAction::Cleanup(_), drop: Some(_), .. } => {
+                vec!["return".into(), "unwind".into(), "drop".into()]
+            }
+            Drop { unwind: UnwindAction::Cleanup(_), drop: None, .. } => {
+                vec!["return".into(), "unwind".into()]
+            }
+            Drop { unwind: _, drop: Some(_), .. } => vec!["return".into(), "drop".into()],
             Drop { unwind: _, .. } => vec!["return".into()],
             Assert { unwind: UnwindAction::Cleanup(_), .. } => {
                 vec!["success".into(), "unwind".into()]

compiler/rustc_middle/src/mir/syntax.rs

@@ -73,6 +73,8 @@ pub enum MirPhase {
     ///    is fields of packed structs. In analysis MIR, `Drop(P)` for a `P` that might be misaligned
     ///    for this reason implicitly moves `P` to a temporary before dropping. Runtime MIR has no such
     ///    rules, and dropping a misaligned place is simply UB.
+    ///  - Async drops: after drop elaboration some drops may become async (`drop`, `async_fut` fields).
+    ///    StateTransform pass will expand those async drops or reset to sync.
     ///  - Unwinding: in analysis MIR, unwinding from a function which may not unwind aborts. In runtime
     ///    MIR, this is UB.
     ///  - Retags: If `-Zmir-emit-retag` is enabled, analysis MIR has "implicit" retags in the same way
@@ -704,7 +706,34 @@ pub enum TerminatorKind<'tcx> {
     /// The `replace` flag indicates whether this terminator was created as part of an assignment.
     /// This should only be used for diagnostic purposes, and does not have any operational
     /// meaning.
-    Drop { place: Place<'tcx>, target: BasicBlock, unwind: UnwindAction, replace: bool },
+    ///
+    /// Async drop processing:
+    /// In compiler/rustc_mir_build/src/build/scope.rs we detect possible async drop:
+    ///   drop of object with `needs_async_drop`.
+    /// Async drop later, in StateTransform pass, may be expanded into additional yield-point
+    ///   for poll-loop of async drop future.
+    /// So we need prepared 'drop' target block in the similar way as for `Yield` terminator
+    ///   (see `drops.build_mir::<CoroutineDrop>` in scopes.rs).
+    /// In compiler/rustc_mir_transform/src/elaborate_drops.rs for object implementing `AsyncDrop` trait
+    ///   we need to prepare async drop feature - resolve `AsyncDrop::drop` and codegen call.
+    /// `async_fut` is set to the corresponding local.
+    /// For coroutine drop we don't need this logic because coroutine drop works with the same
+    ///   layout object as coroutine itself. So `async_fut` will be `None` for coroutine drop.
+    /// Both `drop` and `async_fut` fields are only used in compiler/rustc_mir_transform/src/coroutine.rs,
+    ///   StateTransform pass. In `expand_async_drops` async drops are expanded
+    ///   into one or two yield points with poll ready/pending switch.
+    /// When a coroutine has any internal async drop, the coroutine drop function will be async
+    ///   (generated by `create_coroutine_drop_shim_async`, not `create_coroutine_drop_shim`).
+    Drop {
+        place: Place<'tcx>,
+        target: BasicBlock,
+        unwind: UnwindAction,
+        replace: bool,
+        /// Cleanup to be done if the coroutine is dropped at this suspend point (for async drop).
+        drop: Option<BasicBlock>,
+        /// Prepared async future local (for async drop)
+        async_fut: Option<Local>,
+    },
 
     /// Roughly speaking, evaluates the `func` operand and the arguments, and starts execution of
     /// the referred to function. The operand types must match the argument types of the function.

compiler/rustc_middle/src/mir/terminator.rs

@@ -418,77 +418,123 @@ mod helper {
         pub fn successors(&self) -> Successors<'_> {
             use self::TerminatorKind::*;
             match *self {
+                // 3-successors for async drop: target, unwind, dropline (parent coroutine drop)
+                Drop { target: ref t, unwind: UnwindAction::Cleanup(u), drop: Some(d), .. } => {
+                    slice::from_ref(t)
+                        .into_iter()
+                        .copied()
+                        .chain(Some(u).into_iter().chain(Some(d)))
+                }
+                // 2-successors
                 Call { target: Some(ref t), unwind: UnwindAction::Cleanup(u), .. }
                 | Yield { resume: ref t, drop: Some(u), .. }
-                | Drop { target: ref t, unwind: UnwindAction::Cleanup(u), .. }
+                | Drop { target: ref t, unwind: UnwindAction::Cleanup(u), drop: None, .. }
+                | Drop { target: ref t, unwind: _, drop: Some(u), .. }
                 | Assert { target: ref t, unwind: UnwindAction::Cleanup(u), .. }
                 | FalseUnwind { real_target: ref t, unwind: UnwindAction::Cleanup(u) } => {
-                    slice::from_ref(t).into_iter().copied().chain(Some(u))
+                    slice::from_ref(t).into_iter().copied().chain(Some(u).into_iter().chain(None))
                 }
+                // single successor
                 Goto { target: ref t }
                 | Call { target: None, unwind: UnwindAction::Cleanup(ref t), .. }
                 | Call { target: Some(ref t), unwind: _, .. }
                 | Yield { resume: ref t, drop: None, .. }
                 | Drop { target: ref t, unwind: _, .. }
                 | Assert { target: ref t, unwind: _, .. }
                 | FalseUnwind { real_target: ref t, unwind: _ } => {
-                    slice::from_ref(t).into_iter().copied().chain(None)
+                    slice::from_ref(t).into_iter().copied().chain(None.into_iter().chain(None))
                 }
+                // No successors
                 UnwindResume
                 | UnwindTerminate(_)
                 | CoroutineDrop
                 | Return
                 | Unreachable
                 | TailCall { .. }
-                | Call { target: None, unwind: _, .. } => (&[]).into_iter().copied().chain(None),
+                | Call { target: None, unwind: _, .. } => {
+                    (&[]).into_iter().copied().chain(None.into_iter().chain(None))
+                }
+                // Multiple successors
                 InlineAsm { ref targets, unwind: UnwindAction::Cleanup(u), .. } => {
-                    targets.iter().copied().chain(Some(u))
+                    targets.iter().copied().chain(Some(u).into_iter().chain(None))
+                }
+                InlineAsm { ref targets, unwind: _, .. } => {
+                    targets.iter().copied().chain(None.into_iter().chain(None))
                 }
-                InlineAsm { ref targets, unwind: _, .. } => targets.iter().copied().chain(None),
-                SwitchInt { ref targets, .. } => targets.targets.iter().copied().chain(None),
-                FalseEdge { ref real_target, imaginary_target } => {
-                    slice::from_ref(real_target).into_iter().copied().chain(Some(imaginary_target))
+                SwitchInt { ref targets, .. } => {
+                    targets.targets.iter().copied().chain(None.into_iter().chain(None))
                 }
+                // FalseEdge
+                FalseEdge { ref real_target, imaginary_target } => slice::from_ref(real_target)
+                    .into_iter()
+                    .copied()
+                    .chain(Some(imaginary_target).into_iter().chain(None)),
             }
         }
 
         #[inline]
         pub fn successors_mut(&mut self) -> SuccessorsMut<'_> {
             use self::TerminatorKind::*;
             match *self {
+                // 3-successors for async drop: target, unwind, dropline (parent coroutine drop)
+                Drop {
+                    target: ref mut t,
+                    unwind: UnwindAction::Cleanup(ref mut u),
+                    drop: Some(ref mut d),
+                    ..
+                } => slice::from_mut(t).into_iter().chain(Some(u).into_iter().chain(Some(d))),
+                // 2-successors
                 Call {
                     target: Some(ref mut t), unwind: UnwindAction::Cleanup(ref mut u), ..
                 }
                 | Yield { resume: ref mut t, drop: Some(ref mut u), .. }
-                | Drop { target: ref mut t, unwind: UnwindAction::Cleanup(ref mut u), .. }
+                | Drop {
+                    target: ref mut t,
+                    unwind: UnwindAction::Cleanup(ref mut u),
+                    drop: None,
+                    ..
+                }
+                | Drop { target: ref mut t, unwind: _, drop: Some(ref mut u), .. }
                 | Assert { target: ref mut t, unwind: UnwindAction::Cleanup(ref mut u), .. }
                 | FalseUnwind {
                     real_target: ref mut t,
                     unwind: UnwindAction::Cleanup(ref mut u),
-                } => slice::from_mut(t).into_iter().chain(Some(u)),
+                } => slice::from_mut(t).into_iter().chain(Some(u).into_iter().chain(None)),
+                // single successor
                 Goto { target: ref mut t }
                 | Call { target: None, unwind: UnwindAction::Cleanup(ref mut t), .. }
                 | Call { target: Some(ref mut t), unwind: _, .. }
                 | Yield { resume: ref mut t, drop: None, .. }
                 | Drop { target: ref mut t, unwind: _, .. }
                 | Assert { target: ref mut t, unwind: _, .. }
                 | FalseUnwind { real_target: ref mut t, unwind: _ } => {
-                    slice::from_mut(t).into_iter().chain(None)
+                    slice::from_mut(t).into_iter().chain(None.into_iter().chain(None))
                 }
+                // No successors
                 UnwindResume
                 | UnwindTerminate(_)
                 | CoroutineDrop
                 | Return
                 | Unreachable
                 | TailCall { .. }
-                | Call { target: None, unwind: _, .. } => (&mut []).into_iter().chain(None),
+                | Call { target: None, unwind: _, .. } => {
+                    (&mut []).into_iter().chain(None.into_iter().chain(None))
+                }
+                // Multiple successors
                 InlineAsm { ref mut targets, unwind: UnwindAction::Cleanup(ref mut u), .. } => {
-                    targets.iter_mut().chain(Some(u))
+                    targets.iter_mut().chain(Some(u).into_iter().chain(None))
+                }
+                InlineAsm { ref mut targets, unwind: _, .. } => {
+                    targets.iter_mut().chain(None.into_iter().chain(None))
+                }
+                SwitchInt { ref mut targets, .. } => {
+                    targets.targets.iter_mut().chain(None.into_iter().chain(None))
                 }
-                InlineAsm { ref mut targets, unwind: _, .. } => targets.iter_mut().chain(None),
-                SwitchInt { ref mut targets, .. } => targets.targets.iter_mut().chain(None),
+                // FalseEdge
                 FalseEdge { ref mut real_target, ref mut imaginary_target } => {
-                    slice::from_mut(real_target).into_iter().chain(Some(imaginary_target))
+                    slice::from_mut(real_target)
+                        .into_iter()
+                        .chain(Some(imaginary_target).into_iter().chain(None))
                 }
             }
         }
@@ -619,8 +665,10 @@ impl<'tcx> TerminatorKind<'tcx> {
 
             Goto { target } => TerminatorEdges::Single(target),
 
+            // FIXME: Maybe we need also TerminatorEdges::Trio for async drop
+            // (target + unwind + dropline)
             Assert { target, unwind, expected: _, msg: _, cond: _ }
-            | Drop { target, unwind, place: _, replace: _ }
+            | Drop { target, unwind, place: _, replace: _, drop: _, async_fut: _ }
             | FalseUnwind { real_target: target, unwind } => match unwind {
                 UnwindAction::Cleanup(unwind) => TerminatorEdges::Double(target, unwind),
                 UnwindAction::Continue | UnwindAction::Terminate(_) | UnwindAction::Unreachable => {

compiler/rustc_middle/src/mir/visit.rs

@@ -511,6 +511,8 @@ macro_rules! make_mir_visitor {
                         target: _,
                         unwind: _,
                         replace: _,
+                        drop: _,
+                        async_fut: _,
                     } => {
                         self.visit_place(
                             place,

compiler/rustc_mir_build/src/build/custom/parse/instruction.rs

@@ -70,6 +70,8 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
                     target: self.parse_return_to(args[1])?,
                     unwind: self.parse_unwind_action(args[2])?,
                     replace: false,
+                    drop: None,
+                    async_fut: None,
                 })
             },
             @call(mir_call, args) => {

compiler/rustc_mir_build/src/build/expr/as_rvalue.rs

@@ -729,6 +729,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                         target: success,
                         unwind: UnwindAction::Continue,
                         replace: false,
+                        drop: None,
+                        async_fut: None,
                     },
                 );
                 this.diverge_from(block);

compiler/rustc_mir_build/src/build/scope.rs

@@ -401,6 +401,8 @@ impl DropTree {
                         unwind: UnwindAction::Terminate(UnwindTerminateReason::InCleanup),
                         place: drop_node.data.local.into(),
                         replace: false,
+                        drop: None,
+                        async_fut: None,
                     };
                     cfg.terminate(block, drop_node.data.source_info, terminator);
                 }
@@ -814,6 +816,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                                 target: next,
                                 unwind: UnwindAction::Continue,
                                 replace: false,
+                                drop: None,
+                                async_fut: None,
                             },
                         );
                         block = next;
@@ -1291,6 +1295,8 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 target: assign,
                 unwind: UnwindAction::Cleanup(assign_unwind),
                 replace: true,
+                drop: None,
+                async_fut: None,
             },
         );
         self.diverge_from(block);
@@ -1405,6 +1411,8 @@ fn build_scope_drops<'tcx>(
                         target: next,
                         unwind: UnwindAction::Continue,
                         replace: false,
+                        drop: None,
+                        async_fut: None,
                     },
                 );
                 block = next;