Auto merge of #155216 - jakubadamw:issue-110870-103073, r=<try>

match: Use an aggregate equality comparison for constant array/slice patterns when possible

Author: bors

compiler/rustc_mir_build/src/builder/matches/buckets.rs

@@ -323,6 +323,10 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                     value: case_val,
                     kind: PatConstKind::Float | PatConstKind::Other,
                 },
+            )
+            | (
+                TestKind::AggregateEq { value: test_val, .. },
+                TestableCase::Constant { value: case_val, kind: PatConstKind::Aggregate },
             ) => {
                 if test_val == case_val {
                     fully_matched = true;
@@ -353,6 +357,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 | TestKind::Range { .. }
                 | TestKind::StringEq { .. }
                 | TestKind::ScalarEq { .. }
+                | TestKind::AggregateEq { .. }
                 | TestKind::Deref { .. },
                 _,
             ) => {

compiler/rustc_mir_build/src/builder/matches/match_pair.rs

@@ -4,15 +4,56 @@ use rustc_abi::FieldIdx;
 use rustc_middle::mir::*;
 use rustc_middle::span_bug;
 use rustc_middle::thir::*;
-use rustc_middle::ty::{self, Ty, TypeVisitableExt};
+use rustc_middle::ty::{self, Ty, TyCtxt, TypeVisitableExt};
+use rustc_span::sym;
 
 use crate::builder::Builder;
 use crate::builder::expr::as_place::{PlaceBase, PlaceBuilder};
 use crate::builder::matches::{
     FlatPat, MatchPairTree, PatConstKind, PatternExtraData, SliceLenOp, TestableCase,
 };
 
+/// Checks whether every pattern in `elements` is a `PatKind::Constant` and,
+/// if so, reconstructs a single aggregate `ty::Value` that represents the whole
+/// array or slice. Returns `None` when any element is not a constant or the
+/// sequence is too short to benefit from an aggregate comparison.
+fn try_reconstruct_aggregate_constant<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    aggregate_ty: Ty<'tcx>,
+    elements: &[Pat<'tcx>],
+) -> Option<ty::Value<'tcx>> {
+    // A single element (or empty array) is not worth an aggregate comparison.
+    if elements.len() <= 1 {
+        return None;
+    }
+    let branches = elements
+        .iter()
+        .map(|pat| {
+            if let PatKind::Constant { value } = pat.kind {
+                Some(ty::Const::new_value(tcx, value.valtree, value.ty))
+            } else {
+                None
+            }
+        })
+        .collect::<Option<Vec<_>>>()?;
+    let valtree = ty::ValTree::from_branches(tcx, branches);
+    Some(ty::Value { ty: aggregate_ty, valtree })
+}
+
 impl<'a, 'tcx> Builder<'a, 'tcx> {
+    /// Check if we can use aggregate `PartialEq::eq` comparisons for constant array/slice patterns.
+    /// This is not possible in const contexts unless `#![feature(const_cmp, const_trait_impl)]` are enabled,
+    /// because`PartialEq` is not const-stable.
+    fn can_use_aggregate_eq(&self) -> bool {
+        let const_partial_eq_enabled = {
+            let features = self.tcx.features();
+            features.enabled(sym::const_trait_impl) && features.enabled(sym::const_cmp)
+        };
+        let in_const_context = self.tcx.is_const_fn(self.def_id.to_def_id())
+            || !self.tcx.hir_body_owner_kind(self.def_id).is_fn_or_closure();
+        !in_const_context || const_partial_eq_enabled
+    }
+
     /// Builds and pushes [`MatchPairTree`] subtrees, one for each pattern in
     /// `subpatterns`, representing the fields of a [`PatKind::Variant`] or
     /// [`PatKind::Leaf`].
@@ -239,15 +280,31 @@ impl<'tcx> MatchPairTree<'tcx> {
                     _ => None,
                 };
                 if let Some(array_len) = array_len {
-                    cx.prefix_slice_suffix(
-                        &mut subpairs,
-                        extra_data,
-                        &place_builder,
-                        Some(array_len),
-                        prefix,
-                        slice,
-                        suffix,
-                    );
+                    // When all elements are constants and there is no `..`
+                    // subpattern, compare the whole array at once via
+                    // `PartialEq::eq` rather than element by element.
+                    if slice.is_none()
+                        && suffix.is_empty()
+                        && cx.can_use_aggregate_eq()
+                        && let Some(aggregate_value) =
+                            try_reconstruct_aggregate_constant(cx.tcx, pattern.ty, prefix)
+                    {
+                        Some(TestableCase::Constant {
+                            value: aggregate_value,
+                            kind: PatConstKind::Aggregate,
+                        })
+                    } else {
+                        cx.prefix_slice_suffix(
+                            &mut subpairs,
+                            extra_data,
+                            &place_builder,
+                            Some(array_len),
+                            prefix,
+                            slice,
+                            suffix,
+                        );
+                        None
+                    }
                 } else {
                     // If the array length couldn't be determined, ignore the
                     // subpatterns and delayed-assert that compilation will fail.
@@ -258,37 +315,61 @@ impl<'tcx> MatchPairTree<'tcx> {
                             pattern.ty
                         ),
                     );
+                    None
                 }
-
-                None
             }
             PatKind::Slice { ref prefix, ref slice, ref suffix } => {
-                cx.prefix_slice_suffix(
-                    &mut subpairs,
-                    extra_data,
-                    &place_builder,
-                    None,
-                    prefix,
-                    slice,
-                    suffix,
-                );
-
-                if prefix.is_empty() && slice.is_some() && suffix.is_empty() {
-                    // This pattern is shaped like `[..]`. It can match a slice
-                    // of any length, so no length test is needed.
-                    None
-                } else {
-                    // Any other shape of slice pattern requires a length test.
-                    // Slice patterns with a `..` subpattern require a minimum
-                    // length; those without `..` require an exact length.
-                    Some(TestableCase::Slice {
-                        len: u64::try_from(prefix.len() + suffix.len()).unwrap(),
-                        op: if slice.is_some() {
-                            SliceLenOp::GreaterOrEqual
-                        } else {
-                            SliceLenOp::Equal
+                // When there is no `..`, all elements are constants, and
+                // there are at least two of them, collapse the individual
+                // element subpairs into a single aggregate comparison that
+                // is performed after the length check.
+                if slice.is_none()
+                    && suffix.is_empty()
+                    && cx.can_use_aggregate_eq()
+                    && let Some(aggregate_value) =
+                        try_reconstruct_aggregate_constant(cx.tcx, pattern.ty, prefix)
+                {
+                    subpairs.push(MatchPairTree {
+                        place,
+                        testable_case: TestableCase::Constant {
+                            value: aggregate_value,
+                            kind: PatConstKind::Aggregate,
                         },
+                        subpairs: Vec::new(),
+                        pattern_span: pattern.span,
+                    });
+                    Some(TestableCase::Slice {
+                        len: u64::try_from(prefix.len()).unwrap(),
+                        op: SliceLenOp::Equal,
                     })
+                } else {
+                    cx.prefix_slice_suffix(
+                        &mut subpairs,
+                        extra_data,
+                        &place_builder,
+                        None,
+                        prefix,
+                        slice,
+                        suffix,
+                    );
+
+                    if prefix.is_empty() && slice.is_some() && suffix.is_empty() {
+                        // This pattern is shaped like `[..]`. It can match
+                        // a slice of any length, so no length test is needed.
+                        None
+                    } else {
+                        // Any other shape of slice pattern requires a length test.
+                        // Slice patterns with a `..` subpattern require a minimum
+                        // length; those without `..` require an exact length.
+                        Some(TestableCase::Slice {
+                            len: u64::try_from(prefix.len() + suffix.len()).unwrap(),
+                            op: if slice.is_some() {
+                                SliceLenOp::GreaterOrEqual
+                            } else {
+                                SliceLenOp::Equal
+                            },
+                        })
+                    }
                 }
             }
 

compiler/rustc_mir_build/src/builder/matches/mod.rs

@@ -1266,6 +1266,10 @@ enum PatConstKind {
     Float,
     /// Constant string values, tested via string equality.
     String,
+    /// Constant array or slice values where every element is a constant.
+    /// Tested by calling `PartialEq::eq` on the whole aggregate at once,
+    /// rather than comparing element by element.
+    Aggregate,
     /// Any other constant-pattern is usually tested via some kind of equality
     /// check. Types that might be encountered here include:
     /// - raw pointers derived from integer values
@@ -1351,6 +1355,10 @@ enum TestKind<'tcx> {
     /// Tests the place against a constant using scalar equality.
     ScalarEq { value: ty::Value<'tcx> },
 
+    /// Tests the place against a constant array or slice using `PartialEq::eq`,
+    /// comparing the whole aggregate at once rather than element by element.
+    AggregateEq { value: ty::Value<'tcx> },
+
     /// Test whether the value falls within an inclusive or exclusive range.
     Range(Arc<PatRange<'tcx>>),
 

compiler/rustc_mir_build/src/builder/matches/test.rs

@@ -40,6 +40,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             TestableCase::Constant { value, kind: PatConstKind::String } => {
                 TestKind::StringEq { value }
             }
+            TestableCase::Constant { value, kind: PatConstKind::Aggregate } => {
+                TestKind::AggregateEq { value }
+            }
             TestableCase::Constant { value, kind: PatConstKind::Float | PatConstKind::Other } => {
                 TestKind::ScalarEq { value }
             }
@@ -168,16 +171,54 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 // Compare two strings using `<str as std::cmp::PartialEq>::eq`.
                 // (Interestingly this means that exhaustiveness analysis relies, for soundness,
                 // on the `PartialEq` impl for `str` to be correct!)
-                self.string_compare(
+                self.non_scalar_compare(
                     block,
                     success_block,
                     fail_block,
                     source_info,
+                    tcx.types.str_,
                     expected_value_operand,
                     Operand::Copy(actual_value_ref_place),
                 );
             }
 
+            TestKind::AggregateEq { value } => {
+                let tcx = self.tcx;
+                let success_block = target_block(TestBranch::Success);
+                let fail_block = target_block(TestBranch::Failure);
+
+                let aggregate_ty = value.ty;
+                let ref_ty = Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, aggregate_ty);
+
+                // The constant has type `[T; N]` (or `[T]`), but calling
+                // `PartialEq::eq` requires `&[T; N]` (or `&[T]`) operands.
+                // Valtree representations are the same with or without the
+                // reference wrapper, so we can reinterpret by replacing the type.
+                let expected_value = ty::Value { ty: ref_ty, valtree: value.valtree };
+                let expected_operand =
+                    self.literal_operand(test.span, Const::from_ty_value(tcx, expected_value));
+
+                // Create a reference to the scrutinee place.
+                let actual_ref_place = self.temp(ref_ty, test.span);
+                self.cfg.push_assign(
+                    block,
+                    self.source_info(test.span),
+                    actual_ref_place,
+                    Rvalue::Ref(tcx.lifetimes.re_erased, BorrowKind::Shared, place),
+                );
+
+                // Compare using `<T as PartialEq>::eq` where `T` is the array or slice type.
+                self.non_scalar_compare(
+                    block,
+                    success_block,
+                    fail_block,
+                    source_info,
+                    aggregate_ty,
+                    expected_operand,
+                    Operand::Copy(actual_ref_place),
+                );
+            }
+
             TestKind::ScalarEq { value } => {
                 let tcx = self.tcx;
                 let success_block = target_block(TestBranch::Success);
@@ -404,19 +445,22 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         );
     }
 
-    /// Compare two values of type `&str` using `<str as std::cmp::PartialEq>::eq`.
-    fn string_compare(
+    /// Compare two reference values using `<T as PartialEq>::eq`.
+    ///
+    /// `compared_ty` is the *inner* type (e.g. `str`, `[u8; 64]`);
+    /// `expect` and `val` must already be references to that type.
+    fn non_scalar_compare(
         &mut self,
         block: BasicBlock,
         success_block: BasicBlock,
         fail_block: BasicBlock,
         source_info: SourceInfo,
+        compared_ty: Ty<'tcx>,
         expect: Operand<'tcx>,
         val: Operand<'tcx>,
     ) {
-        let str_ty = self.tcx.types.str_;
         let eq_def_id = self.tcx.require_lang_item(LangItem::PartialEq, source_info.span);
-        let method = trait_method(self.tcx, eq_def_id, sym::eq, [str_ty, str_ty]);
+        let method = trait_method(self.tcx, eq_def_id, sym::eq, [compared_ty, compared_ty]);
 
         let bool_ty = self.tcx.types.bool;
         let eq_result = self.temp(bool_ty, source_info.span);

compiler/rustc_span/src/symbol.rs

@@ -649,6 +649,7 @@ symbols! {
         const_block_items,
         const_c_variadic,
         const_closures,
+        const_cmp,
         const_compare_raw_pointers,
         const_constructor,
         const_continue,

tests/mir-opt/building/match/aggregate_array_eq.array_match.built.after.panic-abort.mir

@@ -0,0 +1,51 @@
+// MIR for `array_match` after built
+
+fn array_match(_1: [u8; 4]) -> bool {
+    debug x => _1;
+    let mut _0: bool;
+    let mut _2: &[u8; 4];
+    let mut _3: bool;
+    scope 1 {
+    }
+
+    bb0: {
+        PlaceMention(_1);
+        _2 = &_1;
+        _3 = <[u8; 4] as PartialEq>::eq(copy _2, const &*b"\x01\x02\x03\x04") -> [return: bb4, unwind: bb8];
+    }
+
+    bb1: {
+        _0 = const false;
+        goto -> bb7;
+    }
+
+    bb2: {
+        falseEdge -> [real: bb6, imaginary: bb1];
+    }
+
+    bb3: {
+        goto -> bb1;
+    }
+
+    bb4: {
+        switchInt(move _3) -> [0: bb1, otherwise: bb2];
+    }
+
+    bb5: {
+        FakeRead(ForMatchedPlace(None), _1);
+        unreachable;
+    }
+
+    bb6: {
+        _0 = const true;
+        goto -> bb7;
+    }
+
+    bb7: {
+        return;
+    }
+
+    bb8 (cleanup): {
+        resume;
+    }
+}