Skip to content

Commit

Permalink
Merge DerefArgVisitor and PinArgVisitor.
Browse files Browse the repository at this point in the history
They are almost identical, differing only in the `ProjectionElem` they
insert. This commit merges them into a new type `SelfArgVisitor`.
  • Loading branch information
nnethercote committed Aug 30, 2024
1 parent 5331280 commit d7cb118
Showing 1 changed file with 7 additions and 47 deletions.
54 changes: 7 additions & 47 deletions compiler/rustc_mir_transform/src/coroutine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -112,11 +112,12 @@ impl<'tcx> MutVisitor<'tcx> for RenameLocalVisitor<'tcx> {
}
}

struct DerefArgVisitor<'tcx> {
struct SelfArgVisitor<'tcx> {
elem: ProjectionElem<Local, Ty<'tcx>>,
tcx: TyCtxt<'tcx>,
}

impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor<'tcx> {
impl<'tcx> MutVisitor<'tcx> for SelfArgVisitor<'tcx> {
fn tcx(&self) -> TyCtxt<'tcx> {
self.tcx
}
Expand All @@ -129,49 +130,7 @@ impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor<'tcx> {
if place.local == SELF_ARG {
replace_base(
place,
Place {
local: SELF_ARG,
projection: self.tcx().mk_place_elems(&[ProjectionElem::Deref]),
},
self.tcx,
);
} else {
self.visit_local(&mut place.local, context, location);

for elem in place.projection.iter() {
if let PlaceElem::Index(local) = elem {
assert_ne!(local, SELF_ARG);
}
}
}
}
}

struct PinArgVisitor<'tcx> {
ref_coroutine_ty: Ty<'tcx>,
tcx: TyCtxt<'tcx>,
}

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

fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
assert_ne!(*local, SELF_ARG);
}

fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
if place.local == SELF_ARG {
replace_base(
place,
Place {
local: SELF_ARG,
projection: self.tcx().mk_place_elems(&[ProjectionElem::Field(
FieldIdx::ZERO,
self.ref_coroutine_ty,
)]),
},
Place { local: SELF_ARG, projection: self.tcx().mk_place_elems(&[self.elem]) },
self.tcx,
);
} else {
Expand Down Expand Up @@ -568,7 +527,7 @@ fn make_coroutine_state_argument_indirect<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Bo
body.local_decls.raw[1].ty = ref_coroutine_ty;

// Add a deref to accesses of the coroutine state
DerefArgVisitor { tcx }.visit_body(body);
SelfArgVisitor { tcx, elem: ProjectionElem::Deref }.visit_body(body);
}

fn make_coroutine_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
Expand All @@ -583,7 +542,8 @@ fn make_coroutine_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body
body.local_decls.raw[1].ty = pin_ref_coroutine_ty;

// Add the Pin field access to accesses of the coroutine state
PinArgVisitor { ref_coroutine_ty, tcx }.visit_body(body);
SelfArgVisitor { tcx, elem: ProjectionElem::Field(FieldIdx::ZERO, ref_coroutine_ty) }
.visit_body(body);
}

/// Allocates a new local and replaces all references of `local` with it. Returns the new local.
Expand Down

0 comments on commit d7cb118

Please sign in to comment.