Limit storage duration of inlined always live locals

Author: tmiasko

compiler/rustc_middle/src/mir/mod.rs

@@ -420,7 +420,9 @@ impl<'tcx> Body<'tcx> {
     /// Returns an iterator over all user-defined variables and compiler-generated temporaries (all
     /// locals that are neither arguments nor the return place).
     #[inline]
-    pub fn vars_and_temps_iter(&self) -> impl Iterator<Item = Local> + ExactSizeIterator {
+    pub fn vars_and_temps_iter(
+        &self,
+    ) -> impl DoubleEndedIterator<Item = Local> + ExactSizeIterator {
         let arg_count = self.arg_count;
         let local_count = self.local_decls.len();
         (arg_count + 1..local_count).map(Local::new)

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).

src/test/mir-opt/inline/inline_diverging.h.Inline.diff

@@ -40,6 +40,7 @@
 -                                          // mir::Constant
                                            // + span: $DIR/inline-diverging.rs:22:16: 22:21
                                            // + literal: Const { ty: fn() -> ! {sleep}, val: Value(Scalar(<ZST>)) }
++         StorageLive(_6);                 // scope 0 at $DIR/inline-diverging.rs:22:5: 22:22
 +         StorageLive(_3);                 // scope 1 at $DIR/inline-diverging.rs:22:5: 22:22
 +         StorageLive(_4);                 // scope 1 at $DIR/inline-diverging.rs:22:5: 22:22
 +         _4 = &_2;                        // scope 1 at $DIR/inline-diverging.rs:22:5: 22:22

src/test/mir-opt/inline/inline_generator.main.Inline.diff

@@ -65,12 +65,16 @@
 -                                          // + literal: Const { ty: for<'r> fn(std::pin::Pin<&'r mut impl std::ops::Generator<bool>>, bool) -> std::ops::GeneratorState<<impl std::ops::Generator<bool> as std::ops::Generator<bool>>::Yield, <impl std::ops::Generator<bool> as std::ops::Generator<bool>>::Return> {<impl std::ops::Generator<bool> as std::ops::Generator<bool>>::resume}, val: Value(Scalar(<ZST>)) }
 +         StorageLive(_7);                 // scope 0 at $DIR/inline-generator.rs:9:14: 9:46
 +         _7 = const false;                // scope 0 at $DIR/inline-generator.rs:9:14: 9:46
++         StorageLive(_8);                 // scope 0 at $DIR/inline-generator.rs:9:14: 9:46
++         StorageLive(_9);                 // scope 0 at $DIR/inline-generator.rs:9:14: 9:46
 +         _9 = discriminant((*(_2.0: &mut [generator@$DIR/inline-generator.rs:15:5: 15:41 {bool, i32}]))); // scope 6 at $DIR/inline-generator.rs:9:14: 9:46
 +         switchInt(move _9) -> [0_u32: bb3, 1_u32: bb8, 3_u32: bb7, otherwise: bb9]; // scope 6 at $DIR/inline-generator.rs:9:14: 9:46
       }
 
 -     bb3: {
 +     bb1: {
++         StorageDead(_9);                 // scope 0 at $DIR/inline-generator.rs:9:14: 9:46
++         StorageDead(_8);                 // scope 0 at $DIR/inline-generator.rs:9:14: 9:46
 +         StorageDead(_7);                 // scope 0 at $DIR/inline-generator.rs:9:14: 9:46
           StorageDead(_2);                 // scope 0 at $DIR/inline-generator.rs:9:45: 9:46
           StorageDead(_4);                 // scope 0 at $DIR/inline-generator.rs:9:46: 9:47

src/test/mir-opt/inline/inline_shims.drop.Inline.diff

@@ -36,11 +36,15 @@
 -                                          // mir::Constant
 -                                          // + span: $DIR/inline-shims.rs:12:14: 12:37
 -                                          // + literal: Const { ty: unsafe fn(*mut std::option::Option<B>) {std::intrinsics::drop_in_place::<std::option::Option<B>>}, val: Value(Scalar(<ZST>)) }
++         StorageLive(_6);                 // scope 2 at $DIR/inline-shims.rs:12:14: 12:40
++         StorageLive(_7);                 // scope 2 at $DIR/inline-shims.rs:12:14: 12:40
 +         _6 = discriminant((*_5));        // scope 3 at $DIR/inline-shims.rs:12:14: 12:40
 +         switchInt(move _6) -> [0_isize: bb2, otherwise: bb3]; // scope 3 at $DIR/inline-shims.rs:12:14: 12:40
       }
 
       bb2: {
++         StorageDead(_7);                 // scope 2 at $DIR/inline-shims.rs:12:14: 12:40
++         StorageDead(_6);                 // scope 2 at $DIR/inline-shims.rs:12:14: 12:40
           StorageDead(_5);                 // scope 2 at $DIR/inline-shims.rs:12:39: 12:40
           return;                          // scope 0 at $DIR/inline-shims.rs:13:2: 13:2
 +     }

src/test/ui/mir/auxiliary/issue_76375_aux.rs

@@ -1,8 +1,8 @@
 // edition:2018
-// compile-flags: -Z mir-opt-level=2 -Z unsound-mir-opts
+// compile-flags: -Z mir-opt-level=2
 
 #[inline(always)]
-pub fn f(s: bool) -> String {
+pub fn copy_prop(s: bool) -> String {
     let a = "Hello world!".to_string();
     let b = a;
     let c = b;
@@ -12,3 +12,9 @@ pub fn f(s: bool) -> String {
         String::new()
     }
 }
+
+#[inline(always)]
+pub fn dest_prop(x: &[u8]) -> &[u8] {
+    let y = &x[..x.len()];
+    y
+}

src/test/ui/mir/issue-76375.rs

@@ -1,15 +1,27 @@
+// Regression test for issue #76375.
+//
 // edition:2018
 // build-pass
-// compile-flags: -Z mir-opt-level=2 -L.
+// compile-flags: -Z mir-opt-level=2
 // aux-build:issue_76375_aux.rs
 
 #![crate_type = "lib"]
 
 extern crate issue_76375_aux;
 
 pub async fn g() {
-    issue_76375_aux::f(true);
+    issue_76375_aux::copy_prop(true);
     h().await;
 }
 
+pub async fn u() {
+    let b = [0u8; 32];
+    let mut i = 0;
+    while i != 10 {
+        issue_76375_aux::dest_prop(&b);
+        h().await;
+        i += 1;
+    }
+}
+
 pub async fn h() {}