Auto merge of #155598 - makai410:partial-ord, r=nnethercote

Implement fast path for `derive(PartialOrd)` when deriving `Ord`



Closes: #155537

Unfortunately, this PR shares the same issue with #124794, which doesn't work when `derive(PartialOrd)` is before `derive(Ord)`.

Author: bors

compiler/rustc_builtin_macros/src/deriving/cmp/partial_ord.rs

@@ -1,7 +1,7 @@
-use rustc_ast::{ExprKind, ItemKind, MetaItem, PatKind, Safety};
+use rustc_ast::{ExprKind, ItemKind, MetaItem, PatKind, Safety, ast};
 use rustc_expand::base::{Annotatable, ExtCtxt};
 use rustc_span::{Ident, Span, sym};
-use thin_vec::thin_vec;
+use thin_vec::{ThinVec, thin_vec};
 
 use crate::deriving::generic::ty::*;
 use crate::deriving::generic::*;
@@ -41,6 +41,45 @@ pub(crate) fn expand_deriving_partial_ord(
     } else {
         true
     };
+
+    let container_id = cx.current_expansion.id.expn_data().parent.expect_local();
+    let has_derive_ord = cx.resolver.has_derive_ord(container_id);
+    let is_simple_candidate = |params: &ThinVec<ast::GenericParam>| -> bool {
+        has_derive_ord
+            && !params.iter().any(|param| matches!(param.kind, ast::GenericParamKind::Type { .. }))
+    };
+
+    let default_substructure = combine_substructure(Box::new(|cx, span, substr| {
+        cs_partial_cmp(cx, span, substr, discr_then_data)
+    }));
+    let simple_substructure = combine_substructure(Box::new(|cx, span, _| {
+        cs_partial_cmp_simple(cx, span, cx.expr_ident(span, Ident::new(sym::other, span)))
+    }));
+    let (is_simple, substructure) = match item {
+        Annotatable::Item(annitem) => match &annitem.kind {
+            // For unit structs/zero-variant enums, the default generated code is better.
+            ItemKind::Struct(.., ast::VariantData::Unit(..)) => (false, default_substructure),
+            // Also for single fieldless variant enum
+            ItemKind::Enum(.., enum_def) if enum_def.variants.is_empty() => {
+                (false, default_substructure)
+            }
+            ItemKind::Enum(.., enum_def)
+                if enum_def.variants.len() == 1
+                    && matches!(enum_def.variants[0].data, ast::VariantData::Unit(..)) =>
+            {
+                (false, default_substructure)
+            }
+            ItemKind::Struct(_, ast::Generics { params, .. }, _)
+            | ItemKind::Enum(_, ast::Generics { params, .. }, _)
+                if is_simple_candidate(params) =>
+            {
+                (true, simple_substructure)
+            }
+            _ => (false, default_substructure),
+        },
+        _ => (false, default_substructure),
+    };
+
     let partial_cmp_def = MethodDef {
         name: sym::partial_cmp,
         generics: Bounds::empty(),
@@ -49,9 +88,7 @@ pub(crate) fn expand_deriving_partial_ord(
         ret_ty,
         attributes: thin_vec![cx.attr_word(sym::inline, span)],
         fieldless_variants_strategy: FieldlessVariantsStrategy::Unify,
-        combine_substructure: combine_substructure(Box::new(|cx, span, substr| {
-            cs_partial_cmp(cx, span, substr, discr_then_data)
-        })),
+        combine_substructure: substructure,
     };
 
     let trait_def = TraitDef {
@@ -68,7 +105,18 @@ pub(crate) fn expand_deriving_partial_ord(
         safety: Safety::Default,
         document: true,
     };
-    trait_def.expand(cx, mitem, item, push)
+    trait_def.expand_ext(cx, mitem, item, push, is_simple)
+}
+
+// Special case for the type deriving both `PartialOrd` and `Ord`. Builds:
+// ```
+// Some(::core::cmp::Ord::cmp(self, other))
+// ```
+fn cs_partial_cmp_simple(cx: &ExtCtxt<'_>, span: Span, other_expr: Box<ast::Expr>) -> BlockOrExpr {
+    let ord_cmp_path = cx.std_path(&[sym::cmp, sym::Ord, sym::cmp]);
+    let cmp_expr =
+        cx.expr_call_global(span, ord_cmp_path, thin_vec![cx.expr_self(span), other_expr]);
+    BlockOrExpr::new_expr(cx.expr_some(span, cmp_expr))
 }
 
 fn cs_partial_cmp(
@@ -98,7 +146,8 @@ fn cs_partial_cmp(
         |cx, fold| match fold {
             CsFold::Single(field) => {
                 let [other_expr] = &field.other_selflike_exprs[..] else {
-                    cx.dcx().span_bug(field.span, "not exactly 2 arguments in `derive(Ord)`");
+                    cx.dcx()
+                        .span_bug(field.span, "not exactly 2 arguments in `derive(PartialOrd)`");
                 };
                 let args = thin_vec![field.self_expr.clone(), other_expr.clone()];
                 cx.expr_call_global(field.span, partial_cmp_path.clone(), args)

compiler/rustc_expand/src/base.rs

@@ -1117,6 +1117,8 @@ pub trait ResolverExpand {
     // Resolver interfaces for specific built-in macros.
     /// Does `#[derive(...)]` attribute with the given `ExpnId` have built-in `Copy` inside it?
     fn has_derive_copy(&self, expn_id: LocalExpnId) -> bool;
+    /// Does `#[derive(...)]` attribute with the given `ExpnId` have built-in `Ord` inside it?
+    fn has_derive_ord(&self, expn_id: LocalExpnId) -> bool;
     /// Resolve paths inside the `#[derive(...)]` attribute with the given `ExpnId`.
     fn resolve_derives(
         &mut self,

compiler/rustc_resolve/src/lib.rs

@@ -1323,7 +1323,10 @@ impl ExternPreludeEntry<'_> {
 struct DeriveData {
     resolutions: Vec<DeriveResolution>,
     helper_attrs: Vec<(usize, IdentKey, Span)>,
+    // if this list keeps getting extended, we could use `bitflags`,
+    // something like what [`rustc_type_ir::flags::TypeFlags`] is doing.
     has_derive_copy: bool,
+    has_derive_ord: bool,
 }
 
 pub struct ResolverOutputs<'tcx> {
@@ -1467,6 +1470,7 @@ pub struct Resolver<'ra, 'tcx> {
     /// Derive macros cannot modify the item themselves and have to store the markers in the global
     /// context, so they attach the markers to derive container IDs using this resolver table.
     containers_deriving_copy: FxHashSet<LocalExpnId> = default::fx_hash_set(),
+    containers_deriving_ord: FxHashSet<LocalExpnId> = default::fx_hash_set(),
     /// Parent scopes in which the macros were invoked.
     /// FIXME: `derives` are missing in these parent scopes and need to be taken from elsewhere.
     invocation_parent_scopes: FxHashMap<LocalExpnId, ParentScope<'ra>> = default::fx_hash_map(),

compiler/rustc_resolve/src/macros.rs

@@ -379,6 +379,10 @@ impl<'ra, 'tcx> ResolverExpand for Resolver<'ra, 'tcx> {
         self.containers_deriving_copy.contains(&expn_id)
     }
 
+    fn has_derive_ord(&self, expn_id: LocalExpnId) -> bool {
+        self.containers_deriving_ord.contains(&expn_id)
+    }
+
     fn resolve_derives(
         &mut self,
         expn_id: LocalExpnId,
@@ -398,6 +402,7 @@ impl<'ra, 'tcx> ResolverExpand for Resolver<'ra, 'tcx> {
             resolutions: derive_paths(),
             helper_attrs: Vec::new(),
             has_derive_copy: false,
+            has_derive_ord: false,
         });
         let parent_scope = self.invocation_parent_scopes[&expn_id];
         for (i, resolution) in entry.resolutions.iter_mut().enumerate() {
@@ -420,6 +425,7 @@ impl<'ra, 'tcx> ResolverExpand for Resolver<'ra, 'tcx> {
                                 );
                             }
                             entry.has_derive_copy |= ext.builtin_name == Some(sym::Copy);
+                            entry.has_derive_ord |= ext.builtin_name == Some(sym::Ord);
                             ext
                         }
                         Ok(_) | Err(Determinacy::Determined) => self.dummy_ext(MacroKind::Derive),
@@ -455,6 +461,12 @@ impl<'ra, 'tcx> ResolverExpand for Resolver<'ra, 'tcx> {
         if entry.has_derive_copy || self.has_derive_copy(parent_scope.expansion) {
             self.containers_deriving_copy.insert(expn_id);
         }
+        // Similar to the above `Copy` and `Clone` case, the code generated for
+        // `derive(PartialOrd)` changes if `derive(Ord)` is also present.
+        // FIXME(makai410): this also doesn't work with `#[derive(PartialOrd)] #[derive(Ord)]`.
+        if entry.has_derive_ord || self.has_derive_ord(parent_scope.expansion) {
+            self.containers_deriving_ord.insert(expn_id);
+        }
         assert!(self.derive_data.is_empty());
         self.derive_data = derive_data;
         Ok(())

tests/mir-opt/pre-codegen/derived_ord_debug.rs

@@ -10,8 +10,8 @@ pub struct MultiField(char, i16);
 // EMIT_MIR derived_ord_debug.{impl#1}-cmp.runtime-optimized.after.mir
 
 // CHECK-LABEL: partial_cmp(_1: &MultiField, _2: &MultiField) -> Option<std::cmp::Ordering>
-// CHECK: = <char as PartialOrd>::partial_cmp(
-// CHECK: = <i16 as PartialOrd>::partial_cmp(
+// CHECK: = <MultiField as Ord>::cmp(
+// CHECK: = Option::<std::cmp::Ordering>::Some(
 
 // CHECK-LABEL: cmp(_1: &MultiField, _2: &MultiField) -> std::cmp::Ordering
 // CHECK: = <char as Ord>::cmp(

tests/mir-opt/pre-codegen/derived_ord_debug.{impl#0}-partial_cmp.runtime-optimized.after.panic-abort.mir

@@ -4,49 +4,14 @@ fn <impl at $DIR/derived_ord_debug.rs:6:10: 6:20>::partial_cmp(_1: &MultiField,
     debug self => _1;
     debug other => _2;
     let mut _0: std::option::Option<std::cmp::Ordering>;
-    let _3: &char;
-    let _4: &char;
-    let mut _5: std::option::Option<std::cmp::Ordering>;
-    let mut _6: isize;
-    let mut _7: i8;
-    let _8: &i16;
-    let _9: &i16;
-    scope 1 {
-        debug cmp => _5;
-    }
+    let mut _3: std::cmp::Ordering;
 
     bb0: {
-        _3 = &((*_1).0: char);
-        _4 = &((*_2).0: char);
-        _5 = <char as PartialOrd>::partial_cmp(copy _3, copy _4) -> [return: bb1, unwind unreachable];
+        _3 = <MultiField as Ord>::cmp(copy _1, copy _2) -> [return: bb1, unwind unreachable];
     }
 
     bb1: {
-        _6 = discriminant(_5);
-        switchInt(move _6) -> [1: bb2, 0: bb4, otherwise: bb6];
-    }
-
-    bb2: {
-        _7 = discriminant(((_5 as Some).0: std::cmp::Ordering));
-        switchInt(move _7) -> [0: bb3, otherwise: bb4];
-    }
-
-    bb3: {
-        _8 = &((*_1).1: i16);
-        _9 = &((*_2).1: i16);
-        _0 = <i16 as PartialOrd>::partial_cmp(copy _8, copy _9) -> [return: bb5, unwind unreachable];
-    }
-
-    bb4: {
-        _0 = copy _5;
-        goto -> bb5;
-    }
-
-    bb5: {
+        _0 = Option::<std::cmp::Ordering>::Some(move _3);
         return;
     }
-
-    bb6: {
-        unreachable;
-    }
 }

tests/mir-opt/pre-codegen/derived_ord_debug.{impl#0}-partial_cmp.runtime-optimized.after.panic-unwind.mir

@@ -4,49 +4,14 @@ fn <impl at $DIR/derived_ord_debug.rs:6:10: 6:20>::partial_cmp(_1: &MultiField,
     debug self => _1;
     debug other => _2;
     let mut _0: std::option::Option<std::cmp::Ordering>;
-    let _3: &char;
-    let _4: &char;
-    let mut _5: std::option::Option<std::cmp::Ordering>;
-    let mut _6: isize;
-    let mut _7: i8;
-    let _8: &i16;
-    let _9: &i16;
-    scope 1 {
-        debug cmp => _5;
-    }
+    let mut _3: std::cmp::Ordering;
 
     bb0: {
-        _3 = &((*_1).0: char);
-        _4 = &((*_2).0: char);
-        _5 = <char as PartialOrd>::partial_cmp(copy _3, copy _4) -> [return: bb1, unwind continue];
+        _3 = <MultiField as Ord>::cmp(copy _1, copy _2) -> [return: bb1, unwind continue];
     }
 
     bb1: {
-        _6 = discriminant(_5);
-        switchInt(move _6) -> [1: bb2, 0: bb4, otherwise: bb6];
-    }
-
-    bb2: {
-        _7 = discriminant(((_5 as Some).0: std::cmp::Ordering));
-        switchInt(move _7) -> [0: bb3, otherwise: bb4];
-    }
-
-    bb3: {
-        _8 = &((*_1).1: i16);
-        _9 = &((*_2).1: i16);
-        _0 = <i16 as PartialOrd>::partial_cmp(copy _8, copy _9) -> [return: bb5, unwind continue];
-    }
-
-    bb4: {
-        _0 = copy _5;
-        goto -> bb5;
-    }
-
-    bb5: {
+        _0 = Option::<std::cmp::Ordering>::Some(move _3);
         return;
     }
-
-    bb6: {
-        unreachable;
-    }
 }

tests/ui/derives/deriving-all-codegen.rs

@@ -235,3 +235,21 @@ struct FooCopyAndClone(i32);
 #[derive(Clone)]
 #[derive(Copy)]
 struct FooCloneAndCopy(i32);
+
+#[derive(PartialOrd, Ord)]
+struct FooPartialOrdOrd(i32);
+
+#[derive(Ord, PartialOrd)]
+struct FooOrdPartialOrd(i32);
+
+#[derive(Ord)]
+#[derive(PartialOrd)]
+struct FooOrdBeforePartialOrd(i32);
+
+// FIXME: this case should also have a trivial `PartialOrd` impl.
+#[derive(PartialOrd)]
+#[derive(Ord)]
+struct FooPartialOrdBeforeOrd(i32);
+
+#[derive(PartialOrd, Ord)]
+struct UnitStruct;