Rollup merge of #138135 - scottmcm:chaining-ord, r=Mark-Simulacrum

Simplify `PartialOrd` on tuples containing primitives

We noticed in https://github.com/rust-lang/rust/pull/133984#issuecomment-2704011800 that currently the tuple comparison code, while it [does optimize down](https://github.com/rust-lang/rust/blob/master/tests/codegen/comparison-operators-2-tuple.rs) today, is kinda huge: <https://rust.godbolt.org/z/xqMoeYbhE>

This PR changes the tuple code to go through an overridable "chaining" version of the comparison functions, so that for simple things like `(i16, u16)` and `(f32, f32)` (as seen in the new MIR pre-codegen test) we just directly get the
```rust
if lhs.0 == rhs.0 { lhs.0 OP rhs.0 }
else { lhs.1 OP rhs.1 }
```
version in MIR, rather than emitting a mess for LLVM to have to clean up.

Test added in the first commit, so you can see the MIR diff in the second one.

5 files changed, 265 insertions, 11 deletions

diff --git a/library/core/src/cmp.rs b/library/core/src/cmp.rs
index 25bd17d5802..0b0dbf723b6 100644
--- a/library/core/src/cmp.rs
+++ b/library/core/src/cmp.rs
@@ -29,6 +29,7 @@ mod bytewise;
 pub(crate) use bytewise::BytewiseEq;
 
 use self::Ordering::*;
+use crate::ops::ControlFlow;
 
 /// Trait for comparisons using the equality operator.
 ///
@@ -1435,6 +1436,67 @@ pub trait PartialOrd<Rhs: ?Sized = Self>: PartialEq<Rhs> {
     fn ge(&self, other: &Rhs) -> bool {
         self.partial_cmp(other).is_some_and(Ordering::is_ge)
     }
+
+    /// If `self == other`, returns `ControlFlow::Continue(())`.
+    /// Otherwise, returns `ControlFlow::Break(self < other)`.
+    ///
+    /// This is useful for chaining together calls when implementing a lexical
+    /// `PartialOrd::lt`, as it allows types (like primitives) which can cheaply
+    /// check `==` and `<` separately to do rather than needing to calculate
+    /// (then optimize out) the three-way `Ordering` result.
+    #[inline]
+    #[must_use]
+    // Added to improve the behaviour of tuples; not necessarily stabilization-track.
+    #[unstable(feature = "partial_ord_chaining_methods", issue = "none")]
+    #[doc(hidden)]
+    fn __chaining_lt(&self, other: &Rhs) -> ControlFlow<bool> {
+        default_chaining_impl(self, other, Ordering::is_lt)
+    }
+
+    /// Same as `__chaining_lt`, but for `<=` instead of `<`.
+    #[inline]
+    #[must_use]
+    #[unstable(feature = "partial_ord_chaining_methods", issue = "none")]
+    #[doc(hidden)]
+    fn __chaining_le(&self, other: &Rhs) -> ControlFlow<bool> {
+        default_chaining_impl(self, other, Ordering::is_le)
+    }
+
+    /// Same as `__chaining_lt`, but for `>` instead of `<`.
+    #[inline]
+    #[must_use]
+    #[unstable(feature = "partial_ord_chaining_methods", issue = "none")]
+    #[doc(hidden)]
+    fn __chaining_gt(&self, other: &Rhs) -> ControlFlow<bool> {
+        default_chaining_impl(self, other, Ordering::is_gt)
+    }
+
+    /// Same as `__chaining_lt`, but for `>=` instead of `<`.
+    #[inline]
+    #[must_use]
+    #[unstable(feature = "partial_ord_chaining_methods", issue = "none")]
+    #[doc(hidden)]
+    fn __chaining_ge(&self, other: &Rhs) -> ControlFlow<bool> {
+        default_chaining_impl(self, other, Ordering::is_ge)
+    }
+}
+
+fn default_chaining_impl<T: ?Sized, U: ?Sized>(
+    lhs: &T,
+    rhs: &U,
+    p: impl FnOnce(Ordering) -> bool,
+) -> ControlFlow<bool>
+where
+    T: PartialOrd<U>,
+{
+    // It's important that this only call `partial_cmp` once, not call `eq` then
+    // one of the relational operators.  We don't want to `bcmp`-then-`memcp` a
+    // `String`, for example, or similarly for other data structures (#108157).
+    match <T as PartialOrd<U>>::partial_cmp(lhs, rhs) {
+        Some(Equal) => ControlFlow::Continue(()),
+        Some(c) => ControlFlow::Break(p(c)),
+        None => ControlFlow::Break(false),
+    }
 }
 
 /// Derive macro generating an impl of the trait [`PartialOrd`].
@@ -1741,6 +1803,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 +1842,35 @@ mod impls {
 
     eq_impl! { () bool char usize u8 u16 u32 u64 u128 isize i8 i16 i32 i64 i128 }
 
+    macro_rules! chaining_methods_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`.
+
+            #[inline]
+            fn __chaining_lt(&self, other: &Self) -> ControlFlow<bool> {
+                let (lhs, rhs) = (*self, *other);
+                if lhs == rhs { Continue(()) } else { Break(lhs < rhs) }
+            }
+            #[inline]
+            fn __chaining_le(&self, other: &Self) -> ControlFlow<bool> {
+                let (lhs, rhs) = (*self, *other);
+                if lhs == rhs { Continue(()) } else { Break(lhs <= rhs) }
+            }
+            #[inline]
+            fn __chaining_gt(&self, other: &Self) -> ControlFlow<bool> {
+                let (lhs, rhs) = (*self, *other);
+                if lhs == rhs { Continue(()) } else { Break(lhs > rhs) }
+            }
+            #[inline]
+            fn __chaining_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")]
@@ -1800,6 +1892,8 @@ mod impls {
                 fn ge(&self, other: &$t) -> bool { (*self) >= (*other) }
                 #[inline(always)]
                 fn gt(&self, other: &$t) -> bool { (*self) > (*other) }
+
+                chaining_methods_impl!($t);
             }
         )*)
     }
@@ -1838,6 +1932,8 @@ mod impls {
                 fn ge(&self, other: &$t) -> bool { (*self) >= (*other) }
                 #[inline(always)]
                 fn gt(&self, other: &$t) -> bool { (*self) > (*other) }
+
+                chaining_methods_impl!($t);
             }
 
             #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/library/core/src/tuple.rs b/library/core/src/tuple.rs
index 206b5b9e2c2..d754bb90343 100644
--- a/library/core/src/tuple.rs
+++ b/library/core/src/tuple.rs
@@ -2,6 +2,7 @@
 
 use crate::cmp::Ordering::{self, *};
 use crate::marker::{ConstParamTy_, StructuralPartialEq, UnsizedConstParamTy};
+use crate::ops::ControlFlow::{Break, Continue};
 
 // Recursive macro for implementing n-ary tuple functions and operations
 //
@@ -80,19 +81,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, __chaining_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, __chaining_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, __chaining_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, __chaining_gt, $( ${ignore($T)} self.${index()}, other.${index()} ),+)
                 }
             }
         }
@@ -171,15 +172,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 chaining method from `PartialOrd` 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 PartialOrd::$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)
     };
diff --git a/tests/mir-opt/pre-codegen/tuple_ord.demo_ge_partial.PreCodegen.after.mir b/tests/mir-opt/pre-codegen/tuple_ord.demo_ge_partial.PreCodegen.after.mir
new file mode 100644
index 00000000000..dd2eebc8f4a
--- /dev/null
+++ b/tests/mir-opt/pre-codegen/tuple_ord.demo_ge_partial.PreCodegen.after.mir
@@ -0,0 +1,70 @@
+// MIR for `demo_ge_partial` after PreCodegen
+
+fn demo_ge_partial(_1: &(f32, f32), _2: &(f32, f32)) -> bool {
+    debug a => _1;
+    debug b => _2;
+    let mut _0: bool;
+    scope 1 (inlined std::cmp::impls::<impl PartialOrd for &(f32, f32)>::ge) {
+        scope 2 (inlined core::tuple::<impl PartialOrd for (f32, f32)>::ge) {
+            let mut _7: std::ops::ControlFlow<bool>;
+            let _8: bool;
+            scope 3 {
+            }
+            scope 4 (inlined std::cmp::impls::<impl PartialOrd for f32>::__chaining_ge) {
+                let mut _3: f32;
+                let mut _4: f32;
+                let mut _5: bool;
+                let mut _6: bool;
+                scope 5 {
+                }
+            }
+            scope 6 (inlined std::cmp::impls::<impl PartialOrd for f32>::ge) {
+                let mut _9: f32;
+                let mut _10: f32;
+            }
+        }
+    }
+
+    bb0: {
+        StorageLive(_7);
+        StorageLive(_3);
+        StorageLive(_4);
+        _3 = copy ((*_1).0: f32);
+        _4 = copy ((*_2).0: f32);
+        StorageLive(_5);
+        _5 = Eq(copy _3, copy _4);
+        switchInt(move _5) -> [0: bb1, otherwise: bb2];
+    }
+
+    bb1: {
+        StorageLive(_6);
+        _6 = Ge(copy _3, copy _4);
+        _7 = ControlFlow::<bool>::Break(move _6);
+        StorageDead(_6);
+        StorageDead(_5);
+        StorageDead(_4);
+        StorageDead(_3);
+        _8 = copy ((_7 as Break).0: bool);
+        _0 = copy _8;
+        goto -> bb3;
+    }
+
+    bb2: {
+        StorageDead(_5);
+        StorageDead(_4);
+        StorageDead(_3);
+        StorageLive(_9);
+        _9 = copy ((*_1).1: f32);
+        StorageLive(_10);
+        _10 = copy ((*_2).1: f32);
+        _0 = Ge(move _9, move _10);
+        StorageDead(_10);
+        StorageDead(_9);
+        goto -> bb3;
+    }
+
+    bb3: {
+        StorageDead(_7);
+        return;
+    }
+}
diff --git a/tests/mir-opt/pre-codegen/tuple_ord.demo_le_total.PreCodegen.after.mir b/tests/mir-opt/pre-codegen/tuple_ord.demo_le_total.PreCodegen.after.mir
new file mode 100644
index 00000000000..ea1d164cefa
--- /dev/null
+++ b/tests/mir-opt/pre-codegen/tuple_ord.demo_le_total.PreCodegen.after.mir
@@ -0,0 +1,70 @@
+// MIR for `demo_le_total` after PreCodegen
+
+fn demo_le_total(_1: &(u16, i16), _2: &(u16, i16)) -> bool {
+    debug a => _1;
+    debug b => _2;
+    let mut _0: bool;
+    scope 1 (inlined std::cmp::impls::<impl PartialOrd for &(u16, i16)>::le) {
+        scope 2 (inlined core::tuple::<impl PartialOrd for (u16, i16)>::le) {
+            let mut _7: std::ops::ControlFlow<bool>;
+            let _8: bool;
+            scope 3 {
+            }
+            scope 4 (inlined std::cmp::impls::<impl PartialOrd for u16>::__chaining_le) {
+                let mut _3: u16;
+                let mut _4: u16;
+                let mut _5: bool;
+                let mut _6: bool;
+                scope 5 {
+                }
+            }
+            scope 6 (inlined std::cmp::impls::<impl PartialOrd for i16>::le) {
+                let mut _9: i16;
+                let mut _10: i16;
+            }
+        }
+    }
+
+    bb0: {
+        StorageLive(_7);
+        StorageLive(_3);
+        StorageLive(_4);
+        _3 = copy ((*_1).0: u16);
+        _4 = copy ((*_2).0: u16);
+        StorageLive(_5);
+        _5 = Eq(copy _3, copy _4);
+        switchInt(move _5) -> [0: bb1, otherwise: bb2];
+    }
+
+    bb1: {
+        StorageLive(_6);
+        _6 = Le(copy _3, copy _4);
+        _7 = ControlFlow::<bool>::Break(move _6);
+        StorageDead(_6);
+        StorageDead(_5);
+        StorageDead(_4);
+        StorageDead(_3);
+        _8 = copy ((_7 as Break).0: bool);
+        _0 = copy _8;
+        goto -> bb3;
+    }
+
+    bb2: {
+        StorageDead(_5);
+        StorageDead(_4);
+        StorageDead(_3);
+        StorageLive(_9);
+        _9 = copy ((*_1).1: i16);
+        StorageLive(_10);
+        _10 = copy ((*_2).1: i16);
+        _0 = Le(move _9, move _10);
+        StorageDead(_10);
+        StorageDead(_9);
+        goto -> bb3;
+    }
+
+    bb3: {
+        StorageDead(_7);
+        return;
+    }
+}
diff --git a/tests/mir-opt/pre-codegen/tuple_ord.rs b/tests/mir-opt/pre-codegen/tuple_ord.rs
new file mode 100644
index 00000000000..74a919e5424
--- /dev/null
+++ b/tests/mir-opt/pre-codegen/tuple_ord.rs
@@ -0,0 +1,16 @@
+//@ compile-flags: -O -Zmir-opt-level=2 -Cdebuginfo=0
+//@ needs-unwind
+
+#![crate_type = "lib"]
+
+// EMIT_MIR tuple_ord.demo_le_total.PreCodegen.after.mir
+pub fn demo_le_total(a: &(u16, i16), b: &(u16, i16)) -> bool {
+    // CHECK-LABEL: demo_le_total
+    a <= b
+}
+
+// EMIT_MIR tuple_ord.demo_ge_partial.PreCodegen.after.mir
+pub fn demo_ge_partial(a: &(f32, f32), b: &(f32, f32)) -> bool {
+    // CHECK-LABEL: demo_ge_partial
+    a >= b
+}