Unnecesary discriminant checks in PartialEq/PartialOrd style code on enums

I tried this code:
compiler explorer

pub enum Enum {
    A(u32),
    B(u32),
    C(u32),
}

#[no_mangle]
pub fn f1(lhs: &Enum, rhs: &Enum) -> bool {
    match (lhs, rhs) {
        (Enum::A(lhs), Enum::A(rhs)) => lhs == rhs,
        (Enum::B(lhs), Enum::B(rhs)) => lhs == rhs,
        (Enum::C(lhs), Enum::C(rhs)) => lhs == rhs,
        _ => false,
    }
}

#[no_mangle]
pub fn f2(lhs: &Enum, rhs: &Enum) -> bool {
    if std::mem::discriminant(lhs) != std::mem::discriminant(rhs) {
        return false;
    }

    match (lhs, rhs) {
        (Enum::A(lhs), Enum::A(rhs)) => lhs == rhs,
        (Enum::B(lhs), Enum::B(rhs)) => lhs == rhs,
        (Enum::C(lhs), Enum::C(rhs)) => lhs == rhs,
        _ => false,
    }
}

#[no_mangle]
pub fn f3(lhs: &Enum, rhs: &Enum) -> bool {
    if std::mem::discriminant(lhs) != std::mem::discriminant(rhs) {
        return false;
    }

    match (lhs, rhs) {
        (Enum::A(lhs), Enum::A(rhs)) => lhs == rhs,
        (Enum::B(lhs), Enum::B(rhs)) => lhs == rhs,
        (Enum::C(lhs), Enum::C(rhs)) => lhs == rhs,
        _ => unsafe { std::hint::unreachable_unchecked() },
    }
}

Code like this is often used for implementations of PartialEq or PartialOrd on enums, so it is important that it be optimized well.

I expected to see this happen:
The functions f1 and f2 should produce equivalent machine code to f3.

Instead, this happened:
f1 and f2 do not produce equivalent machine code.
f1 is missing the early return for when the discriminants do not match.
f2 unnecessarily examines the discriminant of lhs twice!
LLVM should be able to transform f1 into f2, and f2 into f3

Meta

rustc --version --verbose:

rustc 1.76.0-nightly (a96d57bdb 2023-12-15)
binary: rustc
commit-hash: a96d57bdb6d2bb6d233d7d5aaefc2995ab99be01
commit-date: 2023-12-15
host: x86_64-unknown-linux-gnu
release: 1.76.0-nightly
LLVM version: 17.0.6

Backtrace

<backtrace>

@rustbot

@rustbot label I-slow

Interestingly, when Enum is a c-like enum, f1 produces the same code as f3, but f2 doesnt: https://godbolt.org/z/81bGdq6xs

Adding a zero-sized field (like ()) even prevents f1 being optimized into f3 (https://godbolt.org/z/3dqzTfofa), but an empty field list doesn't (https://godbolt.org/z/bTK9qa9z1)

~~Looks similar to #113506, which was fixed upstream.~~
Ah, a little different
Throwing f1 into upstream opt improves codegen, but still not as good as the optimal case in f3

The PartialEq derive uses essentially the same code as f3 because of this issue, so it would be nice to have it work here

Looks similar to #113506, which was fixed upstream

The improvement I'm looking for here is removing unnecessary repeated inspections of enum discriminants, not merging duplicate RHS per se. This improvement would still be useful even if each enum variant has different data in it:

compiler explorer

pub enum Enum {
    A(u32),
    B(u64),
    C(f32),
    D(f64),
}

Here's some more similar strange behaviour:

#[no_mangle]
pub fn f1(lhs: &Enum, rhs: &Enum) -> bool {
    match (lhs, rhs) {
        (Enum::A(lhs), Enum::A(rhs)) => true,
        (Enum::B(lhs), Enum::B(rhs)) => true,
        (Enum::C(lhs), Enum::C(rhs)) => true,
        _ => false,
    }
}

demonstrates the repeated inspections, but if we remove the unused lhs/rhs bindings (or have bindings in only one branch), we get optimal codegen:
https://godbolt.org/z/fPMx5bYz6

Seems like MIR optimisations manage to improve it to a point where LLVM can optimise it better

@rustbot

Update: In fact, we just need to re-enable the early otherwise branch optimization.

Hmm, this requires multiple optimizations to work together.

The first step is to hoist the load instruction, which is already in LLVM 18, as we'll see soon. See Reland "[SimplifyCFG] Hoist common instructions on switch" llvm/llvm-project#67077 and https://llvm.godbolt.org/z/ozfdG1qxb.
The most important thing is to merge the same BB, alive2: https://alive2.llvm.org/ce/z/V39zb2.
After that, we need to hoist the icmp instruction, as Transforms match into an assignment statement #120614 does. Alive2: https://alive2.llvm.org/ce/z/s1e6ez.

As we merge the two i32, we'll end up with better instructions. See https://llvm.godbolt.org/z/EYG46EaxY.

mov     rax, qword ptr [rsi]
cmp     rax, qword ptr [rdi]
sete    al
ret

It may be that part of the optimization fits in MIR, I'll trade off.

@rustbot claim

Add test case for rust-lang#119014

Auto merge of rust-lang#121397 - DianQK:early_otherwise_branch_sound,…

Add test case for rust-lang#119014

Auto merge of rust-lang#121397 - DianQK:early_otherwise_branch_sound,…

Add test case for rust-lang#119014

Auto merge of rust-lang#122387 - DianQK:re-enable-early-otherwise-bra…

Add test case for rust-lang#119014

Auto merge of rust-lang#122387 - DianQK:re-enable-early-otherwise-bra…

Add test case for rust-lang#119014

Auto merge of rust-lang#122387 - DianQK:re-enable-early-otherwise-bra…

Navigation Menu

Repository navigation

Unnecesary discriminant checks in `PartialEq`/`PartialOrd` style code on enums #119014

Meta

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Unnecesary discriminant checks in PartialEq/PartialOrd style code on enums #119014

Description

Meta

Activity

Kmeakin commented on Dec 16, 2023

Kmeakin commented on Dec 16, 2023

Kmeakin commented on Dec 16, 2023

clubby789 commented on Dec 16, 2023

Kmeakin commented on Dec 17, 2023

clubby789 commented on Jan 10, 2024

dianqk commented on Feb 7, 2024

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Unnecesary discriminant checks in `PartialEq`/`PartialOrd` style code on enums #119014