Unnecesary discriminant checks in PartialEq/PartialOrd style code on enums

[Kmeakin]

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>

[Kmeakin]

@rustbot label I-slow

[Kmeakin]

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

[Kmeakin]

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)

[clubby789]

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

[Kmeakin]

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),
}

[clubby789]

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

[DianQK]

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