Avoid modifying succ block in build_async_drop.

Author: cjgillot

compiler/rustc_mir_transform/src/elaborate_drop.rs

@@ -96,8 +96,6 @@ pub(crate) trait DropElaborator<'a, 'tcx>: fmt::Debug {
     fn typing_env(&self) -> ty::TypingEnv<'tcx>;
     fn allow_async_drops(&self) -> bool;
 
-    fn terminator_loc(&self, bb: BasicBlock) -> Location;
-
     // Drop logic
 
     /// Returns how `path` should be dropped, given `mode`.
@@ -200,19 +198,18 @@ where
         self.elaborator.tcx()
     }
 
-    // Generates three blocks:
-    // * #1:pin_obj_bb:   call Pin<ObjTy>::new_unchecked(&mut obj)
-    // * #2:call_drop_bb: fut = call obj.<AsyncDrop::drop>() OR call async_drop_in_place<T>(obj)
-    // * #3:drop_term_bb: drop (obj, fut, ...)
-    // We keep async drop unexpanded to poll-loop here, to expand it later, at StateTransform -
-    //   into states expand.
-    // call_destructor_only - to call only AsyncDrop::drop, not full async_drop_in_place glue
+    /// Generates three blocks:
+    /// * #1:pin_obj_bb:   call Pin<ObjTy>::new_unchecked(&mut obj)
+    /// * #2:call_drop_bb: fut = call obj.<AsyncDrop::drop>() OR call async_drop_in_place<T>(obj)
+    /// * #3:drop_term_bb: drop (obj, fut, ...)
+    /// We keep async drop unexpanded to poll-loop here, to expand it later, at StateTransform -
+    ///   into states expand.
+    /// call_destructor_only - to call only AsyncDrop::drop, not full async_drop_in_place glue
     #[instrument(level = "debug", skip(self), ret)]
     fn build_async_drop(
         &mut self,
         place: Place<'tcx>,
         drop_ty: Ty<'tcx>,
-        bb: Option<BasicBlock>,
         succ: BasicBlock,
         unwind: Unwind,
         dropline: Option<BasicBlock>,
@@ -221,11 +218,6 @@ where
         let tcx = self.tcx();
         let span = self.source_info.span;
 
-        let pin_obj_bb = bb.unwrap_or_else(|| {
-            // Temporary terminator, will be replaced by patch
-            self.new_block(unwind, TerminatorKind::Return)
-        });
-
         let (fut_ty, drop_fn_def_id, trait_args) = if call_destructor_only {
             // Resolving obj.<AsyncDrop::drop>()
             let trait_ref =
@@ -260,8 +252,8 @@ where
                     self.elaborator.body().span,
                     "AsyncDrop type without correct `async fn drop(...)`.",
                 );
-                self.elaborator.patch().patch_terminator(
-                    pin_obj_bb,
+                return self.new_block(
+                    unwind,
                     TerminatorKind::Drop {
                         place,
                         target: succ,
@@ -271,7 +263,6 @@ where
                         async_fut: None,
                     },
                 );
-                return pin_obj_bb;
             };
             let drop_fn = Ty::new_fn_def(tcx, drop_fn_def_id, trait_args);
             let sig = drop_fn.fn_sig(tcx);
@@ -291,10 +282,7 @@ where
         // #1:pin_obj_bb >>> obj_ref = &mut obj
         let obj_ref_ty = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, drop_ty);
         let obj_ref_place = Place::from(self.new_temp(obj_ref_ty));
-
-        let term_loc = self.elaborator.terminator_loc(pin_obj_bb);
-        self.elaborator.patch().add_assign(
-            term_loc,
+        let assign_obj_ref_place = self.assign(
             obj_ref_place,
             Rvalue::Ref(
                 tcx.lifetimes.re_erased,
@@ -394,8 +382,9 @@ where
         }
 
         // #1:pin_obj_bb >>> call Pin<ObjTy>::new_unchecked(&mut obj)
-        self.elaborator.patch().patch_terminator(
-            pin_obj_bb,
+        self.new_block_with_statements(
+            unwind,
+            vec![assign_obj_ref_place],
             TerminatorKind::Call {
                 func: pin_obj_new_unchecked_fn,
                 args: [dummy_spanned(Operand::Move(obj_ref_place))].into(),
@@ -405,24 +394,25 @@ where
                 call_source: CallSource::Misc,
                 fn_span: span,
             },
-        );
-        pin_obj_bb
+        )
     }
 
     fn build_drop(&mut self, bb: BasicBlock) {
         let drop_ty = self.place_ty(self.place);
         if !self.elaborator.patch_ref().block(self.elaborator.body(), bb).is_cleanup
             && self.check_if_can_async_drop(drop_ty, false)
         {
-            self.build_async_drop(
+            let async_drop_bb = self.build_async_drop(
                 self.place,
                 drop_ty,
-                Some(bb),
                 self.succ,
                 self.unwind,
                 self.dropline,
                 false,
             );
+            self.elaborator
+                .patch()
+                .patch_terminator(bb, TerminatorKind::Goto { target: async_drop_bb });
         } else {
             self.elaborator.patch().patch_terminator(
                 bb,
@@ -1005,7 +995,7 @@ where
     ) -> BasicBlock {
         let ty = self.place_ty(self.place);
         if !unwind.is_cleanup() && self.check_if_can_async_drop(ty, true) {
-            self.build_async_drop(self.place, ty, None, succ, unwind, dropline, true)
+            self.build_async_drop(self.place, ty, succ, unwind, dropline, true)
         } else {
             self.destructor_call_block_sync(succ, unwind)
         }
@@ -1067,15 +1057,11 @@ where
 
         let place = tcx.mk_place_deref(ptr);
         if !unwind.is_cleanup() && self.check_if_can_async_drop(ety, false) {
-            self.build_async_drop(
-                place,
-                ety,
-                Some(drop_block),
-                loop_block,
-                unwind,
-                dropline,
-                false,
-            );
+            let async_drop_bb =
+                self.build_async_drop(place, ety, loop_block, unwind, dropline, false);
+            self.elaborator
+                .patch()
+                .patch_terminator(drop_block, TerminatorKind::Goto { target: async_drop_bb });
         } else {
             self.elaborator.patch().patch_terminator(
                 drop_block,
@@ -1331,15 +1317,7 @@ where
     fn drop_block(&mut self, target: BasicBlock, unwind: Unwind) -> BasicBlock {
         let drop_ty = self.place_ty(self.place);
         if !unwind.is_cleanup() && self.check_if_can_async_drop(drop_ty, false) {
-            self.build_async_drop(
-                self.place,
-                drop_ty,
-                None,
-                self.succ,
-                unwind,
-                self.dropline,
-                false,
-            )
+            self.build_async_drop(self.place, drop_ty, self.succ, unwind, self.dropline, false)
         } else {
             self.new_block(
                 unwind,

compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -162,10 +162,6 @@ impl<'a, 'tcx> DropElaborator<'a, 'tcx> for ElaborateDropsCtxt<'a, 'tcx> {
         true
     }
 
-    fn terminator_loc(&self, bb: BasicBlock) -> Location {
-        self.patch.terminator_loc(self.body, bb)
-    }
-
     #[instrument(level = "debug", skip(self), ret)]
     fn drop_style(&self, path: Self::Path, mode: DropFlagMode) -> DropStyle {
         let ((maybe_init, maybe_uninit), multipart) = match mode {

compiler/rustc_mir_transform/src/patch.rs

@@ -166,11 +166,6 @@ impl<'tcx> MirPatch<'tcx> {
         }
     }
 
-    pub(crate) fn terminator_loc(&self, body: &Body<'tcx>, bb: BasicBlock) -> Location {
-        let offset = self.block(body, bb).statements.len();
-        Location { block: bb, statement_index: offset }
-    }
-
     /// Queues the addition of a new temporary with additional local info.
     pub(crate) fn new_local_with_info(
         &mut self,

compiler/rustc_mir_transform/src/shim.rs

@@ -429,9 +429,6 @@ impl<'a, 'tcx> DropElaborator<'a, 'tcx> for DropShimElaborator<'a, 'tcx> {
         self.typing_env
     }
 
-    fn terminator_loc(&self, bb: BasicBlock) -> Location {
-        self.patch.terminator_loc(self.body, bb)
-    }
     fn allow_async_drops(&self) -> bool {
         self.produce_async_drops
     }

tests/mir-opt/coroutine/async_drop.core.future-async_drop-async_drop_in_place-{closure#0}.AsyncEnum.MentionedItems.after.mir

@@ -22,7 +22,7 @@ yields ()
 
     bb0: {
         _3 = move (_1.0: &mut AsyncEnum);
-        goto -> bb21;
+        goto -> bb24;
     }
 
     bb1: {
@@ -49,43 +49,43 @@ yields ()
     }
 
     bb6: {
-        _5 = &mut (((*_3) as A).0: AsyncInt);
-        _6 = Pin::<&mut AsyncInt>::new_unchecked(move _5) -> [return: bb9, unwind: bb2];
+        StorageDead(_4);
+        goto -> bb1;
     }
 
     bb7: {
-        StorageDead(_4);
-        goto -> bb1;
+        async drop((((*_3) as A).0: AsyncInt); poll=_4) -> [return: bb6, unwind: bb2];
     }
 
     bb8: {
-        async drop((((*_3) as A).0: AsyncInt); poll=_4) -> [return: bb7, unwind: bb2];
+        _7 = copy (_6.0: &mut AsyncInt);
+        StorageLive(_4);
+        _4 = async_drop_in_place::<AsyncInt>(move _7) -> [return: bb7, unwind: bb2];
     }
 
     bb9: {
-        _7 = copy (_6.0: &mut AsyncInt);
-        StorageLive(_4);
-        _4 = async_drop_in_place::<AsyncInt>(move _7) -> [return: bb8, unwind: bb2];
+        _5 = &mut (((*_3) as A).0: AsyncInt);
+        _6 = Pin::<&mut AsyncInt>::new_unchecked(move _5) -> [return: bb8, unwind: bb2];
     }
 
     bb10: {
-        _9 = &mut (((*_3) as A).0: AsyncInt);
-        _10 = Pin::<&mut AsyncInt>::new_unchecked(move _9) -> [return: bb13, unwind: bb2];
+        StorageDead(_8);
+        goto -> bb3;
     }
 
     bb11: {
-        StorageDead(_8);
-        goto -> bb3;
+        async drop((((*_3) as A).0: AsyncInt); poll=_8) -> [return: bb10, unwind: bb2, drop: bb1];
     }
 
     bb12: {
-        async drop((((*_3) as A).0: AsyncInt); poll=_8) -> [return: bb11, unwind: bb2, drop: bb1];
+        _11 = copy (_10.0: &mut AsyncInt);
+        StorageLive(_8);
+        _8 = async_drop_in_place::<AsyncInt>(move _11) -> [return: bb11, unwind: bb2];
     }
 
     bb13: {
-        _11 = copy (_10.0: &mut AsyncInt);
-        StorageLive(_8);
-        _8 = async_drop_in_place::<AsyncInt>(move _11) -> [return: bb12, unwind: bb2];
+        _9 = &mut (((*_3) as A).0: AsyncInt);
+        _10 = Pin::<&mut AsyncInt>::new_unchecked(move _9) -> [return: bb12, unwind: bb2];
     }
 
     bb14 (cleanup): {
@@ -106,7 +106,7 @@ yields ()
 
     bb18: {
         _12 = discriminant((*_3));
-        switchInt(move _12) -> [0: bb10, otherwise: bb17];
+        switchInt(move _12) -> [0: bb13, otherwise: bb17];
     }
 
     bb19 (cleanup): {
@@ -118,25 +118,25 @@ yields ()
     bb20: {
         StorageDead(_15);
         _14 = discriminant((*_3));
-        switchInt(move _14) -> [0: bb6, otherwise: bb16];
+        switchInt(move _14) -> [0: bb9, otherwise: bb16];
     }
 
     bb21: {
-        _16 = &mut (*_3);
-        _17 = Pin::<&mut AsyncEnum>::new_unchecked(move _16) -> [return: bb24, unwind: bb19];
+        StorageDead(_15);
+        goto -> bb18;
     }
 
     bb22: {
-        StorageDead(_15);
-        goto -> bb18;
+        async drop((*_3); poll=_15) -> [return: bb21, unwind: bb19, drop: bb20];
     }
 
     bb23: {
-        async drop((*_3); poll=_15) -> [return: bb22, unwind: bb19, drop: bb20];
+        StorageLive(_15);
+        _15 = <AsyncEnum as AsyncDrop>::drop(move _17) -> [return: bb22, unwind: bb19];
     }
 
     bb24: {
-        StorageLive(_15);
-        _15 = <AsyncEnum as AsyncDrop>::drop(move _17) -> [return: bb23, unwind: bb19];
+        _16 = &mut (*_3);
+        _17 = Pin::<&mut AsyncEnum>::new_unchecked(move _16) -> [return: bb23, unwind: bb19];
     }
 }

tests/mir-opt/coroutine/async_drop.core.future-async_drop-async_drop_in_place-{closure#0}.AsyncEnum.StateTransform.after.mir

@@ -100,25 +100,25 @@ fn async_drop_in_place::{closure#0}(_1: Pin<&mut {async fn body of async_drop_in
     }
 
     bb5: {
-        _46 = no_retag copy (((*_43) as variant#7).0: &mut AsyncEnum);
-        _9 = &mut (((*_46) as A).0: AsyncInt);
-        _10 = Pin::<&mut AsyncInt>::new_unchecked(move _9) -> [return: bb8, unwind: bb3];
-    }
-
-    bb6: {
         nop;
         goto -> bb2;
     }
 
-    bb7: {
+    bb6: {
         _24 = move _2;
         goto -> bb17;
     }
 
-    bb8: {
+    bb7: {
         _11 = copy (_10.0: &mut AsyncInt);
         nop;
-        (((*_43) as variant#5).0: impl std::future::Future<Output = ()>) = async_drop_in_place::<AsyncInt>(move _11) -> [return: bb7, unwind: bb3];
+        (((*_43) as variant#5).0: impl std::future::Future<Output = ()>) = async_drop_in_place::<AsyncInt>(move _11) -> [return: bb6, unwind: bb3];
+    }
+
+    bb8: {
+        _46 = no_retag copy (((*_43) as variant#7).0: &mut AsyncEnum);
+        _9 = &mut (((*_46) as A).0: AsyncInt);
+        _10 = Pin::<&mut AsyncInt>::new_unchecked(move _9) -> [return: bb7, unwind: bb3];
     }
 
     bb9 (cleanup): {
@@ -142,7 +142,7 @@ fn async_drop_in_place::{closure#0}(_1: Pin<&mut {async fn body of async_drop_in
         nop;
         _50 = no_retag copy (((*_43) as variant#7).0: &mut AsyncEnum);
         _12 = discriminant((*_50));
-        switchInt(move _12) -> [0: bb5, otherwise: bb10];
+        switchInt(move _12) -> [0: bb8, otherwise: bb10];
     }
 
     bb13: {
@@ -178,7 +178,7 @@ fn async_drop_in_place::{closure#0}(_1: Pin<&mut {async fn body of async_drop_in
     bb19: {
         StorageDead(_26);
         _27 = discriminant(_23);
-        switchInt(move _27) -> [0: bb6, 1: bb16, otherwise: bb18];
+        switchInt(move _27) -> [0: bb5, 1: bb16, otherwise: bb18];
     }
 
     bb20: {