rustc_mir_transform/shim/
async_destructor_ctor.rs

1use rustc_hir::def_id::DefId;
2use rustc_hir::lang_items::LangItem;
3use rustc_hir::{CoroutineDesugaring, CoroutineKind, CoroutineSource, Safety};
4use rustc_index::{Idx, IndexVec};
5use rustc_middle::mir::{
6    BasicBlock, BasicBlockData, Body, Local, LocalDecl, MirSource, Operand, Place, Rvalue,
7    SourceInfo, Statement, StatementKind, Terminator, TerminatorKind,
8};
9use rustc_middle::ty::{self, EarlyBinder, Ty, TyCtxt, TypeVisitableExt};
10
11use super::*;
12use crate::patch::MirPatch;
13
14pub(super) fn build_async_destructor_ctor_shim<'tcx>(
15    tcx: TyCtxt<'tcx>,
16    def_id: DefId,
17    ty: Ty<'tcx>,
18) -> Body<'tcx> {
19    debug!("build_async_destructor_ctor_shim(def_id={:?}, ty={:?})", def_id, ty);
20    debug_assert_eq!(Some(def_id), tcx.lang_items().async_drop_in_place_fn());
21    let generic_body = tcx.optimized_mir(def_id);
22    let args = tcx.mk_args(&[ty.into()]);
23    let mut body = EarlyBinder::bind(generic_body.clone()).instantiate(tcx, args);
24
25    // Minimal shim passes except MentionedItems,
26    // it causes error "mentioned_items for DefId(...async_drop_in_place...) have already been set
27    pm::run_passes(
28        tcx,
29        &mut body,
30        &[
31            &simplify::SimplifyCfg::MakeShim,
32            &abort_unwinding_calls::AbortUnwindingCalls,
33            &add_call_guards::CriticalCallEdges,
34        ],
35        None,
36        pm::Optimizations::Allowed,
37    );
38    body
39}
40
41// build_drop_shim analog for async drop glue (for generated coroutine poll function)
42pub(super) fn build_async_drop_shim<'tcx>(
43    tcx: TyCtxt<'tcx>,
44    def_id: DefId,
45    ty: Ty<'tcx>,
46) -> Body<'tcx> {
47    debug!("build_async_drop_shim(def_id={:?}, ty={:?})", def_id, ty);
48    let ty::Coroutine(_, parent_args) = ty.kind() else {
49        bug!();
50    };
51    let typing_env = ty::TypingEnv::fully_monomorphized();
52
53    let drop_ty = parent_args.first().unwrap().expect_ty();
54    let drop_ptr_ty = Ty::new_mut_ptr(tcx, drop_ty);
55
56    assert!(tcx.is_coroutine(def_id));
57    let coroutine_kind = tcx.coroutine_kind(def_id).unwrap();
58
59    assert!(matches!(
60        coroutine_kind,
61        CoroutineKind::Desugared(CoroutineDesugaring::Async, CoroutineSource::Fn)
62    ));
63
64    let needs_async_drop = drop_ty.needs_async_drop(tcx, typing_env);
65    let needs_sync_drop = !needs_async_drop && drop_ty.needs_drop(tcx, typing_env);
66
67    let resume_adt = tcx.adt_def(tcx.require_lang_item(LangItem::ResumeTy, DUMMY_SP));
68    let resume_ty = Ty::new_adt(tcx, resume_adt, ty::List::empty());
69
70    let fn_sig = ty::Binder::dummy(tcx.mk_fn_sig(
71        [ty, resume_ty],
72        tcx.types.unit,
73        false,
74        Safety::Safe,
75        ExternAbi::Rust,
76    ));
77    let sig = tcx.instantiate_bound_regions_with_erased(fn_sig);
78
79    assert!(!drop_ty.is_coroutine());
80    let span = tcx.def_span(def_id);
81    let source_info = SourceInfo::outermost(span);
82
83    // The first argument (index 0), but add 1 for the return value.
84    let coroutine_layout = Place::from(Local::new(1 + 0));
85    let coroutine_layout_dropee =
86        tcx.mk_place_field(coroutine_layout, FieldIdx::new(0), drop_ptr_ty);
87
88    let return_block = BasicBlock::new(1);
89    let mut blocks = IndexVec::with_capacity(2);
90    let block = |blocks: &mut IndexVec<_, _>, kind| {
91        blocks.push(BasicBlockData::new(Some(Terminator { source_info, kind }), false))
92    };
93    block(
94        &mut blocks,
95        if needs_sync_drop {
96            TerminatorKind::Drop {
97                place: tcx.mk_place_deref(coroutine_layout_dropee),
98                target: return_block,
99                unwind: UnwindAction::Continue,
100                replace: false,
101                drop: None,
102                async_fut: None,
103            }
104        } else {
105            TerminatorKind::Goto { target: return_block }
106        },
107    );
108    block(&mut blocks, TerminatorKind::Return);
109
110    let source = MirSource::from_instance(ty::InstanceKind::AsyncDropGlue(def_id, ty));
111    let mut body =
112        new_body(source, blocks, local_decls_for_sig(&sig, span), sig.inputs().len(), span);
113
114    body.coroutine = Some(Box::new(CoroutineInfo::initial(
115        coroutine_kind,
116        parent_args.as_coroutine().yield_ty(),
117        parent_args.as_coroutine().resume_ty(),
118    )));
119    body.phase = MirPhase::Runtime(RuntimePhase::Initial);
120    if !needs_async_drop || drop_ty.references_error() {
121        // Returning noop body for types without `need async drop`
122        // (or sync Drop in case of !`need async drop` && `need drop`).
123        // And also for error types.
124        return body;
125    }
126
127    let mut dropee_ptr = Place::from(body.local_decls.push(LocalDecl::new(drop_ptr_ty, span)));
128    let st_kind = StatementKind::Assign(Box::new((
129        dropee_ptr,
130        Rvalue::Use(Operand::Move(coroutine_layout_dropee)),
131    )));
132    body.basic_blocks_mut()[START_BLOCK].statements.push(Statement::new(source_info, st_kind));
133    dropee_ptr = dropee_emit_retag(tcx, &mut body, dropee_ptr, span);
134
135    let dropline = body.basic_blocks.last_index();
136
137    let patch = {
138        let mut elaborator = DropShimElaborator {
139            body: &body,
140            patch: MirPatch::new(&body),
141            tcx,
142            typing_env,
143            produce_async_drops: true,
144        };
145        let dropee = tcx.mk_place_deref(dropee_ptr);
146        let resume_block = elaborator.patch.resume_block();
147        elaborate_drop(
148            &mut elaborator,
149            source_info,
150            dropee,
151            (),
152            return_block,
153            Unwind::To(resume_block),
154            START_BLOCK,
155            dropline,
156        );
157        elaborator.patch
158    };
159    patch.apply(&mut body);
160
161    body
162}
163
164// * For async drop a "normal" coroutine:
165// `async_drop_in_place<T>::{closure}.poll()` is converted into `T.future_drop_poll()`.
166// Every coroutine has its `poll` (calculate yourself a little further)
167// and its `future_drop_poll` (drop yourself a little further).
168//
169// * For async drop of "async drop coroutine" (`async_drop_in_place<T>::{closure}`):
170// Correct drop of such coroutine means normal execution of nested async drop.
171// async_drop(async_drop(T))::future_drop_poll() => async_drop(T)::poll().
172pub(super) fn build_future_drop_poll_shim<'tcx>(
173    tcx: TyCtxt<'tcx>,
174    def_id: DefId,
175    proxy_ty: Ty<'tcx>,
176    impl_ty: Ty<'tcx>,
177) -> Body<'tcx> {
178    let instance = ty::InstanceKind::FutureDropPollShim(def_id, proxy_ty, impl_ty);
179    let ty::Coroutine(coroutine_def_id, _) = impl_ty.kind() else {
180        bug!("build_future_drop_poll_shim not for coroutine impl type: ({:?})", instance);
181    };
182
183    let span = tcx.def_span(def_id);
184
185    if tcx.is_async_drop_in_place_coroutine(*coroutine_def_id) {
186        build_adrop_for_adrop_shim(tcx, proxy_ty, impl_ty, span, instance)
187    } else {
188        build_adrop_for_coroutine_shim(tcx, proxy_ty, impl_ty, span, instance)
189    }
190}
191
192// For async drop a "normal" coroutine:
193// `async_drop_in_place<T>::{closure}.poll()` is converted into `T.future_drop_poll()`.
194// Every coroutine has its `poll` (calculate yourself a little further)
195// and its `future_drop_poll` (drop yourself a little further).
196fn build_adrop_for_coroutine_shim<'tcx>(
197    tcx: TyCtxt<'tcx>,
198    proxy_ty: Ty<'tcx>,
199    impl_ty: Ty<'tcx>,
200    span: Span,
201    instance: ty::InstanceKind<'tcx>,
202) -> Body<'tcx> {
203    let ty::Coroutine(coroutine_def_id, impl_args) = impl_ty.kind() else {
204        bug!("build_adrop_for_coroutine_shim not for coroutine impl type: ({:?})", instance);
205    };
206    let proxy_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, proxy_ty);
207    // taking _1.0 (impl from Pin)
208    let pin_proxy_layout_local = Local::new(1);
209    let source_info = SourceInfo::outermost(span);
210    // converting `(_1: Pin<&mut CorLayout>, _2: &mut Context<'_>) -> Poll<()>`
211    // into `(_1: Pin<&mut ProxyLayout>, _2: &mut Context<'_>) -> Poll<()>`
212    // let mut _x: &mut CorLayout = &*_1.0.0;
213    // Replace old _1.0 accesses into _x accesses;
214    let body = tcx.optimized_mir(*coroutine_def_id).future_drop_poll().unwrap();
215    let mut body: Body<'tcx> = EarlyBinder::bind(body.clone()).instantiate(tcx, impl_args);
216    body.source.instance = instance;
217    body.phase = MirPhase::Runtime(RuntimePhase::Initial);
218    body.var_debug_info.clear();
219    let pin_adt_ref = tcx.adt_def(tcx.require_lang_item(LangItem::Pin, span));
220    let args = tcx.mk_args(&[proxy_ref.into()]);
221    let pin_proxy_ref = Ty::new_adt(tcx, pin_adt_ref, args);
222
223    let cor_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, impl_ty);
224
225    let proxy_ref_local = body.local_decls.push(LocalDecl::new(proxy_ref, span));
226    let cor_ref_local = body.local_decls.push(LocalDecl::new(cor_ref, span));
227
228    FixProxyFutureDropVisitor { tcx, replace_to: cor_ref_local }.visit_body(&mut body);
229    // Now changing first arg from Pin<&mut ImplCoroutine> to Pin<&mut ProxyCoroutine>
230    body.local_decls[pin_proxy_layout_local] = LocalDecl::new(pin_proxy_ref, span);
231
232    {
233        let mut idx: usize = 0;
234        // _proxy = _1.0 : Pin<&ProxyLayout> ==> &ProxyLayout
235        let proxy_ref_place = Place::from(pin_proxy_layout_local)
236            .project_deeper(&[PlaceElem::Field(FieldIdx::ZERO, proxy_ref)], tcx);
237        body.basic_blocks_mut()[START_BLOCK].statements.insert(
238            idx,
239            Statement::new(
240                source_info,
241                StatementKind::Assign(Box::new((
242                    Place::from(proxy_ref_local),
243                    Rvalue::CopyForDeref(proxy_ref_place),
244                ))),
245            ),
246        );
247        idx += 1;
248        let mut cor_ptr_local = proxy_ref_local;
249        proxy_ty.find_async_drop_impl_coroutine(tcx, |ty| {
250            if ty != proxy_ty {
251                let ty_ptr = Ty::new_mut_ptr(tcx, ty);
252                let impl_ptr_place = Place::from(cor_ptr_local).project_deeper(
253                    &[PlaceElem::Deref, PlaceElem::Field(FieldIdx::ZERO, ty_ptr)],
254                    tcx,
255                );
256                cor_ptr_local = body.local_decls.push(LocalDecl::new(ty_ptr, span));
257                // _cor_ptr = _proxy.0.0 (... .0)
258                body.basic_blocks_mut()[START_BLOCK].statements.insert(
259                    idx,
260                    Statement::new(
261                        source_info,
262                        StatementKind::Assign(Box::new((
263                            Place::from(cor_ptr_local),
264                            Rvalue::CopyForDeref(impl_ptr_place),
265                        ))),
266                    ),
267                );
268                idx += 1;
269            }
270        });
271
272        // _cor_ref = &*cor_ptr
273        let reborrow = Rvalue::Ref(
274            tcx.lifetimes.re_erased,
275            BorrowKind::Mut { kind: MutBorrowKind::Default },
276            tcx.mk_place_deref(Place::from(cor_ptr_local)),
277        );
278        body.basic_blocks_mut()[START_BLOCK].statements.insert(
279            idx,
280            Statement::new(
281                source_info,
282                StatementKind::Assign(Box::new((Place::from(cor_ref_local), reborrow))),
283            ),
284        );
285    }
286    body
287}
288
289// When dropping async drop coroutine, we continue its execution.
290// async_drop(async_drop(T))::future_drop_poll() => async_drop(T)::poll()
291fn build_adrop_for_adrop_shim<'tcx>(
292    tcx: TyCtxt<'tcx>,
293    proxy_ty: Ty<'tcx>,
294    impl_ty: Ty<'tcx>,
295    span: Span,
296    instance: ty::InstanceKind<'tcx>,
297) -> Body<'tcx> {
298    let source_info = SourceInfo::outermost(span);
299    let proxy_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, proxy_ty);
300    // taking _1.0 (impl from Pin)
301    let pin_proxy_layout_local = Local::new(1);
302    let proxy_ref_place = Place::from(pin_proxy_layout_local)
303        .project_deeper(&[PlaceElem::Field(FieldIdx::ZERO, proxy_ref)], tcx);
304    let cor_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, impl_ty);
305
306    // ret_ty = `Poll<()>`
307    let poll_adt_ref = tcx.adt_def(tcx.require_lang_item(LangItem::Poll, span));
308    let ret_ty = Ty::new_adt(tcx, poll_adt_ref, tcx.mk_args(&[tcx.types.unit.into()]));
309    // env_ty = `Pin<&mut proxy_ty>`
310    let pin_adt_ref = tcx.adt_def(tcx.require_lang_item(LangItem::Pin, span));
311    let env_ty = Ty::new_adt(tcx, pin_adt_ref, tcx.mk_args(&[proxy_ref.into()]));
312    // sig = `fn (Pin<&mut proxy_ty>, &mut Context) -> Poll<()>`
313    let sig = tcx.mk_fn_sig(
314        [env_ty, Ty::new_task_context(tcx)],
315        ret_ty,
316        false,
317        hir::Safety::Safe,
318        ExternAbi::Rust,
319    );
320    // This function will be called with pinned proxy coroutine layout.
321    // We need to extract `Arg0.0` to get proxy layout, and then get `.0`
322    // further to receive impl coroutine (may be needed)
323    let mut locals = local_decls_for_sig(&sig, span);
324    let mut blocks = IndexVec::with_capacity(3);
325
326    let proxy_ref_local = locals.push(LocalDecl::new(proxy_ref, span));
327
328    let call_bb = BasicBlock::new(1);
329    let return_bb = BasicBlock::new(2);
330
331    let mut statements = Vec::new();
332
333    statements.push(Statement::new(
334        source_info,
335        StatementKind::Assign(Box::new((
336            Place::from(proxy_ref_local),
337            Rvalue::CopyForDeref(proxy_ref_place),
338        ))),
339    ));
340
341    let mut cor_ptr_local = proxy_ref_local;
342    proxy_ty.find_async_drop_impl_coroutine(tcx, |ty| {
343        if ty != proxy_ty {
344            let ty_ptr = Ty::new_mut_ptr(tcx, ty);
345            let impl_ptr_place = Place::from(cor_ptr_local)
346                .project_deeper(&[PlaceElem::Deref, PlaceElem::Field(FieldIdx::ZERO, ty_ptr)], tcx);
347            cor_ptr_local = locals.push(LocalDecl::new(ty_ptr, span));
348            // _cor_ptr = _proxy.0.0 (... .0)
349            statements.push(Statement::new(
350                source_info,
351                StatementKind::Assign(Box::new((
352                    Place::from(cor_ptr_local),
353                    Rvalue::CopyForDeref(impl_ptr_place),
354                ))),
355            ));
356        }
357    });
358
359    // convert impl coroutine ptr into ref
360    let reborrow = Rvalue::Ref(
361        tcx.lifetimes.re_erased,
362        BorrowKind::Mut { kind: MutBorrowKind::Default },
363        tcx.mk_place_deref(Place::from(cor_ptr_local)),
364    );
365    let cor_ref_place = Place::from(locals.push(LocalDecl::new(cor_ref, span)));
366    statements.push(Statement::new(
367        source_info,
368        StatementKind::Assign(Box::new((cor_ref_place, reborrow))),
369    ));
370
371    // cor_pin_ty = `Pin<&mut cor_ref>`
372    let cor_pin_ty = Ty::new_adt(tcx, pin_adt_ref, tcx.mk_args(&[cor_ref.into()]));
373    let cor_pin_place = Place::from(locals.push(LocalDecl::new(cor_pin_ty, span)));
374
375    let pin_fn = tcx.require_lang_item(LangItem::PinNewUnchecked, span);
376    // call Pin<FutTy>::new_unchecked(&mut impl_cor)
377    blocks.push(BasicBlockData::new_stmts(
378        statements,
379        Some(Terminator {
380            source_info,
381            kind: TerminatorKind::Call {
382                func: Operand::function_handle(tcx, pin_fn, [cor_ref.into()], span),
383                args: [dummy_spanned(Operand::Move(cor_ref_place))].into(),
384                destination: cor_pin_place,
385                target: Some(call_bb),
386                unwind: UnwindAction::Continue,
387                call_source: CallSource::Misc,
388                fn_span: span,
389            },
390        }),
391        false,
392    ));
393    // When dropping async drop coroutine, we continue its execution:
394    // we call impl::poll (impl_layout, ctx)
395    let poll_fn = tcx.require_lang_item(LangItem::FuturePoll, span);
396    let resume_ctx = Place::from(Local::new(2));
397    blocks.push(BasicBlockData::new(
398        Some(Terminator {
399            source_info,
400            kind: TerminatorKind::Call {
401                func: Operand::function_handle(tcx, poll_fn, [impl_ty.into()], span),
402                args: [
403                    dummy_spanned(Operand::Move(cor_pin_place)),
404                    dummy_spanned(Operand::Move(resume_ctx)),
405                ]
406                .into(),
407                destination: Place::return_place(),
408                target: Some(return_bb),
409                unwind: UnwindAction::Continue,
410                call_source: CallSource::Misc,
411                fn_span: span,
412            },
413        }),
414        false,
415    ));
416    blocks.push(BasicBlockData::new(
417        Some(Terminator { source_info, kind: TerminatorKind::Return }),
418        false,
419    ));
420
421    let source = MirSource::from_instance(instance);
422    let mut body = new_body(source, blocks, locals, sig.inputs().len(), span);
423    body.phase = MirPhase::Runtime(RuntimePhase::Initial);
424    return body;
425}