Skip to content

Commit 33841e6

Browse files
azhoginazhogin
committed
Subpart11 for async drop (major5) - shims codegen
1 parent c57d5af commit 33841e6

File tree

2 files changed

+568
-12
lines changed

2 files changed

+568
-12
lines changed

‎compiler/rustc_mir_transform/src/shim.rs

+248-12
Original file line numberDiff line numberDiff line change
@@ -6,14 +6,15 @@ use rustc_hir::def_id::DefId;
66
use rustc_hir::lang_items::LangItem;
77
use rustc_index::{Idx, IndexVec};
88
use rustc_middle::mir::patch::MirPatch;
9+
use rustc_middle::mir::visit::{MutVisitor, PlaceContext};
910
use rustc_middle::mir::*;
1011
use rustc_middle::query::Providers;
1112
use rustc_middle::ty::{
1213
self, CoroutineArgs, CoroutineArgsExt, EarlyBinder, GenericArgs, Ty, TyCtxt,
1314
};
1415
use rustc_middle::{bug, span_bug};
1516
use rustc_mir_dataflow::elaborate_drops::{self, DropElaborator, DropFlagMode, DropStyle};
16-
use rustc_span::source_map::Spanned;
17+
use rustc_span::source_map::{dummy_spanned, Spanned};
1718
use rustc_span::{Span, DUMMY_SP};
1819
use rustc_target::abi::{FieldIdx, VariantIdx, FIRST_VARIANT};
1920
use rustc_target::spec::abi::Abi;
@@ -23,10 +24,46 @@ use crate::{
2324
instsimplify, mentioned_items, pass_manager as pm, remove_noop_landing_pads, simplify,
2425
};
2526

27+
mod async_destructor_ctor;
28+
2629
pub fn provide(providers: &mut Providers) {
2730
providers.mir_shims = make_shim;
2831
}
2932

33+
// Replace Pin<&mut ImplCoroutine> accesses (_1.0) into Pin<&mut ProxyCoroutine> acceses
34+
struct FixProxyFutureDropVisitor<'tcx> {
35+
tcx: TyCtxt<'tcx>,
36+
replace_to: Local,
37+
}
38+
39+
impl<'tcx> MutVisitor<'tcx> for FixProxyFutureDropVisitor<'tcx> {
40+
fn tcx(&self) -> TyCtxt<'tcx> {
41+
self.tcx
42+
}
43+
44+
fn visit_place(
45+
&mut self,
46+
place: &mut Place<'tcx>,
47+
_context: PlaceContext,
48+
_location: Location,
49+
) {
50+
if place.local == Local::from_u32(1) {
51+
if place.projection.len() == 1 {
52+
assert!(matches!(
53+
place.projection.first(),
54+
Some(ProjectionElem::Field(FieldIdx::ZERO, _))
55+
));
56+
*place = Place::from(self.replace_to);
57+
} else if place.projection.len() == 2 {
58+
assert!(matches!(place.projection[0], ProjectionElem::Field(FieldIdx::ZERO, _)));
59+
assert!(matches!(place.projection[1], ProjectionElem::Deref));
60+
*place =
61+
Place::from(self.replace_to).project_deeper(&[ProjectionElem::Deref], self.tcx);
62+
}
63+
}
64+
}
65+
}
66+
3067
fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceKind<'tcx>) -> Body<'tcx> {
3168
debug!("make_shim({:?})", instance);
3269

@@ -70,7 +107,6 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceKind<'tcx>) -> Body<
70107

71108
build_call_shim(tcx, instance, Some(Adjustment::RefMut), CallKind::Direct(call_mut))
72109
}
73-
74110
ty::InstanceKind::ConstructCoroutineInClosureShim {
75111
coroutine_closure_def_id,
76112
receiver_by_ref,
@@ -81,8 +117,6 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceKind<'tcx>) -> Body<
81117
}
82118

83119
ty::InstanceKind::DropGlue(def_id, ty) => {
84-
// FIXME(#91576): Drop shims for coroutines aren't subject to the MIR passes at the end
85-
// of this function. Is this intentional?
86120
if let Some(ty::Coroutine(coroutine_def_id, args)) = ty.map(Ty::kind) {
87121
let coroutine_body = tcx.optimized_mir(*coroutine_def_id);
88122

@@ -119,23 +153,225 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceKind<'tcx>) -> Body<
119153
],
120154
Some(MirPhase::Runtime(RuntimePhase::Optimized)),
121155
);
122-
123156
return body;
124157
}
125-
126158
build_drop_shim(tcx, def_id, ty)
127159
}
128160
ty::InstanceKind::ThreadLocalShim(..) => build_thread_local_shim(tcx, instance),
129161
ty::InstanceKind::CloneShim(def_id, ty) => build_clone_shim(tcx, def_id, ty),
130162
ty::InstanceKind::FnPtrAddrShim(def_id, ty) => build_fn_ptr_addr_shim(tcx, def_id, ty),
131-
ty::InstanceKind::FutureDropPollShim(_def_id, _proxy_ty, _impl_ty) => {
132-
todo!()
163+
ty::InstanceKind::FutureDropPollShim(def_id, proxy_ty, impl_ty) => {
164+
let ty::Coroutine(coroutine_def_id, impl_args) = impl_ty.kind() else {
165+
bug!("FutureDropPollShim not for coroutine impl type: ({:?})", instance);
166+
};
167+
168+
let span = tcx.def_span(def_id);
169+
let source_info = SourceInfo::outermost(span);
170+
171+
let pin_proxy_layout_local = Local::new(1);
172+
let cor_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, impl_ty);
173+
let proxy_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, proxy_ty);
174+
// taking _1.0.0 (impl from Pin, impl from proxy)
175+
let proxy_ref_place = Place::from(pin_proxy_layout_local)
176+
.project_deeper(&[PlaceElem::Field(FieldIdx::ZERO, proxy_ref)], tcx);
177+
let impl_ref_place = |proxy_ref_local: Local| {
178+
Place::from(proxy_ref_local).project_deeper(
179+
&[
180+
PlaceElem::Deref,
181+
PlaceElem::Downcast(None, VariantIdx::ZERO),
182+
PlaceElem::Field(FieldIdx::ZERO, cor_ref),
183+
],
184+
tcx,
185+
)
186+
};
187+
188+
if tcx.is_templated_coroutine(*coroutine_def_id) {
189+
// ret_ty = `Poll<()>`
190+
let poll_adt_ref = tcx.adt_def(tcx.require_lang_item(LangItem::Poll, None));
191+
let ret_ty = Ty::new_adt(tcx, poll_adt_ref, tcx.mk_args(&[tcx.types.unit.into()]));
192+
// env_ty = `Pin<&mut proxy_ty>`
193+
let pin_adt_ref = tcx.adt_def(tcx.require_lang_item(LangItem::Pin, None));
194+
let env_ty = Ty::new_adt(tcx, pin_adt_ref, tcx.mk_args(&[proxy_ref.into()]));
195+
// sig = `fn (Pin<&mut proxy_ty>, &mut Context) -> Poll<()>`
196+
let sig = tcx.mk_fn_sig(
197+
[env_ty, Ty::new_task_context(tcx)],
198+
ret_ty,
199+
false,
200+
hir::Safety::Safe,
201+
rustc_target::spec::abi::Abi::Rust,
202+
);
203+
let mut locals = local_decls_for_sig(&sig, span);
204+
let mut blocks = IndexVec::with_capacity(3);
205+
206+
let proxy_ref_local = locals.push(LocalDecl::new(proxy_ref, span));
207+
let cor_ref_local = locals.push(LocalDecl::new(cor_ref, span));
208+
let cor_ref_place = Place::from(cor_ref_local);
209+
210+
let call_bb = BasicBlock::new(1);
211+
let return_bb = BasicBlock::new(2);
212+
213+
let assign1 = Statement {
214+
source_info,
215+
kind: StatementKind::Assign(Box::new((
216+
Place::from(proxy_ref_local),
217+
Rvalue::CopyForDeref(proxy_ref_place),
218+
))),
219+
};
220+
let assign2 = Statement {
221+
source_info,
222+
kind: StatementKind::Assign(Box::new((
223+
cor_ref_place,
224+
Rvalue::CopyForDeref(impl_ref_place(proxy_ref_local)),
225+
))),
226+
};
227+
228+
// cor_pin_ty = `Pin<&mut cor_ref>`
229+
let cor_pin_ty = Ty::new_adt(tcx, pin_adt_ref, tcx.mk_args(&[cor_ref.into()]));
230+
let cor_pin_place = Place::from(locals.push(LocalDecl::new(cor_pin_ty, span)));
231+
232+
let pin_fn = tcx.require_lang_item(LangItem::PinNewUnchecked, Some(span));
233+
// call Pin<FutTy>::new_unchecked(&mut impl_cor)
234+
blocks.push(BasicBlockData {
235+
statements: vec![assign1, assign2],
236+
terminator: Some(Terminator {
237+
source_info,
238+
kind: TerminatorKind::Call {
239+
func: Operand::function_handle(tcx, pin_fn, [cor_ref.into()], span),
240+
args: [dummy_spanned(Operand::Move(cor_ref_place))].into(),
241+
destination: cor_pin_place,
242+
target: Some(call_bb),
243+
unwind: UnwindAction::Continue,
244+
call_source: CallSource::Misc,
245+
fn_span: span,
246+
},
247+
}),
248+
is_cleanup: false,
249+
});
250+
// When dropping async drop coroutine, we continue its execution:
251+
// we call impl::poll (impl_layout, ctx)
252+
let poll_fn = tcx.require_lang_item(LangItem::FuturePoll, None);
253+
let resume_ctx = Place::from(Local::new(2));
254+
blocks.push(BasicBlockData {
255+
statements: vec![],
256+
terminator: Some(Terminator {
257+
source_info,
258+
kind: TerminatorKind::Call {
259+
func: Operand::function_handle(tcx, poll_fn, [impl_ty.into()], span),
260+
args: [
261+
dummy_spanned(Operand::Move(cor_pin_place)),
262+
dummy_spanned(Operand::Move(resume_ctx)),
263+
]
264+
.into(),
265+
destination: Place::return_place(),
266+
target: Some(return_bb),
267+
unwind: UnwindAction::Continue,
268+
call_source: CallSource::Misc,
269+
fn_span: span,
270+
},
271+
}),
272+
is_cleanup: false,
273+
});
274+
blocks.push(BasicBlockData {
275+
statements: vec![],
276+
terminator: Some(Terminator { source_info, kind: TerminatorKind::Return }),
277+
is_cleanup: false,
278+
});
279+
280+
let source = MirSource::from_instance(instance);
281+
let mut body = new_body(source, blocks, locals, sig.inputs().len(), span);
282+
pm::run_passes(
283+
tcx,
284+
&mut body,
285+
&[
286+
&mentioned_items::MentionedItems,
287+
&abort_unwinding_calls::AbortUnwindingCalls,
288+
&add_call_guards::CriticalCallEdges,
289+
],
290+
Some(MirPhase::Runtime(RuntimePhase::Optimized)),
291+
);
292+
return body;
293+
}
294+
// future drop poll for async drop must be resolved to standart poll (AsyncDropGlue)
295+
assert!(!tcx.is_templated_coroutine(*coroutine_def_id));
296+
297+
// converting `(_1: Pin<&mut CorLayout>, _2: &mut Context<'_>) -> Poll<()>`
298+
// into `(_1: Pin<&mut ProxyLayout>, _2: &mut Context<'_>) -> Poll<()>`
299+
// let mut _x: &mut CorLayout = &*_1.0.0;
300+
// Replace old _1.0 accesses into _x accesses;
301+
let body = tcx.optimized_mir(*coroutine_def_id).future_drop_poll().unwrap();
302+
let mut body: Body<'tcx> = EarlyBinder::bind(body.clone()).instantiate(tcx, impl_args);
303+
body.source.instance = instance;
304+
body.var_debug_info.clear();
305+
let pin_adt_ref = tcx.adt_def(tcx.require_lang_item(LangItem::Pin, Some(span)));
306+
let args = tcx.mk_args(&[proxy_ref.into()]);
307+
let pin_proxy_ref = Ty::new_adt(tcx, pin_adt_ref, args);
308+
309+
let proxy_ref_local = body.local_decls.push(LocalDecl::new(proxy_ref, span));
310+
let cor_ref_local = body.local_decls.push(LocalDecl::new(cor_ref, span));
311+
FixProxyFutureDropVisitor { tcx, replace_to: cor_ref_local }.visit_body(&mut body);
312+
// Now changing first arg from Pin<&mut ImplCoroutine> to Pin<&mut ProxyCoroutine>
313+
body.local_decls[pin_proxy_layout_local] = LocalDecl::new(pin_proxy_ref, span);
314+
315+
{
316+
let bb: &mut BasicBlockData<'tcx> = &mut body.basic_blocks_mut()[START_BLOCK];
317+
// _tmp = _1.0 : Pin<&ProxyLayout> ==> &ProxyLayout
318+
bb.statements.insert(
319+
0,
320+
Statement {
321+
source_info,
322+
kind: StatementKind::Assign(Box::new((
323+
Place::from(proxy_ref_local),
324+
Rvalue::CopyForDeref(proxy_ref_place),
325+
))),
326+
},
327+
);
328+
bb.statements.insert(
329+
1,
330+
Statement {
331+
source_info,
332+
kind: StatementKind::Assign(Box::new((
333+
Place::from(cor_ref_local),
334+
Rvalue::CopyForDeref(impl_ref_place(proxy_ref_local)),
335+
))),
336+
},
337+
);
338+
}
339+
340+
pm::run_passes(
341+
tcx,
342+
&mut body,
343+
&[
344+
&mentioned_items::MentionedItems,
345+
&abort_unwinding_calls::AbortUnwindingCalls,
346+
&add_call_guards::CriticalCallEdges,
347+
],
348+
Some(MirPhase::Runtime(RuntimePhase::Optimized)),
349+
);
350+
debug!("make_shim({:?}) = {:?}", instance, body);
351+
return body;
133352
}
134-
ty::InstanceKind::AsyncDropGlue(_def_id, _ty) => {
135-
todo!()
353+
ty::InstanceKind::AsyncDropGlue(def_id, ty) => {
354+
let mut body = async_destructor_ctor::build_async_drop_shim(tcx, def_id, ty);
355+
356+
pm::run_passes(
357+
tcx,
358+
&mut body,
359+
&[
360+
&mentioned_items::MentionedItems,
361+
&simplify::SimplifyCfg::MakeShim,
362+
&crate::reveal_all::RevealAll,
363+
&crate::coroutine::StateTransform,
364+
],
365+
Some(MirPhase::Runtime(RuntimePhase::PostCleanup)),
366+
);
367+
debug!("make_shim({:?}) = {:?}", instance, body);
368+
return body;
136369
}
137-
ty::InstanceKind::AsyncDropGlueCtorShim(_def_id, _ty) => {
138-
bug!("AsyncDropGlueCtorShim in re-working ({:?})", instance)
370+
371+
ty::InstanceKind::AsyncDropGlueCtorShim(def_id, ty) => {
372+
let body = async_destructor_ctor::build_async_destructor_ctor_shim(tcx, def_id, ty);
373+
debug!("make_shim({:?}) = {:?}", instance, body);
374+
return body;
139375
}
140376
ty::InstanceKind::Virtual(..) => {
141377
bug!("InstanceKind::Virtual ({:?}) is for direct calls only", instance)

0 commit comments

Comments
 (0)