Rollup merge of #78899 - tmiasko:inline-diverging, r=oli-obk

Support inlining diverging function calls

The existing heuristic does penalize diverging calls to some degree, but since
it never inlined them previously it might need some further modifications.

Additionally introduce storage markers for all temporaries created by
the inliner. The temporary introduced for destination rebrorrow, didn't
use them previously.

Author: jonas-schievink

compiler/rustc_mir/src/transform/inline.rs

@@ -28,7 +28,8 @@ pub struct Inline;
 #[derive(Copy, Clone, Debug)]
 struct CallSite<'tcx> {
     callee: Instance<'tcx>,
-    bb: BasicBlock,
+    block: BasicBlock,
+    target: Option<BasicBlock>,
     source_info: SourceInfo,
 }
 
@@ -172,8 +173,7 @@ impl Inliner<'tcx> {
 
         // Only consider direct calls to functions
         let terminator = bb_data.terminator();
-        // FIXME: Handle inlining of diverging calls
-        if let TerminatorKind::Call { func: ref op, destination: Some(_), .. } = terminator.kind {
+        if let TerminatorKind::Call { func: ref op, ref destination, .. } = terminator.kind {
             if let ty::FnDef(callee_def_id, substs) = *op.ty(caller_body, self.tcx).kind() {
                 // To resolve an instance its substs have to be fully normalized, so
                 // we do this here.
@@ -187,7 +187,12 @@ impl Inliner<'tcx> {
                     return None;
                 }
 
-                return Some(CallSite { callee, bb, source_info: terminator.source_info });
+                return Some(CallSite {
+                    callee,
+                    block: bb,
+                    target: destination.map(|(_, target)| target),
+                    source_info: terminator.source_info,
+                });
             }
         }
 
@@ -399,9 +404,9 @@ impl Inliner<'tcx> {
         caller_body: &mut Body<'tcx>,
         mut callee_body: Body<'tcx>,
     ) {
-        let terminator = caller_body[callsite.bb].terminator.take().unwrap();
+        let terminator = caller_body[callsite.block].terminator.take().unwrap();
         match terminator.kind {
-            TerminatorKind::Call { args, destination: Some(destination), cleanup, .. } => {
+            TerminatorKind::Call { args, destination, cleanup, .. } => {
                 // If the call is something like `a[*i] = f(i)`, where
                 // `i : &mut usize`, then just duplicating the `a[*i]`
                 // Place could result in two different locations if `f`
@@ -418,43 +423,39 @@ impl Inliner<'tcx> {
                     false
                 }
 
-                let dest = if dest_needs_borrow(destination.0) {
-                    trace!("creating temp for return destination");
-                    let dest = Rvalue::Ref(
-                        self.tcx.lifetimes.re_erased,
-                        BorrowKind::Mut { allow_two_phase_borrow: false },
-                        destination.0,
-                    );
-
-                    let ty = dest.ty(caller_body, self.tcx);
-
-                    let temp = LocalDecl::new(ty, callsite.source_info.span);
-
-                    let tmp = caller_body.local_decls.push(temp);
-                    let tmp = Place::from(tmp);
-
-                    let stmt = Statement {
-                        source_info: callsite.source_info,
-                        kind: StatementKind::Assign(box (tmp, dest)),
-                    };
-                    caller_body[callsite.bb].statements.push(stmt);
-                    self.tcx.mk_place_deref(tmp)
+                let dest = if let Some((destination_place, _)) = destination {
+                    if dest_needs_borrow(destination_place) {
+                        trace!("creating temp for return destination");
+                        let dest = Rvalue::Ref(
+                            self.tcx.lifetimes.re_erased,
+                            BorrowKind::Mut { allow_two_phase_borrow: false },
+                            destination_place,
+                        );
+                        let dest_ty = dest.ty(caller_body, self.tcx);
+                        let temp = Place::from(self.new_call_temp(caller_body, &callsite, dest_ty));
+                        caller_body[callsite.block].statements.push(Statement {
+                            source_info: callsite.source_info,
+                            kind: StatementKind::Assign(box (temp, dest)),
+                        });
+                        self.tcx.mk_place_deref(temp)
+                    } else {
+                        destination_place
+                    }
                 } else {
-                    destination.0
+                    trace!("creating temp for return place");
+                    Place::from(self.new_call_temp(caller_body, &callsite, callee_body.return_ty()))
                 };
 
-                let return_block = destination.1;
-
                 // Copy the arguments if needed.
-                let args: Vec<_> = self.make_call_args(args, &callsite, caller_body, return_block);
+                let args: Vec<_> = self.make_call_args(args, &callsite, caller_body);
 
                 let mut integrator = Integrator {
                     args: &args,
                     new_locals: Local::new(caller_body.local_decls.len())..,
                     new_scopes: SourceScope::new(caller_body.source_scopes.len())..,
                     new_blocks: BasicBlock::new(caller_body.basic_blocks().len())..,
                     destination: dest,
-                    return_block,
+                    return_block: callsite.target,
                     cleanup_block: cleanup,
                     in_cleanup_block: false,
                     tcx: self.tcx,
@@ -503,7 +504,7 @@ impl Inliner<'tcx> {
                 caller_body.var_debug_info.extend(callee_body.var_debug_info.drain(..));
                 caller_body.basic_blocks_mut().extend(callee_body.basic_blocks_mut().drain(..));
 
-                caller_body[callsite.bb].terminator = Some(Terminator {
+                caller_body[callsite.block].terminator = Some(Terminator {
                     source_info: callsite.source_info,
                     kind: TerminatorKind::Goto { target: integrator.map_block(START_BLOCK) },
                 });
@@ -527,7 +528,6 @@ impl Inliner<'tcx> {
         args: Vec<Operand<'tcx>>,
         callsite: &CallSite<'tcx>,
         caller_body: &mut Body<'tcx>,
-        return_block: BasicBlock,
     ) -> Vec<Local> {
         let tcx = self.tcx;
 
@@ -558,18 +558,8 @@ impl Inliner<'tcx> {
         // `callee_body.spread_arg == None`, instead of special-casing closures.
         if tcx.is_closure(callsite.callee.def_id()) {
             let mut args = args.into_iter();
-            let self_ = self.create_temp_if_necessary(
-                args.next().unwrap(),
-                callsite,
-                caller_body,
-                return_block,
-            );
-            let tuple = self.create_temp_if_necessary(
-                args.next().unwrap(),
-                callsite,
-                caller_body,
-                return_block,
-            );
+            let self_ = self.create_temp_if_necessary(args.next().unwrap(), callsite, caller_body);
+            let tuple = self.create_temp_if_necessary(args.next().unwrap(), callsite, caller_body);
             assert!(args.next().is_none());
 
             let tuple = Place::from(tuple);
@@ -589,13 +579,13 @@ impl Inliner<'tcx> {
                     Operand::Move(tcx.mk_place_field(tuple, Field::new(i), ty.expect_ty()));
 
                 // Spill to a local to make e.g., `tmp0`.
-                self.create_temp_if_necessary(tuple_field, callsite, caller_body, return_block)
+                self.create_temp_if_necessary(tuple_field, callsite, caller_body)
             });
 
             closure_ref_arg.chain(tuple_tmp_args).collect()
         } else {
             args.into_iter()
-                .map(|a| self.create_temp_if_necessary(a, callsite, caller_body, return_block))
+                .map(|a| self.create_temp_if_necessary(a, callsite, caller_body))
                 .collect()
         }
     }
@@ -607,46 +597,52 @@ impl Inliner<'tcx> {
         arg: Operand<'tcx>,
         callsite: &CallSite<'tcx>,
         caller_body: &mut Body<'tcx>,
-        return_block: BasicBlock,
     ) -> Local {
-        // FIXME: Analysis of the usage of the arguments to avoid
-        // unnecessary temporaries.
-
+        // Reuse the operand if it is a moved temporary.
         if let Operand::Move(place) = &arg {
             if let Some(local) = place.as_local() {
                 if caller_body.local_kind(local) == LocalKind::Temp {
-                    // Reuse the operand if it's a temporary already
                     return local;
                 }
             }
         }
 
+        // Otherwise, create a temporary for the argument.
         trace!("creating temp for argument {:?}", arg);
-        // Otherwise, create a temporary for the arg
-        let arg = Rvalue::Use(arg);
-
-        let ty = arg.ty(caller_body, self.tcx);
-
-        let arg_tmp = LocalDecl::new(ty, callsite.source_info.span);
-        let arg_tmp = caller_body.local_decls.push(arg_tmp);
-
-        caller_body[callsite.bb].statements.push(Statement {
+        let arg_ty = arg.ty(caller_body, self.tcx);
+        let local = self.new_call_temp(caller_body, callsite, arg_ty);
+        caller_body[callsite.block].statements.push(Statement {
             source_info: callsite.source_info,
-            kind: StatementKind::StorageLive(arg_tmp),
+            kind: StatementKind::Assign(box (Place::from(local), Rvalue::Use(arg))),
         });
-        caller_body[callsite.bb].statements.push(Statement {
+        local
+    }
+
+    /// Introduces a new temporary into the caller body that is live for the duration of the call.
+    fn new_call_temp(
+        &self,
+        caller_body: &mut Body<'tcx>,
+        callsite: &CallSite<'tcx>,
+        ty: Ty<'tcx>,
+    ) -> Local {
+        let local = caller_body.local_decls.push(LocalDecl::new(ty, callsite.source_info.span));
+
+        caller_body[callsite.block].statements.push(Statement {
             source_info: callsite.source_info,
-            kind: StatementKind::Assign(box (Place::from(arg_tmp), arg)),
+            kind: StatementKind::StorageLive(local),
         });
-        caller_body[return_block].statements.insert(
-            0,
-            Statement {
-                source_info: callsite.source_info,
-                kind: StatementKind::StorageDead(arg_tmp),
-            },
-        );
-
-        arg_tmp
+
+        if let Some(block) = callsite.target {
+            caller_body[block].statements.insert(
+                0,
+                Statement {
+                    source_info: callsite.source_info,
+                    kind: StatementKind::StorageDead(local),
+                },
+            );
+        }
+
+        local
     }
 }
 
@@ -671,7 +667,7 @@ struct Integrator<'a, 'tcx> {
     new_scopes: RangeFrom<SourceScope>,
     new_blocks: RangeFrom<BasicBlock>,
     destination: Place<'tcx>,
-    return_block: BasicBlock,
+    return_block: Option<BasicBlock>,
     cleanup_block: Option<BasicBlock>,
     in_cleanup_block: bool,
     tcx: TyCtxt<'tcx>,
@@ -817,7 +813,11 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Integrator<'a, 'tcx> {
                 }
             }
             TerminatorKind::Return => {
-                terminator.kind = TerminatorKind::Goto { target: self.return_block };
+                terminator.kind = if let Some(tgt) = self.return_block {
+                    TerminatorKind::Goto { target: tgt }
+                } else {
+                    TerminatorKind::Unreachable
+                }
             }
             TerminatorKind::Resume => {
                 if let Some(tgt) = self.cleanup_block {

src/test/mir-opt/inline/inline-diverging.rs

@@ -0,0 +1,40 @@
+// Tests inlining of diverging calls.
+//
+// ignore-wasm32-bare compiled with panic=abort by default
+#![crate_type = "lib"]
+
+// EMIT_MIR inline_diverging.f.Inline.diff
+pub fn f() {
+    sleep();
+}
+
+// EMIT_MIR inline_diverging.g.Inline.diff
+pub fn g(i: i32) -> u32 {
+    if i > 0 {
+        i as u32
+    } else {
+        panic();
+    }
+}
+
+// EMIT_MIR inline_diverging.h.Inline.diff
+pub fn h() {
+    call_twice(sleep);
+}
+
+#[inline(always)]
+pub fn call_twice<R, F: Fn() -> R>(f: F) -> (R, R) {
+    let a = f();
+    let b = f();
+    (a, b)
+}
+
+#[inline(always)]
+fn panic() -> ! {
+    panic!();
+}
+
+#[inline(always)]
+fn sleep() -> ! {
+    loop {}
+}

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

@@ -0,0 +1,26 @@
+- // MIR for `f` before Inline
++ // MIR for `f` after Inline
+  
+  fn f() -> () {
+      let mut _0: ();                      // return place in scope 0 at $DIR/inline-diverging.rs:7:12: 7:12
+      let mut _1: !;                       // in scope 0 at $DIR/inline-diverging.rs:7:12: 9:2
+      let _2: !;                           // in scope 0 at $DIR/inline-diverging.rs:8:5: 8:12
++     let mut _3: !;                       // in scope 0 at $DIR/inline-diverging.rs:8:5: 8:12
++     scope 1 (inlined sleep) {            // at $DIR/inline-diverging.rs:8:5: 8:12
++     }
+  
+      bb0: {
+          StorageLive(_2);                 // scope 0 at $DIR/inline-diverging.rs:8:5: 8:12
+-         sleep();                         // scope 0 at $DIR/inline-diverging.rs:8:5: 8:12
+-                                          // mir::Constant
+-                                          // + span: $DIR/inline-diverging.rs:8:5: 8:10
+-                                          // + literal: Const { ty: fn() -> ! {sleep}, val: Value(Scalar(<ZST>)) }
++         StorageLive(_3);                 // scope 0 at $DIR/inline-diverging.rs:8:5: 8:12
++         goto -> bb1;                     // scope 0 at $DIR/inline-diverging.rs:8:5: 8:12
++     }
++ 
++     bb1: {
++         goto -> bb1;                     // scope 1 at $DIR/inline-diverging.rs:8:5: 8:12
+      }
+  }
+  

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

@@ -0,0 +1,52 @@
+- // MIR for `g` before Inline
++ // MIR for `g` after Inline
+  
+  fn g(_1: i32) -> u32 {
+      debug i => _1;                       // in scope 0 at $DIR/inline-diverging.rs:12:10: 12:11
+      let mut _0: u32;                     // return place in scope 0 at $DIR/inline-diverging.rs:12:21: 12:24
+      let mut _2: bool;                    // in scope 0 at $DIR/inline-diverging.rs:13:8: 13:13
+      let mut _3: i32;                     // in scope 0 at $DIR/inline-diverging.rs:13:8: 13:9
+      let mut _4: i32;                     // in scope 0 at $DIR/inline-diverging.rs:14:9: 14:10
+      let mut _5: !;                       // in scope 0 at $DIR/inline-diverging.rs:15:12: 17:6
+      let _6: !;                           // in scope 0 at $DIR/inline-diverging.rs:16:9: 16:16
++     let mut _7: !;                       // in scope 0 at $DIR/inline-diverging.rs:16:9: 16:16
++     scope 1 (inlined panic) {            // at $DIR/inline-diverging.rs:16:9: 16:16
++     }
+  
+      bb0: {
+          StorageLive(_2);                 // scope 0 at $DIR/inline-diverging.rs:13:8: 13:13
+          StorageLive(_3);                 // scope 0 at $DIR/inline-diverging.rs:13:8: 13:9
+          _3 = _1;                         // scope 0 at $DIR/inline-diverging.rs:13:8: 13:9
+          _2 = Gt(move _3, const 0_i32);   // scope 0 at $DIR/inline-diverging.rs:13:8: 13:13
+          StorageDead(_3);                 // scope 0 at $DIR/inline-diverging.rs:13:12: 13:13
+          switchInt(_2) -> [false: bb1, otherwise: bb2]; // scope 0 at $DIR/inline-diverging.rs:13:5: 17:6
+      }
+  
+      bb1: {
+          StorageLive(_6);                 // scope 0 at $DIR/inline-diverging.rs:16:9: 16:16
+-         panic();                         // scope 0 at $DIR/inline-diverging.rs:16:9: 16:16
++         StorageLive(_7);                 // scope 0 at $DIR/inline-diverging.rs:16:9: 16:16
++         begin_panic::<&str>(const "explicit panic"); // scope 1 at $DIR/inline-diverging.rs:16:9: 16:16
+                                           // mir::Constant
+-                                          // + span: $DIR/inline-diverging.rs:16:9: 16:14
+-                                          // + literal: Const { ty: fn() -> ! {panic}, val: Value(Scalar(<ZST>)) }
++                                          // + span: $DIR/inline-diverging.rs:16:9: 16:16
++                                          // + literal: Const { ty: fn(&str) -> ! {std::rt::begin_panic::<&str>}, val: Value(Scalar(<ZST>)) }
++                                          // ty::Const
++                                          // + ty: &str
++                                          // + val: Value(Slice { data: Allocation { bytes: [101, 120, 112, 108, 105, 99, 105, 116, 32, 112, 97, 110, 105, 99], relocations: Relocations(SortedMap { data: [] }), init_mask: InitMask { blocks: [16383], len: Size { raw: 14 } }, size: Size { raw: 14 }, align: Align { pow2: 0 }, mutability: Not, extra: () }, start: 0, end: 14 })
++                                          // mir::Constant
++                                          // + span: $DIR/inline-diverging.rs:16:9: 16:16
++                                          // + literal: Const { ty: &str, val: Value(Slice { data: Allocation { bytes: [101, 120, 112, 108, 105, 99, 105, 116, 32, 112, 97, 110, 105, 99], relocations: Relocations(SortedMap { data: [] }), init_mask: InitMask { blocks: [16383], len: Size { raw: 14 } }, size: Size { raw: 14 }, align: Align { pow2: 0 }, mutability: Not, extra: () }, start: 0, end: 14 }) }
+      }
+  
+      bb2: {
+          StorageLive(_4);                 // scope 0 at $DIR/inline-diverging.rs:14:9: 14:10
+          _4 = _1;                         // scope 0 at $DIR/inline-diverging.rs:14:9: 14:10
+          _0 = move _4 as u32 (Misc);      // scope 0 at $DIR/inline-diverging.rs:14:9: 14:17
+          StorageDead(_4);                 // scope 0 at $DIR/inline-diverging.rs:14:16: 14:17
+          StorageDead(_2);                 // scope 0 at $DIR/inline-diverging.rs:18:1: 18:2
+          return;                          // scope 0 at $DIR/inline-diverging.rs:18:2: 18:2
+      }
+  }
+