Simplify type computation.

Author: cjgillot

compiler/rustc_mir_transform/src/elaborate_drop.rs

@@ -6,9 +6,9 @@ use rustc_index::Idx;
 use rustc_middle::mir::*;
 use rustc_middle::ty::adjustment::PointerCoercion;
 use rustc_middle::ty::util::IntTypeExt;
-use rustc_middle::ty::{self, GenericArg, GenericArgsRef, Ty, TyCtxt};
+use rustc_middle::ty::{self, GenericArgsRef, Ty, TyCtxt};
 use rustc_middle::{bug, span_bug};
-use rustc_span::{DUMMY_SP, Spanned, dummy_spanned};
+use rustc_span::{DUMMY_SP, dummy_spanned};
 use tracing::{debug, instrument};
 
 use crate::patch::MirPatch;
@@ -216,6 +216,7 @@ where
     ) -> BasicBlock {
         let tcx = self.tcx();
         let span = self.source_info.span;
+        let obj_ref_ty = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, drop_ty);
 
         let async_drop_fn_def_id = if call_destructor_only {
             // Resolving obj.<AsyncDrop::drop>()
@@ -233,37 +234,9 @@ where
             .output();
         let fut = self.new_temp(fut_ty);
 
-        // #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 assign_obj_ref_place = self.assign(
-            obj_ref_place,
-            Rvalue::Ref(
-                tcx.lifetimes.re_erased,
-                BorrowKind::Mut { kind: MutBorrowKind::Default },
-                place,
-            ),
-        );
-
-        // pin_obj_place preparation
-        let pin_obj_new_unchecked_fn = Ty::new_fn_def(
-            tcx,
-            tcx.require_lang_item(LangItem::PinNewUnchecked, span),
-            [GenericArg::from(obj_ref_ty)],
-        );
-        let pin_obj_ty = pin_obj_new_unchecked_fn.fn_sig(tcx).output().no_bound_vars().unwrap();
-        let pin_obj_place = Place::from(self.new_temp(pin_obj_ty));
-        let pin_obj_new_unchecked_fn = Operand::Constant(Box::new(ConstOperand {
-            span,
-            user_ty: None,
-            const_: Const::zero_sized(pin_obj_new_unchecked_fn),
-        }));
-
         // Create an intermediate block that does StorageDead(fut) then jumps to succ.
-        // This is necessary because `succ` may differ from `self.succ` (e.g. when
-        // build_async_drop is called from drop_loop, `succ` is the loop header).
-        // Placing StorageDead directly at `self.succ` would miss the loop-back edge,
-        // causing StorageLive(fut) to fire again without a preceding StorageDead.
+        // This is necessary because we do not want to modify statements
+        // in existing blocks, in case those are used somewhere else in MIR.
         let succ_with_dead = self.new_block_with_statements(
             unwind,
             vec![self.storage_dead(fut)],
@@ -297,35 +270,24 @@ where
             },
         );
 
-        // #2:call_drop_bb
-        let mut call_statements = Vec::new();
+        // #2:call_drop_bb >>> call AsyncDrop::drop(pin_obj) OR call async_drop_in_place(pin_obj.pointer)
+        let pin_adt_def = tcx.adt_def(tcx.require_lang_item(LangItem::Pin, span));
+        let pin_obj_ty = Ty::new_adt(tcx, pin_adt_def, tcx.mk_args(&[obj_ref_ty.into()]));
+        // Where we store the result of Pin<&drop_ty>::new_unchecked(&mut place).
+        let pin_obj_local = self.new_temp(pin_obj_ty);
         let drop_arg = if call_destructor_only {
-            pin_obj_place
+            // `AsyncDrop::drop` takes `self: Pin<&mut Self>`.
+            Operand::Move(pin_obj_local.into())
         } else {
-            let ty::Adt(adt_def, adt_args) = pin_obj_ty.kind() else {
-                bug!();
-            };
-            let unwrap_ty =
-                adt_def.non_enum_variant().fields[FieldIdx::ZERO].ty(tcx, adt_args).skip_norm_wip();
-            let obj_ref_place = Place::from(self.new_temp(unwrap_ty));
-            call_statements.push(self.assign(
-                obj_ref_place,
-                Rvalue::Use(
-                    Operand::Copy(tcx.mk_place_field(pin_obj_place, FieldIdx::ZERO, unwrap_ty)),
-                    WithRetag::Yes,
-                ),
-            ));
-
-            obj_ref_place
+            // `async_drop_in_place` takes `obj: &mut T`.
+            Operand::Copy(tcx.mk_place_field(pin_obj_local.into(), FieldIdx::ZERO, obj_ref_ty))
         };
-        call_statements.push(self.storage_live(fut));
-
         let call_drop_bb = self.new_block_with_statements(
             unwind,
-            call_statements,
+            vec![self.storage_live(fut)],
             TerminatorKind::Call {
                 func: Operand::function_handle(tcx, async_drop_fn_def_id, [drop_ty.into()], span),
-                args: [Spanned { node: Operand::Move(drop_arg), span: DUMMY_SP }].into(),
+                args: [dummy_spanned(drop_arg)].into(),
                 destination: fut.into(),
                 target: Some(drop_term_bb),
                 unwind: unwind_with_dead.into_action(),
@@ -335,13 +297,28 @@ where
         );
 
         // #1:pin_obj_bb >>> call Pin<ObjTy>::new_unchecked(&mut obj)
+        let obj_ref_place = Place::from(self.new_temp(obj_ref_ty));
+        let pin_obj_new_unchecked_fn = tcx.require_lang_item(LangItem::PinNewUnchecked, span);
+        let assign_obj_ref_place = self.assign(
+            obj_ref_place,
+            Rvalue::Ref(
+                tcx.lifetimes.re_erased,
+                BorrowKind::Mut { kind: MutBorrowKind::Default },
+                place,
+            ),
+        );
         self.new_block_with_statements(
             unwind,
             vec![assign_obj_ref_place],
             TerminatorKind::Call {
-                func: pin_obj_new_unchecked_fn,
+                func: Operand::function_handle(
+                    tcx,
+                    pin_obj_new_unchecked_fn,
+                    [obj_ref_ty.into()],
+                    span,
+                ),
                 args: [dummy_spanned(Operand::Move(obj_ref_place))].into(),
-                destination: pin_obj_place,
+                destination: pin_obj_local.into(),
                 target: Some(call_drop_bb),
                 unwind: unwind.into_action(),
                 call_source: CallSource::Misc,
@@ -928,8 +905,7 @@ where
             )],
             TerminatorKind::Call {
                 func: Operand::function_handle(tcx, drop_fn, [ty.into()], self.source_info.span),
-                args: [Spanned { node: Operand::Move(Place::from(ref_place)), span: DUMMY_SP }]
-                    .into(),
+                args: [dummy_spanned(Operand::Move(Place::from(ref_place)))].into(),
                 destination: unit_temp,
                 target: Some(succ),
                 unwind: unwind.into_action(),

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

@@ -6,19 +6,17 @@ yields ()
     let mut _0: ();
     let mut _3: &mut AsyncEnum;
     let mut _4: impl std::future::Future<Output = ()>;
-    let mut _5: &mut AsyncInt;
-    let mut _6: std::pin::Pin<&mut AsyncInt>;
-    let mut _7: &mut AsyncInt;
-    let mut _8: impl std::future::Future<Output = ()>;
+    let mut _5: std::pin::Pin<&mut AsyncInt>;
+    let mut _6: &mut AsyncInt;
+    let mut _7: impl std::future::Future<Output = ()>;
+    let mut _8: std::pin::Pin<&mut AsyncInt>;
     let mut _9: &mut AsyncInt;
-    let mut _10: std::pin::Pin<&mut AsyncInt>;
-    let mut _11: &mut AsyncInt;
+    let mut _10: isize;
+    let mut _11: isize;
     let mut _12: isize;
-    let mut _13: isize;
-    let mut _14: isize;
-    let mut _15: impl std::future::Future<Output = ()>;
-    let mut _16: &mut AsyncEnum;
-    let mut _17: std::pin::Pin<&mut AsyncEnum>;
+    let mut _13: impl std::future::Future<Output = ()>;
+    let mut _14: std::pin::Pin<&mut AsyncEnum>;
+    let mut _15: &mut AsyncEnum;
 
     bb0: {
         _3 = move (_1.0: &mut AsyncEnum);
@@ -60,44 +58,42 @@ yields ()
     }
 
     bb9: {
-        _7 = copy (_6.0: &mut AsyncInt);
         StorageLive(_4);
-        _4 = async_drop_in_place::<AsyncInt>(move _7) -> [return: bb8, unwind: bb7];
+        _4 = async_drop_in_place::<AsyncInt>(copy (_5.0: &mut AsyncInt)) -> [return: bb8, unwind: bb7];
     }
 
     bb10: {
-        _5 = &mut (((*_3) as A).0: AsyncInt);
-        _6 = Pin::<&mut AsyncInt>::new_unchecked(move _5) -> [return: bb9, unwind: bb2];
+        _6 = &mut (((*_3) as A).0: AsyncInt);
+        _5 = Pin::<&mut AsyncInt>::new_unchecked(move _6) -> [return: bb9, unwind: bb2];
     }
 
     bb11: {
-        StorageDead(_8);
+        StorageDead(_7);
         goto -> bb3;
     }
 
     bb12: {
-        StorageDead(_8);
+        StorageDead(_7);
         goto -> bb1;
     }
 
     bb13 (cleanup): {
-        StorageDead(_8);
+        StorageDead(_7);
         goto -> bb2;
     }
 
     bb14: {
-        async drop((((*_3) as A).0: AsyncInt); poll=_8) -> [return: bb11, unwind: bb13, drop: bb12];
+        async drop((((*_3) as A).0: AsyncInt); poll=_7) -> [return: bb11, unwind: bb13, drop: bb12];
     }
 
     bb15: {
-        _11 = copy (_10.0: &mut AsyncInt);
-        StorageLive(_8);
-        _8 = async_drop_in_place::<AsyncInt>(move _11) -> [return: bb14, unwind: bb13];
+        StorageLive(_7);
+        _7 = async_drop_in_place::<AsyncInt>(copy (_8.0: &mut AsyncInt)) -> [return: bb14, unwind: bb13];
     }
 
     bb16: {
         _9 = &mut (((*_3) as A).0: AsyncInt);
-        _10 = Pin::<&mut AsyncInt>::new_unchecked(move _9) -> [return: bb15, unwind: bb2];
+        _8 = Pin::<&mut AsyncInt>::new_unchecked(move _9) -> [return: bb15, unwind: bb2];
     }
 
     bb17 (cleanup): {
@@ -117,46 +113,46 @@ yields ()
     }
 
     bb21: {
-        _12 = discriminant((*_3));
-        switchInt(move _12) -> [0: bb16, otherwise: bb20];
+        _10 = discriminant((*_3));
+        switchInt(move _10) -> [0: bb16, otherwise: bb20];
     }
 
     bb22 (cleanup): {
-        _13 = discriminant((*_3));
-        switchInt(move _13) -> [0: bb4, otherwise: bb17];
+        _11 = discriminant((*_3));
+        switchInt(move _11) -> [0: bb4, otherwise: bb17];
     }
 
     bb23: {
-        _14 = discriminant((*_3));
-        switchInt(move _14) -> [0: bb10, otherwise: bb19];
+        _12 = discriminant((*_3));
+        switchInt(move _12) -> [0: bb10, otherwise: bb19];
     }
 
     bb24: {
-        StorageDead(_15);
+        StorageDead(_13);
         goto -> bb21;
     }
 
     bb25: {
-        StorageDead(_15);
+        StorageDead(_13);
         goto -> bb23;
     }
 
     bb26 (cleanup): {
-        StorageDead(_15);
+        StorageDead(_13);
         goto -> bb22;
     }
 
     bb27: {
-        async drop((*_3); poll=_15) -> [return: bb24, unwind: bb26, drop: bb25];
+        async drop((*_3); poll=_13) -> [return: bb24, unwind: bb26, drop: bb25];
     }
 
     bb28: {
-        StorageLive(_15);
-        _15 = <AsyncEnum as AsyncDrop>::drop(move _17) -> [return: bb27, unwind: bb26];
+        StorageLive(_13);
+        _13 = <AsyncEnum as AsyncDrop>::drop(move _14) -> [return: bb27, unwind: bb26];
     }
 
     bb29: {
-        _16 = &mut (*_3);
-        _17 = Pin::<&mut AsyncEnum>::new_unchecked(move _16) -> [return: bb28, unwind: bb22];
+        _15 = &mut (*_3);
+        _14 = Pin::<&mut AsyncEnum>::new_unchecked(move _15) -> [return: bb28, unwind: bb22];
     }
 }