Skip to content

Simplify how coroutine transform and shims handle nested derefs #156422

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
cjgillot wants to merge 3 commits into from
Closed

Simplify how coroutine transform and shims handle nested derefs #156422

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d1280c6
Call deref_finder when building bodies instead of after it.
cjgillot May 9, 2026
191e766
Fix deref separation for async drop shim.
cjgillot May 10, 2026
2c199f2
Simplify build_adrop_for_coroutine_shim.
cjgillot May 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 6 additions & 11 deletions compiler/rustc_mir_transform/src/coroutine.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1292,6 +1292,9 @@ fn create_coroutine_resume_function<'tcx>(

pm::run_passes_no_validate(tcx, body, &[&abort_unwinding_calls::AbortUnwindingCalls], None);

// Run derefer to fix Derefs that are not in the first place
deref_finder(tcx, body, false);

if let Some(dumper) = MirDumper::new(tcx, "coroutine_resume", body) {
dumper.dump_mir(body);
}
Expand Down Expand Up @@ -1628,30 +1631,22 @@ impl<'tcx> crate::MirPass<'tcx> for StateTransform {
// Create a copy of our MIR and use it to create the drop shim for the coroutine
if has_async_drops {
// If coroutine has async drops, generating async drop shim
let mut drop_shim =
let drop_shim =
create_coroutine_drop_shim_async(tcx, &transform, body, drop_clean, can_unwind);
// Run derefer to fix Derefs that are not in the first place
deref_finder(tcx, &mut drop_shim, false);
body.coroutine.as_mut().unwrap().coroutine_drop_async = Some(drop_shim);
} else {
// If coroutine has no async drops, generating sync drop shim
let mut drop_shim =
let drop_shim =
create_coroutine_drop_shim(tcx, &transform, coroutine_ty, body, drop_clean);
// Run derefer to fix Derefs that are not in the first place
deref_finder(tcx, &mut drop_shim, false);
body.coroutine.as_mut().unwrap().coroutine_drop = Some(drop_shim);

// For coroutine with sync drop, generating async proxy for `future_drop_poll` call
let mut proxy_shim = create_coroutine_drop_shim_proxy_async(tcx, body);
deref_finder(tcx, &mut proxy_shim, false);
let proxy_shim = create_coroutine_drop_shim_proxy_async(tcx, body);
body.coroutine.as_mut().unwrap().coroutine_drop_proxy_async = Some(proxy_shim);
}

// Create the Coroutine::resume / Future::poll function
create_coroutine_resume_function(tcx, transform, body, can_return, can_unwind);

// Run derefer to fix Derefs that are not in the first place
deref_finder(tcx, body, false);
}

fn is_required(&self) -> bool {
Expand Down
9 changes: 9 additions & 0 deletions compiler/rustc_mir_transform/src/coroutine/drop.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -600,6 +600,9 @@ pub(super) fn create_coroutine_drop_shim<'tcx>(
// unrelated code from the resume part of the function
simplify::remove_dead_blocks(&mut body);

// Run derefer to fix Derefs that are not in the first place
deref_finder(tcx, &mut body, false);

// Update the body's def to become the drop glue.
let coroutine_instance = body.source.instance;
let drop_glue = tcx.require_lang_item(LangItem::DropGlue, body.span);
Expand Down Expand Up @@ -703,6 +706,9 @@ pub(super) fn create_coroutine_drop_shim_async<'tcx>(
None,
);

// Run derefer to fix Derefs that are not in the first place
deref_finder(tcx, &mut body, false);

if let Some(dumper) = MirDumper::new(tcx, "coroutine_drop_async", &body) {
dumper.dump_mir(&body);
}
Expand Down Expand Up @@ -750,6 +756,9 @@ pub(super) fn create_coroutine_drop_shim_proxy_async<'tcx>(
};
body.basic_blocks_mut()[call_bb].terminator = Some(Terminator { source_info, kind });

// Run derefer to fix Derefs that are not in the first place
deref_finder(tcx, &mut body, false);

if let Some(dumper) = MirDumper::new(tcx, "coroutine_drop_proxy_async", &body) {
dumper.dump_mir(&body);
}
Expand Down
34 changes: 0 additions & 34 deletions compiler/rustc_mir_transform/src/shim.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,40 +30,6 @@ pub(super) fn provide(providers: &mut Providers) {
providers.mir_shims = make_shim;
}

// Replace Pin<&mut ImplCoroutine> accesses (_1.0) into Pin<&mut ProxyCoroutine> accesses
struct FixProxyFutureDropVisitor<'tcx> {
tcx: TyCtxt<'tcx>,
replace_to: Local,
}

impl<'tcx> MutVisitor<'tcx> for FixProxyFutureDropVisitor<'tcx> {
fn tcx(&self) -> TyCtxt<'tcx> {
self.tcx
}

fn visit_place(
&mut self,
place: &mut Place<'tcx>,
_context: PlaceContext,
_location: Location,
) {
if place.local == Local::from_u32(1) {
if place.projection.len() == 1 {
assert!(matches!(
place.projection.first(),
Some(ProjectionElem::Field(FieldIdx::ZERO, _))
));
*place = Place::from(self.replace_to);
} else if place.projection.len() == 2 {
assert!(matches!(place.projection[0], ProjectionElem::Field(FieldIdx::ZERO, _)));
assert!(matches!(place.projection[1], ProjectionElem::Deref));
*place =
Place::from(self.replace_to).project_deeper(&[ProjectionElem::Deref], self.tcx);
}
}
}
}

fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceKind<'tcx>) -> Body<'tcx> {
debug!("make_shim({:?})", instance);

Expand Down
212 changes: 107 additions & 105 deletions compiler/rustc_mir_transform/src/shim/async_destructor_ctor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,11 @@ use rustc_middle::mir::{
use rustc_middle::ty::{self, EarlyBinder, Ty, TyCtxt, TypeVisitableExt};

use super::*;
use crate::deref_separator::deref_finder;
use crate::patch::MirPatch;

const SELF_ARG: Local = Local::arg(0);

pub(super) fn build_async_destructor_ctor_shim<'tcx>(
tcx: TyCtxt<'tcx>,
def_id: DefId,
Expand Down Expand Up @@ -39,12 +42,12 @@ pub(super) fn build_async_destructor_ctor_shim<'tcx>(
}

// build_drop_shim analog for async drop glue (for generated coroutine poll function)
#[tracing::instrument(level = "trace", skip(tcx), ret)]
pub(super) fn build_async_drop_shim<'tcx>(
tcx: TyCtxt<'tcx>,
def_id: DefId,
ty: Ty<'tcx>,
) -> Body<'tcx> {
debug!("build_async_drop_shim(def_id={:?}, ty={:?})", def_id, ty);
let ty::Coroutine(_, parent_args) = ty.kind() else {
bug!();
};
Expand Down Expand Up @@ -75,7 +78,7 @@ pub(super) fn build_async_drop_shim<'tcx>(
let source_info = SourceInfo::outermost(span);

// The first argument (index 0) which will be local 1 (after the return value).
let coroutine_layout = Place::from(Local::arg(0));
let coroutine_layout = Place::from(SELF_ARG);
let coroutine_layout_dropee =
tcx.mk_place_field(coroutine_layout, FieldIdx::new(0), drop_ptr_ty);

Expand Down Expand Up @@ -111,45 +114,47 @@ pub(super) fn build_async_drop_shim<'tcx>(
parent_args.as_coroutine().resume_ty(),
)));
body.phase = MirPhase::Runtime(RuntimePhase::Initial);
if !needs_async_drop || drop_ty.references_error() {
// Returning noop body for types without `need async drop`
// (or sync Drop in case of !`need async drop` && `need drop`).
// And also for error types.
return body;
}

let dropee_ptr = Place::from(body.local_decls.push(LocalDecl::new(drop_ptr_ty, span)));
let st_kind = StatementKind::Assign(Box::new((
dropee_ptr,
Rvalue::Use(Operand::Move(coroutine_layout_dropee), WithRetag::Yes),
)));
body.basic_blocks_mut()[START_BLOCK].statements.push(Statement::new(source_info, st_kind));

let dropline = body.basic_blocks.last_index();

let patch = {
let mut elaborator = DropShimElaborator {
body: &body,
patch: MirPatch::new(&body),
tcx,
typing_env,
produce_async_drops: true,
// Returning noop body for types without `need async drop`
// (or sync Drop in case of !`need async drop` && `need drop`).
// And also for error types.
if needs_async_drop && !drop_ty.references_error() {
let dropee_ptr = Place::from(body.local_decls.push(LocalDecl::new(drop_ptr_ty, span)));
let st_kind = StatementKind::Assign(Box::new((
dropee_ptr,
Rvalue::Use(Operand::Move(coroutine_layout_dropee), WithRetag::Yes),
)));
body.basic_blocks_mut()[START_BLOCK].statements.push(Statement::new(source_info, st_kind));

let dropline = body.basic_blocks.last_index();

let patch = {
let mut elaborator = DropShimElaborator {
body: &body,
patch: MirPatch::new(&body),
tcx,
typing_env,
produce_async_drops: true,
};
let dropee = tcx.mk_place_deref(dropee_ptr);
let resume_block = elaborator.patch.resume_block();
elaborate_drop(
&mut elaborator,
source_info,
dropee,
(),
return_block,
Unwind::To(resume_block),
START_BLOCK,
dropline,
);
elaborator.patch
};
let dropee = tcx.mk_place_deref(dropee_ptr);
let resume_block = elaborator.patch.resume_block();
elaborate_drop(
&mut elaborator,
source_info,
dropee,
(),
return_block,
Unwind::To(resume_block),
START_BLOCK,
dropline,
);
elaborator.patch
};
patch.apply(&mut body);
patch.apply(&mut body);
}

// We did not bother respectig deref separation, do it here.
deref_finder(tcx, &mut body, false);

body
}
Expand Down Expand Up @@ -196,87 +201,85 @@ fn build_adrop_for_coroutine_shim<'tcx>(
let ty::Coroutine(coroutine_def_id, impl_args) = impl_ty.kind() else {
bug!("build_adrop_for_coroutine_shim not for coroutine impl type: ({:?})", instance);
};
let proxy_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, proxy_ty);
// taking _1.0 (impl from Pin)
let pin_proxy_layout_local = Local::new(1);
let source_info = SourceInfo::outermost(span);
// converting `(_1: Pin<&mut CorLayout>, _2: &mut Context<'_>) -> Poll<()>`
// into `(_1: Pin<&mut ProxyLayout>, _2: &mut Context<'_>) -> Poll<()>`
// let mut _x: &mut CorLayout = &mut *_1.0.0;
// Replace old _1.0 accesses into _x accesses;
let body = tcx.optimized_mir(*coroutine_def_id).future_drop_poll().unwrap();
let mut body: Body<'tcx> =
EarlyBinder::bind(body.clone()).instantiate(tcx, impl_args).skip_norm_wip();
body.source.instance = instance;
body.phase = MirPhase::Runtime(RuntimePhase::Initial);
body.var_debug_info.clear();

// converting `(_1: Pin<&mut CorLayout>, _2: &mut Context<'_>) -> Poll<()>`
// into `(_1: Pin<&mut ProxyLayout>, _2: &mut Context<'_>) -> Poll<()>`
// let mut _x: &mut CorLayout = &mut *_1.0.0;
// Replace old _1.0 accesses into _x accesses;
let proxy_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, proxy_ty);

let pin_adt_ref = tcx.adt_def(tcx.require_lang_item(LangItem::Pin, span));
let args = tcx.mk_args(&[proxy_ref.into()]);
let pin_proxy_ref = Ty::new_adt(tcx, pin_adt_ref, args);

let cor_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, impl_ty);

let proxy_ref_local = body.local_decls.push(LocalDecl::new(proxy_ref, span));
let cor_ref_local = body.local_decls.push(LocalDecl::new(cor_ref, span));

FixProxyFutureDropVisitor { tcx, replace_to: cor_ref_local }.visit_body(&mut body);

// Now changing first arg from Pin<&mut ImplCoroutine> to Pin<&mut ProxyCoroutine>
body.local_decls[pin_proxy_layout_local] = LocalDecl::new(pin_proxy_ref, span);

{
let mut idx: usize = 0;
// _proxy = _1.0 : Pin<&mut ProxyLayout> ==> &mut ProxyLayout
let proxy_ref_place = Place::from(pin_proxy_layout_local)
.project_deeper(&[PlaceElem::Field(FieldIdx::ZERO, proxy_ref)], tcx);
body.basic_blocks_mut()[START_BLOCK].statements.insert(
idx,
Statement::new(
source_info,
StatementKind::Assign(Box::new((
Place::from(proxy_ref_local),
Rvalue::Use(Operand::Copy(proxy_ref_place), WithRetag::Yes),
))),
),
);
idx += 1;

// _cor_ref_tmp = (*(*_proxy).0).0...
let mut cor_ref_tmp_local = proxy_ref_local;
proxy_ty.find_async_drop_impl_coroutine(tcx, |ty| {
if ty != proxy_ty {
let ty_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, ty);
let impl_ptr_place = Place::from(cor_ref_tmp_local).project_deeper(
&[PlaceElem::Deref, PlaceElem::Field(FieldIdx::ZERO, ty_ref)],
tcx,
);
cor_ref_tmp_local = body.local_decls.push(LocalDecl::new(ty_ref, span));
body.basic_blocks_mut()[START_BLOCK].statements.insert(
idx,
Statement::new(
source_info,
StatementKind::Assign(Box::new((
Place::from(cor_ref_tmp_local),
Rvalue::Use(Operand::Copy(impl_ptr_place), WithRetag::Yes),
))),
),
);
idx += 1;
}
});
body.local_decls[SELF_ARG] = LocalDecl::new(pin_proxy_ref, span);

// _cor_ref = cor_ref_tmp
body.basic_blocks_mut()[START_BLOCK].statements.insert(
idx,
Statement::new(
source_info,
StatementKind::Assign(Box::new((
Place::from(cor_ref_local),
Rvalue::Use(Operand::Move(Place::from(cor_ref_tmp_local)), WithRetag::Yes),
))),
),
);
// Build the projection to assign `cor_ref_local = _1.<projection>`.
let mut pin_proxy_to_cor_projection = vec![
// _1.0 : Pin<&mut ProxyLayout> ==> &mut ProxyLayout
PlaceElem::Field(FieldIdx::ZERO, proxy_ref),
];

// _cor_ref_tmp = (*(*_proxy).0).0...
proxy_ty.find_async_drop_impl_coroutine(tcx, |ty| {
if ty != proxy_ty {
let ty_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, ty);
pin_proxy_to_cor_projection.push(PlaceElem::Deref);
pin_proxy_to_cor_projection.push(PlaceElem::Field(FieldIdx::ZERO, ty_ref));
}
});

// _cor_ref = cor_ref_tmp
let projected_pin = Place::from(SELF_ARG).project_deeper(&pin_proxy_to_cor_projection, tcx);
body.basic_blocks_mut()[START_BLOCK].statements.insert(
0,
Statement::new(
source_info,
StatementKind::Assign(Box::new((
Place::from(cor_ref_local),
Rvalue::Use(Operand::Move(projected_pin), WithRetag::Yes),
))),
),
);

// We did not bother respectig deref separation, do it here.
deref_finder(tcx, &mut body, false);

return body;

/// Replace Pin<&mut ImplCoroutine> accesses (_1.0) into Pin<&mut ProxyCoroutine> accesses
struct FixProxyFutureDropVisitor<'tcx> {
tcx: TyCtxt<'tcx>,
replace_to: Local,
}

impl<'tcx> MutVisitor<'tcx> for FixProxyFutureDropVisitor<'tcx> {
fn tcx(&self) -> TyCtxt<'tcx> {
self.tcx
}

fn visit_place(&mut self, place: &mut Place<'tcx>, _: PlaceContext, _: Location) {
if place.local == SELF_ARG
&& let Some((first, rest)) = place.projection.split_first()
{
assert!(matches!(first, ProjectionElem::Field(FieldIdx::ZERO, _)));
*place = Place::from(self.replace_to).project_deeper(rest, self.tcx);
}
}
}
body
}

// When dropping async drop coroutine, we continue its execution.
Expand All @@ -291,9 +294,8 @@ fn build_adrop_for_adrop_shim<'tcx>(
let source_info = SourceInfo::outermost(span);
let proxy_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, proxy_ty);
// taking _1.0 (impl from Pin)
let pin_proxy_layout_local = Local::new(1);
let proxy_ref_place = Place::from(pin_proxy_layout_local)
.project_deeper(&[PlaceElem::Field(FieldIdx::ZERO, proxy_ref)], tcx);
let proxy_ref_place =
Place::from(SELF_ARG).project_deeper(&[PlaceElem::Field(FieldIdx::ZERO, proxy_ref)], tcx);
let cor_ref = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, impl_ty);

// ret_ty = `Poll<()>`
Expand Down
Loading