Avoid MIR bloat in inlining

compiler/rustc_mir_transform/src/cost_checker.rs

@@ -31,6 +31,37 @@ impl<'b, 'tcx> CostChecker<'b, 'tcx> {
         CostChecker { tcx, param_env, callee_body, instance, penalty: 0, bonus: 0 }
     }
 
+    /// Add function-level costs not well-represented by the block-level costs.
+    ///
+    /// Needed because the `CostChecker` is used sometimes for just blocks,
+    /// and even the full `Inline` doesn't call `visit_body`, so there's nowhere
+    /// to put this logic in the visitor.
+    pub fn add_function_level_costs(&mut self) {
+        fn is_call_like(bbd: &BasicBlockData<'_>) -> bool {
+            use TerminatorKind::*;
+            match bbd.terminator().kind {
+                Call { .. } | Drop { .. } | Assert { .. } | InlineAsm { .. } => true,
+
+                Goto { .. }
+                | SwitchInt { .. }
+                | UnwindResume
+                | UnwindTerminate(_)
+                | Return
+                | Unreachable => false,
+
+                Yield { .. } | CoroutineDrop | FalseEdge { .. } | FalseUnwind { .. } => {
+                    unreachable!()
+                }
+            }
+        }
+
+        // If the only has one Call (or similar), inlining isn't increasing the total
+        // number of calls, so give extra encouragement to inlining that.
+        if self.callee_body.basic_blocks.iter().filter(|bbd| is_call_like(bbd)).count() == 1 {
+            self.bonus += CALL_PENALTY;
+        }
+    }
+
     pub fn cost(&self) -> usize {
         usize::saturating_sub(self.penalty, self.bonus)
     }

compiler/rustc_mir_transform/src/inline.rs

@@ -85,13 +85,18 @@ fn inline<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool {
     }
 
     let param_env = tcx.param_env_reveal_all_normalized(def_id);
+    let codegen_fn_attrs = tcx.codegen_fn_attrs(def_id);
 
     let mut this = Inliner {
         tcx,
         param_env,
-        codegen_fn_attrs: tcx.codegen_fn_attrs(def_id),
+        codegen_fn_attrs,
         history: Vec::new(),
         changed: false,
+        caller_is_inline_forwarder: matches!(
+            codegen_fn_attrs.inline,
+            InlineAttr::Hint | InlineAttr::Always
+        ) && body_is_forwarder(body),
     };
     let blocks = START_BLOCK..body.basic_blocks.next_index();
     this.process_blocks(body, blocks);
@@ -111,6 +116,9 @@ struct Inliner<'tcx> {
     history: Vec<DefId>,
     /// Indicates that the caller body has been modified.
     changed: bool,
+    /// Indicates that the caller is #[inline] and just calls another function,
+    /// and thus we can inline less into it as it'll be inlined itself.
+    caller_is_inline_forwarder: bool,
 }
 
 impl<'tcx> Inliner<'tcx> {
@@ -485,7 +493,9 @@ impl<'tcx> Inliner<'tcx> {
     ) -> Result<(), &'static str> {
         let tcx = self.tcx;
 
-        let mut threshold = if cross_crate_inlinable {
+        let mut threshold = if self.caller_is_inline_forwarder {
+            self.tcx.sess.opts.unstable_opts.inline_mir_forwarder_threshold.unwrap_or(30)
+        } else if cross_crate_inlinable {
             self.tcx.sess.opts.unstable_opts.inline_mir_hint_threshold.unwrap_or(100)
         } else {
             self.tcx.sess.opts.unstable_opts.inline_mir_threshold.unwrap_or(50)
@@ -504,6 +514,8 @@ impl<'tcx> Inliner<'tcx> {
         let mut checker =
             CostChecker::new(self.tcx, self.param_env, Some(callsite.callee), callee_body);
 
+        checker.add_function_level_costs();
+
         // Traverse the MIR manually so we can account for the effects of inlining on the CFG.
         let mut work_list = vec![START_BLOCK];
         let mut visited = BitSet::new_empty(callee_body.basic_blocks.len());
@@ -1091,3 +1103,37 @@ fn try_instance_mir<'tcx>(
     }
     Ok(tcx.instance_mir(instance))
 }
+
+fn body_is_forwarder(body: &Body<'_>) -> bool {
+    let TerminatorKind::Call { target, .. } = body.basic_blocks[START_BLOCK].terminator().kind
+    else {
+        return false;
+    };
+    if let Some(target) = target {
+        let TerminatorKind::Return = body.basic_blocks[target].terminator().kind else {
+            return false;
+        };
+    }
+
+    let max_blocks = if !body.is_polymorphic {
+        2
+    } else if target.is_none() {
+        3
+    } else {
+        4
+    };
+    if body.basic_blocks.len() > max_blocks {
+        return false;
+    }
+
+    body.basic_blocks.iter_enumerated().all(|(bb, bb_data)| {
+        bb == START_BLOCK
+            || matches!(
+                bb_data.terminator().kind,
+                TerminatorKind::Return
+                    | TerminatorKind::Drop { .. }
+                    | TerminatorKind::UnwindResume
+                    | TerminatorKind::UnwindTerminate(_)
+            )
+    })
+}

compiler/rustc_session/src/options.rs

@@ -1766,6 +1766,8 @@ options! {
         "enable LLVM inlining (default: yes)"),
     inline_mir: Option<bool> = (None, parse_opt_bool, [TRACKED],
         "enable MIR inlining (default: no)"),
+    inline_mir_forwarder_threshold: Option<usize> = (None, parse_opt_number, [TRACKED],
+        "inlining threshold when the caller is a simple forwarding function (default: 30)"),
     inline_mir_hint_threshold: Option<usize> = (None, parse_opt_number, [TRACKED],
         "inlining threshold for functions with inline hint (default: 100)"),
     inline_mir_preserve_debug: Option<bool> = (None, parse_opt_bool, [TRACKED],

library/alloc/src/raw_vec.rs

@@ -429,6 +429,7 @@ impl<T, A: Allocator> RawVec<T, A> {
     ///
     /// Aborts on OOM.
     #[cfg(not(no_global_oom_handling))]
+    #[inline]
     pub fn shrink_to_fit(&mut self, cap: usize) {
         if let Err(err) = self.shrink(cap) {
             handle_error(err);
@@ -511,9 +512,25 @@ impl<T, A: Allocator> RawVec<T, A> {
     }
 
     #[cfg(not(no_global_oom_handling))]
+    #[inline]
     fn shrink(&mut self, cap: usize) -> Result<(), TryReserveError> {
         assert!(cap <= self.capacity(), "Tried to shrink to a larger capacity");
+        // SAFETY: Just checked this isn't trying to grow
+        unsafe { self.shrink_unchecked(cap) }
+    }
 
+    /// `shrink`, but without the capacity check.
+    ///
+    /// This is split out so that `shrink` can inline the check, since it
+    /// optimizes out in things like `shrink_to_fit`, without needing to
+    /// also inline all this code, as doing that ends up failing the
+    /// `vec-shrink-panic` codegen test when `shrink_to_fit` ends up being too
+    /// big for LLVM to be willing to inline.
+    ///
+    /// # Safety
+    /// `cap <= self.capacity()`
+    #[cfg(not(no_global_oom_handling))]
+    unsafe fn shrink_unchecked(&mut self, cap: usize) -> Result<(), TryReserveError> {
         let (ptr, layout) = if let Some(mem) = self.current_memory() { mem } else { return Ok(()) };
         // See current_memory() why this assert is here
         const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };

library/alloc/src/vec/mod.rs

@@ -1101,6 +1101,7 @@ impl<T, A: Allocator> Vec<T, A> {
     /// ```
     #[cfg(not(no_global_oom_handling))]
     #[stable(feature = "rust1", since = "1.0.0")]
+    #[inline]
     pub fn shrink_to_fit(&mut self) {
         // The capacity is never less than the length, and there's nothing to do when
         // they are equal, so we can avoid the panic case in `RawVec::shrink_to_fit`

tests/codegen-units/item-collection/drop_in_place_intrinsic.rs

@@ -1,6 +1,7 @@
 //
 //@ compile-flags:-Zprint-mono-items=eager
 //@ compile-flags:-Zinline-in-all-cgus
+//@ compile-flags:-Zinline-mir=no
 
 #![feature(start)]
 

tests/mir-opt/inline/inline_more_in_non_inline.marked_inline_direct.Inline.after.panic-abort.mir

@@ -0,0 +1,22 @@
+// MIR for `marked_inline_direct` after Inline
+
+fn marked_inline_direct(_1: i32) -> () {
+    debug x => _1;
+    let mut _0: ();
+    let _2: ();
+    let mut _3: i32;
+
+    bb0: {
+        StorageLive(_2);
+        StorageLive(_3);
+        _3 = _1;
+        _2 = call_twice(move _3) -> [return: bb1, unwind unreachable];
+    }
+
+    bb1: {
+        StorageDead(_3);
+        StorageDead(_2);
+        _0 = const ();
+        return;
+    }
+}

tests/mir-opt/inline/inline_more_in_non_inline.marked_inline_direct.Inline.after.panic-unwind.mir

@@ -0,0 +1,22 @@
+// MIR for `marked_inline_direct` after Inline
+
+fn marked_inline_direct(_1: i32) -> () {
+    debug x => _1;
+    let mut _0: ();
+    let _2: ();
+    let mut _3: i32;
+
+    bb0: {
+        StorageLive(_2);
+        StorageLive(_3);
+        _3 = _1;
+        _2 = call_twice(move _3) -> [return: bb1, unwind continue];
+    }
+
+    bb1: {
+        StorageDead(_3);
+        StorageDead(_2);
+        _0 = const ();
+        return;
+    }
+}

tests/mir-opt/inline/inline_more_in_non_inline.marked_inline_indirect.Inline.after.panic-abort.mir

@@ -0,0 +1,27 @@
+// MIR for `marked_inline_indirect` after Inline
+
+fn marked_inline_indirect(_1: i32) -> () {
+    debug x => _1;
+    let mut _0: ();
+    let _2: ();
+    let mut _3: i32;
+    scope 1 (inlined marked_inline_direct) {
+        let _4: ();
+    }
+
+    bb0: {
+        StorageLive(_2);
+        StorageLive(_3);
+        _3 = _1;
+        StorageLive(_4);
+        _4 = call_twice(move _3) -> [return: bb1, unwind unreachable];
+    }
+
+    bb1: {
+        StorageDead(_4);
+        StorageDead(_3);
+        StorageDead(_2);
+        _0 = const ();
+        return;
+    }
+}

tests/mir-opt/inline/inline_more_in_non_inline.marked_inline_indirect.Inline.after.panic-unwind.mir

@@ -0,0 +1,27 @@
+// MIR for `marked_inline_indirect` after Inline
+
+fn marked_inline_indirect(_1: i32) -> () {
+    debug x => _1;
+    let mut _0: ();
+    let _2: ();
+    let mut _3: i32;
+    scope 1 (inlined marked_inline_direct) {
+        let _4: ();
+    }
+
+    bb0: {
+        StorageLive(_2);
+        StorageLive(_3);
+        _3 = _1;
+        StorageLive(_4);
+        _4 = call_twice(move _3) -> [return: bb1, unwind continue];
+    }
+
+    bb1: {
+        StorageDead(_4);
+        StorageDead(_3);
+        StorageDead(_2);
+        _0 = const ();
+        return;
+    }
+}

tests/mir-opt/inline/inline_more_in_non_inline.monomorphic_not_inline.Inline.after.panic-abort.mir

@@ -0,0 +1,34 @@
+// MIR for `monomorphic_not_inline` after Inline
+
+fn monomorphic_not_inline(_1: i32) -> () {
+    debug x => _1;
+    let mut _0: ();
+    let _2: ();
+    let mut _3: i32;
+    scope 1 (inlined call_twice) {
+        let _4: ();
+        let _5: ();
+    }
+
+    bb0: {
+        StorageLive(_2);
+        StorageLive(_3);
+        _3 = _1;
+        StorageLive(_4);
+        StorageLive(_5);
+        _4 = other_thing(_3) -> [return: bb2, unwind unreachable];
+    }
+
+    bb1: {
+        StorageDead(_5);
+        StorageDead(_4);
+        StorageDead(_3);
+        StorageDead(_2);
+        _0 = const ();
+        return;
+    }
+
+    bb2: {
+        _5 = other_thing(move _3) -> [return: bb1, unwind unreachable];
+    }
+}

tests/mir-opt/inline/inline_more_in_non_inline.monomorphic_not_inline.Inline.after.panic-unwind.mir

@@ -0,0 +1,34 @@
+// MIR for `monomorphic_not_inline` after Inline
+
+fn monomorphic_not_inline(_1: i32) -> () {
+    debug x => _1;
+    let mut _0: ();
+    let _2: ();
+    let mut _3: i32;
+    scope 1 (inlined call_twice) {
+        let _4: ();
+        let _5: ();
+    }
+
+    bb0: {
+        StorageLive(_2);
+        StorageLive(_3);
+        _3 = _1;
+        StorageLive(_4);
+        StorageLive(_5);
+        _4 = other_thing(_3) -> [return: bb2, unwind continue];
+    }
+
+    bb1: {
+        StorageDead(_5);
+        StorageDead(_4);
+        StorageDead(_3);
+        StorageDead(_2);
+        _0 = const ();
+        return;
+    }
+
+    bb2: {
+        _5 = other_thing(move _3) -> [return: bb1, unwind continue];
+    }
+}