Simplify accessing locals in a MirPatch

Author: oli-obk

compiler/rustc_middle/src/mir/mod.rs

@@ -75,20 +75,20 @@ pub use self::pretty::{
 pub type LocalDecls<'tcx> = IndexVec<Local, LocalDecl<'tcx>>;
 
 pub trait HasLocalDecls<'tcx> {
-    fn local_decls(&self) -> &LocalDecls<'tcx>;
+    fn local_decl(&self, local: Local) -> &LocalDecl<'tcx>;
 }
 
 impl<'tcx> HasLocalDecls<'tcx> for LocalDecls<'tcx> {
     #[inline]
-    fn local_decls(&self) -> &LocalDecls<'tcx> {
-        self
+    fn local_decl(&self, local: Local) -> &LocalDecl<'tcx> {
+        &self[local]
     }
 }
 
 impl<'tcx> HasLocalDecls<'tcx> for Body<'tcx> {
     #[inline]
-    fn local_decls(&self) -> &LocalDecls<'tcx> {
-        &self.local_decls
+    fn local_decl(&self, local: Local) -> &LocalDecl<'tcx> {
+        &self.local_decls[local]
     }
 }
 

compiler/rustc_middle/src/mir/patch.rs

@@ -17,6 +17,12 @@ pub struct MirPatch<'tcx> {
 }
 
 impl<'tcx> MirPatch<'tcx> {
+    pub fn local_decl<'a>(&'a self, body: &'a Body<'tcx>, local: Local) -> &'a LocalDecl<'tcx> {
+        body.local_decls
+            .get(local)
+            .unwrap_or_else(|| &self.new_locals[local.as_usize() - body.local_decls.len()])
+    }
+
     pub fn new(body: &Body<'tcx>) -> Self {
         let mut result = MirPatch {
             patch_map: IndexVec::from_elem(None, &body.basic_blocks),

compiler/rustc_middle/src/mir/tcx.rs

@@ -127,7 +127,7 @@ impl<'tcx> Place<'tcx> {
     {
         projection
             .iter()
-            .fold(PlaceTy::from_ty(local_decls.local_decls()[local].ty), |place_ty, &elem| {
+            .fold(PlaceTy::from_ty(local_decls.local_decl(local).ty), |place_ty, &elem| {
                 place_ty.projection_ty(tcx, elem)
             })
     }

compiler/rustc_mir_dataflow/src/elaborate_drops.rs

@@ -95,7 +95,7 @@ impl Unwind {
     }
 }
 
-pub trait DropElaborator<'a, 'tcx>: fmt::Debug {
+pub trait DropElaborator<'a, 'tcx>: fmt::Debug + HasLocalDecls<'tcx> {
     /// The type representing paths that can be moved out of.
     ///
     /// Users can move out of individual fields of a struct, such as `a.b.c`. This type is used to
@@ -195,7 +195,7 @@ where
     'tcx: 'b,
 {
     fn place_ty(&self, place: Place<'tcx>) -> Ty<'tcx> {
-        place.ty(self.elaborator.body(), self.tcx()).ty
+        place.ty(&*self.elaborator, self.tcx()).ty
     }
 
     fn tcx(&self) -> TyCtxt<'tcx> {

compiler/rustc_mir_dataflow/src/impls/borrowed_locals.rs

@@ -24,7 +24,7 @@ impl<'tcx> AnalysisDomain<'tcx> for MaybeBorrowedLocals {
 
     fn bottom_value(&self, body: &mir::Body<'tcx>) -> Self::Domain {
         // bottom = unborrowed
-        BitSet::new_empty(body.local_decls().len())
+        BitSet::new_empty(body.local_decls.len())
     }
 
     fn initialize_start_block(&self, _: &mir::Body<'tcx>, _: &mut Self::Domain) {

compiler/rustc_mir_transform/src/const_goto.rs

@@ -84,7 +84,7 @@ impl<'tcx> Visitor<'tcx> for ConstGotoOptimizationFinder<'_, 'tcx> {
                 let target_bb_terminator = target_bb.terminator();
                 let (discr, targets) = target_bb_terminator.kind.as_switch()?;
                 if discr.place() == Some(*place) {
-                    let switch_ty = place.ty(self.body.local_decls(), self.tcx).ty;
+                    let switch_ty = place.ty(self.body, self.tcx).ty;
                     // We now know that the Switch matches on the const place, and it is statementless
                     // Now find which value in the Switch matches the const value.
                     let const_value =

compiler/rustc_mir_transform/src/early_otherwise_branch.rs

@@ -131,7 +131,7 @@ impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch {
                 Operand::Copy(x) => Operand::Copy(*x),
                 Operand::Constant(x) => Operand::Constant(x.clone()),
             };
-            let parent_ty = parent_op.ty(body.local_decls(), tcx);
+            let parent_ty = parent_op.ty(body, tcx);
             let statements_before = bbs[parent].statements.len();
             let parent_end = Location { block: parent, statement_index: statements_before };
 
@@ -268,7 +268,7 @@ fn may_hoist<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, place: Place<'tcx>) ->
             //    `otherwise` branch in `BBC, BBD` in the input to our transformation, which would
             //    have invalidated the data when computing `discriminant(P)`
             // So dereferencing here is correct.
-            ProjectionElem::Deref => match place.ty(body.local_decls(), tcx).ty.kind() {
+            ProjectionElem::Deref => match place.ty(body, tcx).ty.kind() {
                 ty::Ref(..) => {}
                 _ => return false,
             },
@@ -317,7 +317,7 @@ fn evaluate_candidate<'tcx>(
     } = &bbs[parent].terminator().kind else {
         return None
     };
-    let parent_ty = parent_discr.ty(body.local_decls(), tcx);
+    let parent_ty = parent_discr.ty(body, tcx);
     let parent_dest = {
         let poss = targets.otherwise();
         // If the fallthrough on the parent is trivially unreachable, we can let the
@@ -338,7 +338,7 @@ fn evaluate_candidate<'tcx>(
     } = &child_terminator.kind else {
         return None
     };
-    let child_ty = child_discr.ty(body.local_decls(), tcx);
+    let child_ty = child_discr.ty(body, tcx);
     if child_ty != parent_ty {
         return None;
     }

compiler/rustc_mir_transform/src/elaborate_drops.rs

@@ -156,6 +156,12 @@ impl fmt::Debug for Elaborator<'_, '_, '_> {
     }
 }
 
+impl<'tcx> HasLocalDecls<'tcx> for Elaborator<'_, '_, 'tcx> {
+    fn local_decl(&self, local: Local) -> &LocalDecl<'tcx> {
+        self.ctxt.patch.local_decl(self.ctxt.body, local)
+    }
+}
+
 impl<'a, 'tcx> DropElaborator<'a, 'tcx> for Elaborator<'a, '_, 'tcx> {
     type Path = MovePathIndex;
 

compiler/rustc_mir_transform/src/shim.rs

@@ -190,7 +190,7 @@ fn build_drop_shim<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, ty: Option<Ty<'tcx>>)
             BorrowKind::Mut { allow_two_phase_borrow: false },
             tcx.mk_place_deref(dropee_ptr),
         );
-        let ref_ty = reborrow.ty(body.local_decls(), tcx);
+        let ref_ty = reborrow.ty(&body, tcx);
         dropee_ptr = body.local_decls.push(LocalDecl::new(ref_ty, span)).into();
         let new_statements = [
             StatementKind::Assign(Box::new((dropee_ptr, reborrow))),
@@ -270,6 +270,11 @@ impl fmt::Debug for DropShimElaborator<'_, '_> {
         Ok(())
     }
 }
+impl<'tcx> HasLocalDecls<'tcx> for DropShimElaborator<'_, 'tcx> {
+    fn local_decl(&self, local: Local) -> &LocalDecl<'tcx> {
+        self.patch.local_decl(self.body, local)
+    }
+}
 
 impl<'a, 'tcx> DropElaborator<'a, 'tcx> for DropShimElaborator<'a, 'tcx> {
     type Path = ();

compiler/rustc_mir_transform/src/sroa.rs

@@ -32,7 +32,7 @@ impl<'tcx> MirPass<'tcx> for ScalarReplacementOfAggregates {
 fn escaping_locals(body: &Body<'_>) -> BitSet<Local> {
     let mut set = BitSet::new_empty(body.local_decls.len());
     set.insert_range(RETURN_PLACE..=Local::from_usize(body.arg_count));
-    for (local, decl) in body.local_decls().iter_enumerated() {
+    for (local, decl) in body.local_decls.iter_enumerated() {
         if decl.ty.is_union() || decl.ty.is_enum() {
             set.insert(local);
         }