From d7cb1181dc29210af3c49e7be267b3423095fd3f Mon Sep 17 00:00:00 2001
From: Nicholas Nethercote <n.nethercote@gmail.com>
Date: Wed, 28 Aug 2024 13:08:02 +1000
Subject: [PATCH] Merge `DerefArgVisitor` and `PinArgVisitor`.

They are almost identical, differing only in the `ProjectionElem` they
insert. This commit merges them into a new type `SelfArgVisitor`.
---
 compiler/rustc_mir_transform/src/coroutine.rs | 54 +++----------------
 1 file changed, 7 insertions(+), 47 deletions(-)

diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs
index 7dbad721d6d..99c6eee40e9 100644
--- a/compiler/rustc_mir_transform/src/coroutine.rs
+++ b/compiler/rustc_mir_transform/src/coroutine.rs
@@ -112,11 +112,12 @@ fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Loca
     }
 }
 
-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
     }
@@ -129,49 +130,7 @@ fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, locati
         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 {
@@ -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>) {
@@ -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.