Only rebuild storage for locals that previously had storage statements

Author: Amanieu

compiler/rustc_mir_transform/src/move_elimination.rs

@@ -53,10 +53,10 @@ impl<'tcx> crate::MirPass<'tcx> for MoveElimination {
             &liveness_matrix,
         );
 
-        apply_mappings(tcx, body, &remapped_locals);
+        let storage_to_rebuild = apply_mappings(tcx, body, &remapped_locals);
 
-        if false && tcx.sess.emit_lifetime_markers() {
-            reconstruct_storage(body, &points, &liveness_matrix);
+        if tcx.sess.emit_lifetime_markers() {
+            reconstruct_storage(body, &points, &liveness_matrix, &storage_to_rebuild);
         }
 
         apply_alias_fixup(tcx, body);
@@ -384,27 +384,31 @@ impl<'tcx> Visitor<'tcx> for PlaceUnification<'_, 'tcx> {
 ////////////////////////////////////////////////////////////////////////////////
 // Apply place mappings to the MIR body.
 
+/// Returns the set of locals whose storage needs to be rebuilt.
 fn apply_mappings<'tcx>(
     tcx: TyCtxt<'tcx>,
     body: &mut Body<'tcx>,
     remapped_locals: &IndexVec<Local, Option<Place<'tcx>>>,
-) {
-    let mut mapped_locals = DenseBitSet::new_empty(body.local_decls.len());
+) -> DenseBitSet<Local> {
+    let mut unified_locals = DenseBitSet::new_empty(body.local_decls.len());
     for (src, dst) in remapped_locals.iter_enumerated() {
         if let Some(dst_place) = dst {
-            mapped_locals.insert(src);
-            mapped_locals.insert(dst_place.local);
+            unified_locals.insert(src);
+            unified_locals.insert(dst_place.local);
         }
     }
 
-    let mut rewriter = PlaceUpdater { tcx, remapped_locals, mapped_locals };
+    let storage_to_rebuild = DenseBitSet::new_empty(body.local_decls.len());
+    let mut rewriter = PlaceUpdater { tcx, remapped_locals, unified_locals, storage_to_rebuild };
     rewriter.visit_body_preserves_cfg(body);
+    rewriter.storage_to_rebuild
 }
 
 struct PlaceUpdater<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     remapped_locals: &'a IndexVec<Local, Option<Place<'tcx>>>,
-    mapped_locals: DenseBitSet<Local>,
+    unified_locals: DenseBitSet<Local>,
+    storage_to_rebuild: DenseBitSet<Local>,
 }
 
 impl<'tcx> MutVisitor<'tcx> for PlaceUpdater<'_, 'tcx> {
@@ -434,8 +438,16 @@ impl<'tcx> MutVisitor<'tcx> for PlaceUpdater<'_, 'tcx> {
     fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
         match statement.kind {
             StatementKind::StorageDead(local) | StatementKind::StorageLive(local)
-                if self.mapped_locals.contains(local) =>
+                if self.unified_locals.contains(local) =>
             {
+                // Only rebuild storage for locals that previously had storage
+                // statements.
+                let mapped_local = match self.remapped_locals[local] {
+                    Some(place) => place.local,
+                    None => local,
+                };
+                self.storage_to_rebuild.insert(mapped_local);
+
                 statement.make_nop(true);
                 return;
             }
@@ -481,11 +493,16 @@ fn reconstruct_storage<'tcx>(
     body: &mut Body<'tcx>,
     points: &DenseLocationMap,
     liveness_matrix: &SparseIntervalMatrix<Local, SplitPointIndex>,
+    storage_to_rebuild: &DenseBitSet<Local>,
 ) {
     let mut patcher = MirPatch::new(body);
     let mut split_edges: FxHashMap<(BasicBlock, BasicBlock), BasicBlock> = Default::default();
 
     for local in body.local_decls.indices() {
+        if !storage_to_rebuild.contains(local) {
+            continue;
+        }
+
         // Arguments and return values don't use storage statements.
         match body.local_kind(local) {
             LocalKind::Arg | LocalKind::ReturnPointer => continue,