Limit storage duration of inlined always live locals

1 files changed, 39 insertions, 0 deletions

diff --git a/compiler/rustc_mir/src/transform/inline.rs b/compiler/rustc_mir/src/transform/inline.rs
index aae98f5b6d8..220fc06da86 100644
--- a/compiler/rustc_mir/src/transform/inline.rs
+++ b/compiler/rustc_mir/src/transform/inline.rs
@@ -459,6 +459,7 @@ impl Inliner<'tcx> {
                     tcx: self.tcx,
                     callsite_span: callsite.source_info.span,
                     body_span: callee_body.span,
+                    always_live_locals: BitSet::new_filled(callee_body.local_decls.len()),
                 };
 
                 // Map all `Local`s, `SourceScope`s and `BasicBlock`s to new ones
@@ -490,6 +491,34 @@ impl Inliner<'tcx> {
                     }
                 }
 
+                // If there are any locals without storage markers, give them storage only for the
+                // duration of the call.
+                for local in callee_body.vars_and_temps_iter() {
+                    if integrator.always_live_locals.contains(local) {
+                        let new_local = integrator.map_local(local);
+                        caller_body[callsite.block].statements.push(Statement {
+                            source_info: callsite.source_info,
+                            kind: StatementKind::StorageLive(new_local),
+                        });
+                    }
+                }
+                if let Some(block) = callsite.target {
+                    // To avoid repeated O(n) insert, push any new statements to the end and rotate
+                    // the slice once.
+                    let mut n = 0;
+                    for local in callee_body.vars_and_temps_iter().rev() {
+                        if integrator.always_live_locals.contains(local) {
+                            let new_local = integrator.map_local(local);
+                            caller_body[block].statements.push(Statement {
+                                source_info: callsite.source_info,
+                                kind: StatementKind::StorageDead(new_local),
+                            });
+                            n += 1;
+                        }
+                    }
+                    caller_body[block].statements.rotate_right(n);
+                }
+
                 // Insert all of the (mapped) parts of the callee body into the caller.
                 caller_body.local_decls.extend(
                     // FIXME(eddyb) make `Range<Local>` iterable so that we can use
@@ -670,6 +699,7 @@ struct Integrator<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     callsite_span: Span,
     body_span: Span,
+    always_live_locals: BitSet<Local>,
 }
 
 impl<'a, 'tcx> Integrator<'a, 'tcx> {
@@ -759,6 +789,15 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Integrator<'a, 'tcx> {
         }
     }
 
+    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
+        if let StatementKind::StorageLive(local) | StatementKind::StorageDead(local) =
+            statement.kind
+        {
+            self.always_live_locals.remove(local);
+        }
+        self.super_statement(statement, location);
+    }
+
     fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, loc: Location) {
         // Don't try to modify the implicit `_0` access on return (`return` terminators are
         // replaced down below anyways).