Missed optimization: re-checking enum variants after `std::mem::replace` is not eliminated (somtimes)

[oxalica]

Minified example: (godbolt)

use std::task::{Poll, Waker};

pub enum State<T> {
    Inactive,
    Active(Waker),
    Signalled(T),
}

#[unsafe(no_mangle)]
pub fn poll_state(st: &mut State<String>, w: &Waker) -> Poll<String> { // a
    match st {
        State::Signalled(_) => {
            // Just checked the variant, take the value out.
            let State::Signalled(v) = std::mem::replace(st, State::Inactive) else {
                unreachable!() // This panic should be eliminated.
            };
            Poll::Ready(v)
        }
        _ => {
            *st = State::Active(w.clone()); // b
            Poll::Pending
        }
    }
}

The optimization is fragile. If we remove the state assignment on the second branch (b), or change function signature to operate on State<u8>, then the unreachable panic in the first branch will be correctly eliminated.

[Berrysoft]

Seems that the panic won't be eliminated if assignment (b) is removed.

[ashivaram23]

I've come across something similar. I think MIR optimizations don't propagate known enum discriminants across blocks so it's left to LLVM, which might be unable to do more with it. In the case I've seen, LLVM's inability to optimize further could be traced down to how it handles the enum's niche optimization, but it could be something totally different here.

A MIR pass that propagates enum discriminants (through dataflow analysis or just by checking the direct predecessors of each block) would probably fix all these cases.

[oxalica]

known enum discriminants

So is it related to that known value ranges are not propagated through some boundaries (functions, for example)? I also encountered an issue about the discriminant of Option<NonZero<usize>> in #49572 (comment)? Not sure if they are related.