Missed optimization: returning a niche-optimized `Option<&T>` leaves a redundant null check after inlining

[wyfo]

The following code:

use core:hint::black_box;

#[inline(always)]
fn foo(cond: bool, ptr: *const usize) -> Option<&'static usize> {
    if cond {
        Some(unsafe { &*ptr })
    } else {
        None
    }
}

#[unsafe(no_mangle)]
fn bar(cond: bool, ptr: *const usize) {
    if let Some(r) = foo(cond, ptr) {
        black_box(r);
    }
}

is compiled with -C opt-level=3 -Zinline-mir=no into the following assembly:

bar:
        push    rax
        xor     cl, 1
        test    rdx, rdx
        sete    al
        or      al, cl
        jne     .LBB0_2
        mov     qword ptr [rsp], rdx
        mov     rax, rsp
.LBB0_2:
        pop     rax
        ret

In the assembly, there are two checks merged in one, the first is the boolean condition, and the second one is whether ptr is null. The second check is redundant and should ideally not be compiled.

https://godbolt.org/z/esxPahrzE

My investigations so far

The issue doesn't happen when foo is inlined in the MIR, that's why MIR inlining is disabled. It can be reproduced without the compiler flag by adding a bit more code to foo, so it is not inlined in the MIR but by LLVM. That's what happened my original reproducing code.

Because totally disabling has some unwanted side effects, in my next godbolt link, I add 17 black_box(cond) at the beginning of foo to make it not MIR-inlinable: https://godbolt.org/z/Gxqhb3cqM

The problem seems to come from the niche optimization of Option<&T>. In fact, I've tried many things in godbolt, like replacing the return type by Option<Cell<&T>> or (MaybeUninit<T>, bool), both removing the niche optimization, and in these cases, the redundant check was not compiled.

Making LLVM understand that the pointer is not null also removes the redundant check. For example, replacing &*ptr by ptr.as_ref().unwrap_unchecked(), or wrapping it into a black_box work.

One strange thing is that adding an assert_unchecked(!ptr.is_null()) has no effect. However, if ptr becomes the field of a struct, assert_unchecked(!self.ptr.is_null()) correctly tag the self.ptr as !nonnull and eliminate the redundant check. My own code is in this case (the pointer is a field), so I use the assert_unchecked solution for now.

[cezarbbb]

I think it is that LLVM SROA lost the !nonnull attribute attached to load when promoting &*ptr to a raw ptr, causing select cond, ptr, null to be unable to be simplified to !cond.
In my opinion, there's not much that can be done in rustc; it would likely require modifications to LLVM to resolve this issue.

[nikic]

The !nonnull metadata is on a load under the != null condition, so it provides no information.

[cezarbbb]

The !nonnull metadata is on a load under the != null condition, so it provides no information.

You are right, !nonnull isn't dropped, it's on the load in the non-null arm, so it's conditional on the null test itself and can't help eliminate it.