Auto merge of rust-lang#138157 - scottmcm:inline-more-tiny-things, r=<try>

Allow more top-down inlining for single-BB callees

This means that things like `<usize as Step>::forward_unchecked` and `<PartialOrd for f32>::le` will inline even if
we've already done a bunch of inlining to find the calls to them.

Fixes rust-lang#138136

Draft as it's built atop rust-lang#138135, which adds a mir-opt test that's a nice demonstration of this.

Author: bors

compiler/rustc_mir_transform/src/inline.rs

@@ -26,6 +26,7 @@ use crate::{check_inline, util};
 
 pub(crate) mod cycle;
 
+const HISTORY_DEPTH_LIMIT: usize = 20;
 const TOP_DOWN_DEPTH_LIMIT: usize = 5;
 
 #[derive(Clone, Debug)]
@@ -128,10 +129,6 @@ trait Inliner<'tcx> {
         callee_attrs: &CodegenFnAttrs,
     ) -> Result<(), &'static str>;
 
-    // How many callsites in a body are we allowed to inline? We need to limit this in order
-    // to prevent super-linear growth in MIR size.
-    fn inline_limit_for_block(&self) -> Option<usize>;
-
     /// Called when inlining succeeds.
     fn on_inline_success(
         &mut self,
@@ -142,9 +139,6 @@ trait Inliner<'tcx> {
 
     /// Called when inlining failed or was not performed.
     fn on_inline_failure(&self, callsite: &CallSite<'tcx>, reason: &'static str);
-
-    /// Called when the inline limit for a body is reached.
-    fn on_inline_limit_reached(&self) -> bool;
 }
 
 struct ForceInliner<'tcx> {
@@ -224,10 +218,6 @@ impl<'tcx> Inliner<'tcx> for ForceInliner<'tcx> {
         }
     }
 
-    fn inline_limit_for_block(&self) -> Option<usize> {
-        Some(usize::MAX)
-    }
-
     fn on_inline_success(
         &mut self,
         callsite: &CallSite<'tcx>,
@@ -261,10 +251,6 @@ impl<'tcx> Inliner<'tcx> for ForceInliner<'tcx> {
             justification: justification.map(|sym| crate::errors::ForceInlineJustification { sym }),
         });
     }
-
-    fn on_inline_limit_reached(&self) -> bool {
-        false
-    }
 }
 
 struct NormalInliner<'tcx> {
@@ -278,6 +264,10 @@ struct NormalInliner<'tcx> {
     /// The number of `DefId`s is finite, so checking history is enough
     /// to ensure that we do not loop endlessly while inlining.
     history: Vec<DefId>,
+    /// How many (multi-call) callsites have we inlined for the top-level call?
+    ///
+    /// We need to limit this in order to prevent super-linear growth in MIR size.
+    top_down_counter: usize,
     /// Indicates that the caller body has been modified.
     changed: bool,
     /// Indicates that the caller is #[inline] and just calls another function,
@@ -295,6 +285,7 @@ impl<'tcx> Inliner<'tcx> for NormalInliner<'tcx> {
             typing_env,
             def_id,
             history: Vec::new(),
+            top_down_counter: 0,
             changed: false,
             caller_is_inline_forwarder: matches!(
                 codegen_fn_attrs.inline,
@@ -351,6 +342,15 @@ impl<'tcx> Inliner<'tcx> for NormalInliner<'tcx> {
             return Err("body has errors");
         }
 
+        if callee_body.basic_blocks.len() > 1 {
+            if self.history.len() > HISTORY_DEPTH_LIMIT {
+                return Err("Too many blocks for deep history");
+            }
+            if self.top_down_counter > TOP_DOWN_DEPTH_LIMIT {
+                return Err("Too many blocks for more top-down inlining");
+            }
+        }
+
         let mut threshold = if self.caller_is_inline_forwarder {
             tcx.sess.opts.unstable_opts.inline_mir_forwarder_threshold.unwrap_or(30)
         } else if tcx.cross_crate_inlinable(callsite.callee.def_id()) {
@@ -431,14 +431,6 @@ impl<'tcx> Inliner<'tcx> for NormalInliner<'tcx> {
         }
     }
 
-    fn inline_limit_for_block(&self) -> Option<usize> {
-        match self.history.len() {
-            0 => Some(usize::MAX),
-            1..=TOP_DOWN_DEPTH_LIMIT => Some(1),
-            _ => None,
-        }
-    }
-
     fn on_inline_success(
         &mut self,
         callsite: &CallSite<'tcx>,
@@ -447,13 +439,26 @@ impl<'tcx> Inliner<'tcx> for NormalInliner<'tcx> {
     ) {
         self.changed = true;
 
+        let new_calls_count = new_blocks
+            .clone()
+            .filter(|&bb| {
+                matches!(
+                    caller_body.basic_blocks[bb].terminator().kind,
+                    TerminatorKind::Call { .. },
+                )
+            })
+            .count();
+        if new_calls_count > 1 {
+            self.top_down_counter += 1;
+        }
+
         self.history.push(callsite.callee.def_id());
         process_blocks(self, caller_body, new_blocks);
         self.history.pop();
-    }
 
-    fn on_inline_limit_reached(&self) -> bool {
-        true
+        if self.history.is_empty() {
+            self.top_down_counter = 0;
+        }
     }
 
     fn on_inline_failure(&self, _: &CallSite<'tcx>, _: &'static str) {}
@@ -482,8 +487,6 @@ fn process_blocks<'tcx, I: Inliner<'tcx>>(
     caller_body: &mut Body<'tcx>,
     blocks: Range<BasicBlock>,
 ) {
-    let Some(inline_limit) = inliner.inline_limit_for_block() else { return };
-    let mut inlined_count = 0;
     for bb in blocks {
         let bb_data = &caller_body[bb];
         if bb_data.is_cleanup {
@@ -505,13 +508,6 @@ fn process_blocks<'tcx, I: Inliner<'tcx>>(
             Ok(new_blocks) => {
                 debug!("inlined {}", callsite.callee);
                 inliner.on_inline_success(&callsite, caller_body, new_blocks);
-
-                inlined_count += 1;
-                if inlined_count == inline_limit {
-                    if inliner.on_inline_limit_reached() {
-                        return;
-                    }
-                }
             }
         }
     }

library/core/src/cmp.rs

@@ -29,6 +29,7 @@ mod bytewise;
 pub(crate) use bytewise::BytewiseEq;
 
 use self::Ordering::*;
+use crate::ops::ControlFlow::{self, Break, Continue};
 
 /// Trait for comparisons using the equality operator.
 ///
@@ -1446,6 +1447,54 @@ pub macro PartialOrd($item:item) {
     /* compiler built-in */
 }
 
+/// Helpers for chaining together field PartialOrds into the full type's ordering.
+///
+/// If the two values are equal, returns `ControlFlow::Continue`.
+/// If the two values are not equal, returns `ControlFlow::Break(self OP other)`.
+///
+/// This allows simple types like `i32` and `f64` to just emit two comparisons
+/// directly, instead of needing to optimize the 3-way comparison.
+///
+/// Currently this is done using specialization, but it doesn't need that:
+/// it could be provided methods on `PartialOrd` instead and work fine.
+pub(crate) trait SpecChainingPartialOrd<Rhs>: PartialOrd<Rhs> {
+    fn spec_chain_lt(&self, other: &Rhs) -> ControlFlow<bool>;
+    fn spec_chain_le(&self, other: &Rhs) -> ControlFlow<bool>;
+    fn spec_chain_gt(&self, other: &Rhs) -> ControlFlow<bool>;
+    fn spec_chain_ge(&self, other: &Rhs) -> ControlFlow<bool>;
+}
+
+impl<T: PartialOrd<U>, U> SpecChainingPartialOrd<U> for T {
+    #[inline]
+    default fn spec_chain_lt(&self, other: &U) -> ControlFlow<bool> {
+        match PartialOrd::partial_cmp(self, other) {
+            Some(Equal) => Continue(()),
+            c => Break(c.is_some_and(Ordering::is_lt)),
+        }
+    }
+    #[inline]
+    default fn spec_chain_le(&self, other: &U) -> ControlFlow<bool> {
+        match PartialOrd::partial_cmp(self, other) {
+            Some(Equal) => Continue(()),
+            c => Break(c.is_some_and(Ordering::is_le)),
+        }
+    }
+    #[inline]
+    default fn spec_chain_gt(&self, other: &U) -> ControlFlow<bool> {
+        match PartialOrd::partial_cmp(self, other) {
+            Some(Equal) => Continue(()),
+            c => Break(c.is_some_and(Ordering::is_gt)),
+        }
+    }
+    #[inline]
+    default fn spec_chain_ge(&self, other: &U) -> ControlFlow<bool> {
+        match PartialOrd::partial_cmp(self, other) {
+            Some(Equal) => Continue(()),
+            c => Break(c.is_some_and(Ordering::is_ge)),
+        }
+    }
+}
+
 /// Compares and returns the minimum of two values.
 ///
 /// Returns the first argument if the comparison determines them to be equal.
@@ -1741,6 +1790,7 @@ where
 mod impls {
     use crate::cmp::Ordering::{self, Equal, Greater, Less};
     use crate::hint::unreachable_unchecked;
+    use crate::ops::ControlFlow::{self, Break, Continue};
 
     macro_rules! partial_eq_impl {
         ($($t:ty)*) => ($(
@@ -1779,6 +1829,36 @@ mod impls {
 
     eq_impl! { () bool char usize u8 u16 u32 u64 u128 isize i8 i16 i32 i64 i128 }
 
+    macro_rules! chaining_impl {
+        ($t:ty) => {
+            // These implementations are the same for `Ord` or `PartialOrd` types
+            // because if either is NAN the `==` test will fail so we end up in
+            // the `Break` case and the comparison will correctly return `false`.
+            impl super::SpecChainingPartialOrd<$t> for $t {
+                #[inline]
+                fn spec_chain_lt(&self, other: &Self) -> ControlFlow<bool> {
+                    let (lhs, rhs) = (*self, *other);
+                    if lhs == rhs { Continue(()) } else { Break(lhs < rhs) }
+                }
+                #[inline]
+                fn spec_chain_le(&self, other: &Self) -> ControlFlow<bool> {
+                    let (lhs, rhs) = (*self, *other);
+                    if lhs == rhs { Continue(()) } else { Break(lhs <= rhs) }
+                }
+                #[inline]
+                fn spec_chain_gt(&self, other: &Self) -> ControlFlow<bool> {
+                    let (lhs, rhs) = (*self, *other);
+                    if lhs == rhs { Continue(()) } else { Break(lhs > rhs) }
+                }
+                #[inline]
+                fn spec_chain_ge(&self, other: &Self) -> ControlFlow<bool> {
+                    let (lhs, rhs) = (*self, *other);
+                    if lhs == rhs { Continue(()) } else { Break(lhs >= rhs) }
+                }
+            }
+        };
+    }
+
     macro_rules! partial_ord_impl {
         ($($t:ty)*) => ($(
             #[stable(feature = "rust1", since = "1.0.0")]
@@ -1801,6 +1881,8 @@ mod impls {
                 #[inline(always)]
                 fn gt(&self, other: &$t) -> bool { (*self) > (*other) }
             }
+
+            chaining_impl!($t);
         )*)
     }
 
@@ -1840,6 +1922,8 @@ mod impls {
                 fn gt(&self, other: &$t) -> bool { (*self) > (*other) }
             }
 
+            chaining_impl!($t);
+
             #[stable(feature = "rust1", since = "1.0.0")]
             impl Ord for $t {
                 #[inline]

library/core/src/tuple.rs

@@ -1,7 +1,9 @@
 // See core/src/primitive_docs.rs for documentation.
 
 use crate::cmp::Ordering::{self, *};
+use crate::cmp::SpecChainingPartialOrd;
 use crate::marker::{ConstParamTy_, StructuralPartialEq, UnsizedConstParamTy};
+use crate::ops::ControlFlow::{Break, Continue};
 
 // Recursive macro for implementing n-ary tuple functions and operations
 //
@@ -80,19 +82,19 @@ macro_rules! tuple_impls {
                 }
                 #[inline]
                 fn lt(&self, other: &($($T,)+)) -> bool {
-                    lexical_ord!(lt, Less, $( ${ignore($T)} self.${index()}, other.${index()} ),+)
+                    lexical_ord!(lt, spec_chain_lt, $( ${ignore($T)} self.${index()}, other.${index()} ),+)
                 }
                 #[inline]
                 fn le(&self, other: &($($T,)+)) -> bool {
-                    lexical_ord!(le, Less, $( ${ignore($T)} self.${index()}, other.${index()} ),+)
+                    lexical_ord!(le, spec_chain_le, $( ${ignore($T)} self.${index()}, other.${index()} ),+)
                 }
                 #[inline]
                 fn ge(&self, other: &($($T,)+)) -> bool {
-                    lexical_ord!(ge, Greater, $( ${ignore($T)} self.${index()}, other.${index()} ),+)
+                    lexical_ord!(ge, spec_chain_ge, $( ${ignore($T)} self.${index()}, other.${index()} ),+)
                 }
                 #[inline]
                 fn gt(&self, other: &($($T,)+)) -> bool {
-                    lexical_ord!(gt, Greater, $( ${ignore($T)} self.${index()}, other.${index()} ),+)
+                    lexical_ord!(gt, spec_chain_gt, $( ${ignore($T)} self.${index()}, other.${index()} ),+)
                 }
             }
         }
@@ -171,15 +173,16 @@ macro_rules! maybe_tuple_doc {
 // `(a1, a2, a3) < (b1, b2, b3)` would be `lexical_ord!(lt, opt_is_lt, a1, b1,
 // a2, b2, a3, b3)` (and similarly for `lexical_cmp`)
 //
-// `$ne_rel` is only used to determine the result after checking that they're
-// not equal, so `lt` and `le` can both just use `Less`.
+// `$chain_rel` is the method from `SpecChainingPartialOrd` to use for all but the
+// final value, to produce better results for simple primitives.
 macro_rules! lexical_ord {
-    ($rel: ident, $ne_rel: ident, $a:expr, $b:expr, $($rest_a:expr, $rest_b:expr),+) => {{
-        let c = PartialOrd::partial_cmp(&$a, &$b);
-        if c != Some(Equal) { c == Some($ne_rel) }
-        else { lexical_ord!($rel, $ne_rel, $($rest_a, $rest_b),+) }
+    ($rel: ident, $chain_rel: ident, $a:expr, $b:expr, $($rest_a:expr, $rest_b:expr),+) => {{
+        match SpecChainingPartialOrd::$chain_rel(&$a, &$b) {
+            Break(val) => val,
+            Continue(()) => lexical_ord!($rel, $chain_rel, $($rest_a, $rest_b),+),
+        }
     }};
-    ($rel: ident, $ne_rel: ident, $a:expr, $b:expr) => {
+    ($rel: ident, $chain_rel: ident, $a:expr, $b:expr) => {
         // Use the specific method for the last element
         PartialOrd::$rel(&$a, &$b)
     };

tests/mir-opt/inline/exponential_runtime.rs

@@ -87,10 +87,15 @@ fn main() {
     // CHECK-LABEL: fn main(
     // CHECK-NOT: inlined
     // CHECK: (inlined <() as G>::call)
+    // CHECK-NOT: inlined
     // CHECK: (inlined <() as F>::call)
+    // CHECK-NOT: inlined
     // CHECK: (inlined <() as E>::call)
+    // CHECK-NOT: inlined
     // CHECK: (inlined <() as D>::call)
+    // CHECK-NOT: inlined
     // CHECK: (inlined <() as C>::call)
+    // CHECK-NOT: inlined
     // CHECK: (inlined <() as B>::call)
     // CHECK-NOT: inlined
     <() as G>::call();

tests/mir-opt/inline/inline_diverging.h.Inline.panic-abort.diff

@@ -6,6 +6,7 @@
       let _1: (!, !);
 +     let mut _2: fn() -> ! {sleep};
 +     let mut _7: ();
++     let mut _8: ();
 +     scope 1 (inlined call_twice::<!, fn() -> ! {sleep}>) {
 +         debug f => _2;
 +         let mut _3: &fn() -> ! {sleep};
@@ -17,6 +18,10 @@
 +             scope 3 {
 +                 debug b => _6;
 +             }
++             scope 6 (inlined <fn() -> ! {sleep} as Fn<()>>::call - shim(fn() -> ! {sleep})) {
++                 scope 7 (inlined sleep) {
++                 }
++             }
 +         }
 +         scope 4 (inlined <fn() -> ! {sleep} as Fn<()>>::call - shim(fn() -> ! {sleep})) {
 +             scope 5 (inlined sleep) {

tests/mir-opt/inline/inline_diverging.h.Inline.panic-unwind.diff

@@ -6,6 +6,7 @@
       let _1: (!, !);
 +     let mut _2: fn() -> ! {sleep};
 +     let mut _8: ();
++     let mut _9: ();
 +     scope 1 (inlined call_twice::<!, fn() -> ! {sleep}>) {
 +         debug f => _2;
 +         let mut _3: &fn() -> ! {sleep};
@@ -18,6 +19,10 @@
 +             scope 3 {
 +                 debug b => _6;
 +             }
++             scope 6 (inlined <fn() -> ! {sleep} as Fn<()>>::call - shim(fn() -> ! {sleep})) {
++                 scope 7 (inlined sleep) {
++                 }
++             }
 +         }
 +         scope 4 (inlined <fn() -> ! {sleep} as Fn<()>>::call - shim(fn() -> ! {sleep})) {
 +             scope 5 (inlined sleep) {