Experiment: spaceship in MIR

As mentioned in <rust-lang#108788 (comment)>, `Ord::cmp` on primitives generates a large amount of MIR, preventing (or at least discouraging) it from mir-inlining.

Let's see whether making it a MIR primitive helps stuff -- derived `(Partial)Ord` in particular, if we're lucky.

Author: scottmcm

compiler/rustc_codegen_cranelift/src/codegen_i128.rs

@@ -148,6 +148,7 @@ pub(crate) fn maybe_codegen<'tcx>(
             assert!(!checked);
             None
         }
+        BinOp::Cmp => None,
         BinOp::Shl | BinOp::Shr => None,
     }
 }

compiler/rustc_codegen_cranelift/src/num.rs

@@ -40,13 +40,34 @@ pub(crate) fn bin_op_to_intcc(bin_op: BinOp, signed: bool) -> Option<IntCC> {
     })
 }
 
+fn codegen_three_way_compare<'tcx>(
+    fx: &mut FunctionCx<'_, '_, 'tcx>,
+    signed: bool,
+    lhs: Value,
+    rhs: Value,
+) -> CValue<'tcx> {
+    let gt_cc = crate::num::bin_op_to_intcc(BinOp::Gt, signed).unwrap();
+    let lt_cc = crate::num::bin_op_to_intcc(BinOp::Lt, signed).unwrap();
+    let gt = fx.bcx.ins().icmp(gt_cc, lhs, rhs);
+    let lt = fx.bcx.ins().icmp(lt_cc, lhs, rhs);
+    // Cranelift no longer has a single-bit type, so the comparison results
+    // are already `I8`s, which we can subtract to get the -1/0/+1 we want.
+    // See <https://github.com/bytecodealliance/wasmtime/pull/5031>.
+    let val = fx.bcx.ins().isub(gt, lt);
+    CValue::by_val(val, fx.layout_of(fx.tcx.types.i8))
+}
+
 fn codegen_compare_bin_op<'tcx>(
     fx: &mut FunctionCx<'_, '_, 'tcx>,
     bin_op: BinOp,
     signed: bool,
     lhs: Value,
     rhs: Value,
 ) -> CValue<'tcx> {
+    if bin_op == BinOp::Cmp {
+        return codegen_three_way_compare(fx, signed, lhs, rhs);
+    }
+
     let intcc = crate::num::bin_op_to_intcc(bin_op, signed).unwrap();
     let val = fx.bcx.ins().icmp(intcc, lhs, rhs);
     CValue::by_val(val, fx.layout_of(fx.tcx.types.bool))
@@ -59,7 +80,7 @@ pub(crate) fn codegen_binop<'tcx>(
     in_rhs: CValue<'tcx>,
 ) -> CValue<'tcx> {
     match bin_op {
-        BinOp::Eq | BinOp::Lt | BinOp::Le | BinOp::Ne | BinOp::Ge | BinOp::Gt => {
+        BinOp::Eq | BinOp::Lt | BinOp::Le | BinOp::Ne | BinOp::Ge | BinOp::Gt | BinOp::Cmp => {
             match in_lhs.layout().ty.kind() {
                 ty::Bool | ty::Uint(_) | ty::Int(_) | ty::Char => {
                     let signed = type_sign(in_lhs.layout().ty);

compiler/rustc_codegen_gcc/src/common.rs

@@ -97,6 +97,10 @@ impl<'gcc, 'tcx> ConstMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
         self.const_int(self.type_i32(), i as i64)
     }
 
+    fn const_i8(&self, i: i8) -> RValue<'gcc> {
+        self.const_int(self.type_i8(), i as i64)
+    }
+
     fn const_u32(&self, i: u32) -> RValue<'gcc> {
         self.const_uint(self.type_u32(), i as u64)
     }

compiler/rustc_codegen_llvm/src/common.rs

@@ -157,6 +157,10 @@ impl<'ll, 'tcx> ConstMethods<'tcx> for CodegenCx<'ll, 'tcx> {
         self.const_int(self.type_i32(), i as i64)
     }
 
+    fn const_i8(&self, i: i8) -> &'ll Value {
+        self.const_int(self.type_i8(), i as i64)
+    }
+
     fn const_u32(&self, i: u32) -> &'ll Value {
         self.const_uint(self.type_i32(), i as u64)
     }

compiler/rustc_codegen_ssa/src/mir/rvalue.rs

@@ -7,6 +7,7 @@ use crate::common::{self, IntPredicate};
 use crate::traits::*;
 use crate::MemFlags;
 
+use rustc_hir as hir;
 use rustc_middle::mir;
 use rustc_middle::mir::Operand;
 use rustc_middle::ty::cast::{CastTy, IntTy};
@@ -599,6 +600,27 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                     bx.icmp(base::bin_op_to_icmp_predicate(op.to_hir_binop(), is_signed), lhs, rhs)
                 }
             }
+            mir::BinOp::Cmp => {
+                use std::cmp::Ordering;
+                debug_assert!(!is_float);
+                // FIXME: To avoid this PR changing behaviour, the operations used
+                // here are those from <https://github.com/rust-lang/rust/pull/63767>,
+                // as tested by `tests/codegen/integer-cmp.rs`.
+                // Something in future might want to pick different ones. For example,
+                // maybe the ones from Clang's `<=>` operator in C++20 (see
+                // <https://github.com/llvm/llvm-project/issues/60012>) or once we
+                // update to new LLVM, something to take advantage of the new folds in
+                // <https://github.com/llvm/llvm-project/issues/59666>.
+                let pred = |op| base::bin_op_to_icmp_predicate(op, is_signed);
+                let is_lt = bx.icmp(pred(hir::BinOpKind::Lt), lhs, rhs);
+                let is_ne = bx.icmp(pred(hir::BinOpKind::Ne), lhs, rhs);
+                let ge = bx.select(
+                    is_ne,
+                    bx.cx().const_i8(Ordering::Greater as i8),
+                    bx.cx().const_i8(Ordering::Equal as i8),
+                );
+                bx.select(is_lt, bx.cx().const_i8(Ordering::Less as i8), ge)
+            }
         }
     }
 

compiler/rustc_codegen_ssa/src/traits/consts.rs

@@ -14,6 +14,7 @@ pub trait ConstMethods<'tcx>: BackendTypes {
     fn const_bool(&self, val: bool) -> Self::Value;
     fn const_i16(&self, i: i16) -> Self::Value;
     fn const_i32(&self, i: i32) -> Self::Value;
+    fn const_i8(&self, i: i8) -> Self::Value;
     fn const_u32(&self, i: u32) -> Self::Value;
     fn const_u64(&self, i: u64) -> Self::Value;
     fn const_usize(&self, i: u64) -> Self::Value;

compiler/rustc_const_eval/src/interpret/operator.rs

@@ -59,6 +59,11 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
 }
 
 impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
+    fn three_way_compare<T: Ord>(&self, lhs: T, rhs: T) -> (Scalar<M::Provenance>, bool, Ty<'tcx>) {
+        let res = Ord::cmp(&lhs, &rhs) as i8;
+        return (Scalar::from_i8(res), false, self.tcx.ty_ordering_enum(None));
+    }
+
     fn binary_char_op(
         &self,
         bin_op: mir::BinOp,
@@ -67,6 +72,10 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
     ) -> (Scalar<M::Provenance>, bool, Ty<'tcx>) {
         use rustc_middle::mir::BinOp::*;
 
+        if bin_op == Cmp {
+            return self.three_way_compare(l, r);
+        }
+
         let res = match bin_op {
             Eq => l == r,
             Ne => l != r,
@@ -211,6 +220,11 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                 let r = self.sign_extend(r, right_layout) as i128;
                 return Ok((Scalar::from_bool(op(&l, &r)), false, self.tcx.types.bool));
             }
+            if bin_op == Cmp {
+                let l = self.sign_extend(l, left_layout) as i128;
+                let r = self.sign_extend(r, right_layout) as i128;
+                return Ok(self.three_way_compare(l, r));
+            }
             let op: Option<fn(i128, i128) -> (i128, bool)> = match bin_op {
                 Div if r == 0 => throw_ub!(DivisionByZero),
                 Rem if r == 0 => throw_ub!(RemainderByZero),
@@ -250,6 +264,10 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             }
         }
 
+        if bin_op == Cmp {
+            return Ok(self.three_way_compare(l, r));
+        }
+
         let (val, ty) = match bin_op {
             Eq => (Scalar::from_bool(l == r), self.tcx.types.bool),
             Ne => (Scalar::from_bool(l != r), self.tcx.types.bool),

compiler/rustc_const_eval/src/interpret/step.rs

@@ -17,7 +17,7 @@ fn binop_left_homogeneous(op: mir::BinOp) -> bool {
     use rustc_middle::mir::BinOp::*;
     match op {
         Add | Sub | Mul | Div | Rem | BitXor | BitAnd | BitOr | Offset | Shl | Shr => true,
-        Eq | Ne | Lt | Le | Gt | Ge => false,
+        Eq | Ne | Lt | Le | Gt | Ge | Cmp => false,
     }
 }
 /// Classify whether an operator is "right-homogeneous", i.e., the RHS has the
@@ -26,7 +26,8 @@ fn binop_left_homogeneous(op: mir::BinOp) -> bool {
 fn binop_right_homogeneous(op: mir::BinOp) -> bool {
     use rustc_middle::mir::BinOp::*;
     match op {
-        Add | Sub | Mul | Div | Rem | BitXor | BitAnd | BitOr | Eq | Ne | Lt | Le | Gt | Ge => true,
+        Add | Sub | Mul | Div | Rem | BitXor | BitAnd | BitOr | Eq | Ne | Lt | Le | Gt | Ge
+        | Cmp => true,
         Offset | Shl | Shr => false,
     }
 }

compiler/rustc_const_eval/src/transform/promote_consts.rs

@@ -568,6 +568,7 @@ impl<'tcx> Validator<'_, 'tcx> {
                     | BinOp::Lt
                     | BinOp::Ge
                     | BinOp::Gt
+                    | BinOp::Cmp
                     | BinOp::Offset
                     | BinOp::Add
                     | BinOp::Sub

compiler/rustc_const_eval/src/transform/validate.rs

@@ -483,6 +483,15 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                             );
                         }
                     }
+                    Cmp => {
+                        for x in [a, b] {
+                            check_kinds!(
+                                x,
+                                "Cannot three-way compare non-integer type {:?}",
+                                ty::Char | ty::Uint(..) | ty::Int(..)
+                            )
+                        }
+                    }
                     Shl | Shr => {
                         for x in [a, b] {
                             check_kinds!(