EarlyOtherwiseBranch transformation is incorrect

[tmiasko]

For details see #77163 (comment), opening an issue since PR was closed:

pub enum E<'a> {
    Empty,
    Some(&'a E<'a>),
}

fn f(e: &E) -> u32 {
   if let E::Some(E::Some(_)) = e { 1 } else { 2 }
}

fn main() {
   println!("{:?}", f(&E::Empty));
}

rustc b.rs -Zmir-opt-level=2 && ./b
"./b" terminated by signal SIGSEGV (Address boundary error)

[tmiasko]

While this one particular test case was fixed, the transformation still suffers from issues hinted at earlier comment:

1. It assumes that terminator always switches on a local that was assigned last in the block. The assumption is false.
2. It moves a read of a discriminant from the second block into first without ensuring that it holds the same value still.
3. In some executions it introduces a read that would not have happen otherwise and could have an undefined behaviour.
4. In some executions it skips execution of statements from one the blocks without checking if they are side-effect free.

cc @oli-obk

[wecing]

@tmiasko ,

1. It assumes that terminator always switches on a local that was assigned last in the block. The assumption is false.

EarlyOtherwiseBranch only applies to basic blocks of which the last statement is a discriminant (and the terminator is a switch):

    fn find_switch_discriminant_info(
        &self,
        bb: &BasicBlockData<'tcx>,
        switch: &Terminator<'tcx>,
    ) -> Option<SwitchDiscriminantInfo<'tcx>> {
        ...
                let place_of_adt_discr_read = match bb.statements.last()?.kind {
                    StatementKind::Assign(box (_, Rvalue::Discriminant(adt_place))) => {
                        Some(adt_place)
                    }
                    _ => None,
                }?;

And,

2. It moves a read of a discriminant from the second block into first without ensuring that it holds the same value still.
4. In some executions it skips execution of statements from one the blocks without checking if they are side-effect free.

We could require the discriminant read in the second basic block to be the only statement within the block. Not sure by how much would the change affect optimization hit rate, but at least the example in #68867 would still be optimized.

---

What makes me really concerned is this one:

3. In some executions it introduces a read that would not have happen otherwise and could have an undefined behaviour.

When would the first discriminant be safe but not the second? The example you constructed is a legit one: the second read assumes the value of the first read and does a downcast; that is the case my patch catches.

The only other possible invalid case I could think of, is that when the second discriminant call's argument is an invalid pointer. Probably something like:

fn f(x: &Option<i32>, y: &Option<i32>) -> i32 {
    if let &Some(_) = x {
        if let &Some(_) = y {
            return 1
        }
    }
    42
}

f(&None, /* invalid ref obtained through unsafe? */)

But today the MIR that FE emits dodges the problem by accident:

fn f(_1: &Option<i32>, _2: &Option<i32>) -> i32 {
    debug x => _1;                       // in scope 0 at t.rs:22:12: 22:13
    debug y => _2;                       // in scope 0 at t.rs:22:19: 22:20
    let mut _0: i32;                     // return place in scope 0 at t.rs:22:29: 22:32
    let _3: ();                          // in scope 0 at t.rs:23:5: 27:6
    let mut _4: isize;                   // in scope 0 at t.rs:23:13: 23:20
    let mut _5: isize;                   // in scope 0 at t.rs:24:17: 24:24
    let mut _6: !;                       // in scope 0 at t.rs:25:13: 25:21

    bb0: {
        StorageLive(_3);                 // scope 0 at t.rs:23:5: 27:6
        _4 = discriminant((*_1));        // scope 0 at t.rs:23:13: 23:20
        switchInt(move _4) -> [0_isize: bb2, otherwise: bb1]; // scope 0 at t.rs:23:13: 23:20
    }

    bb1: {
        _3 = const ();                   // scope 0 at t.rs:27:6: 27:6
        goto -> bb5;                     // scope 0 at t.rs:23:5: 27:6
    }

    bb2: {
        _5 = discriminant((*_2));        // scope 0 at t.rs:24:17: 24:24
        switchInt(move _5) -> [0_isize: bb4, otherwise: bb3]; // scope 0 at t.rs:24:17: 24:24
    }

    bb3: {
        _3 = const ();                   // scope 0 at t.rs:26:10: 26:10
        goto -> bb5;                     // scope 0 at t.rs:24:9: 26:10
    }

    bb4: {
        _0 = const 1_i32;                // scope 0 at t.rs:25:20: 25:21
        StorageDead(_3);                 // scope 0 at t.rs:27:5: 27:6
        goto -> bb6;                     // scope 0 at t.rs:29:2: 29:2
    }

    bb5: {
        StorageDead(_3);                 // scope 0 at t.rs:27:5: 27:6
        _0 = const 42_i32;               // scope 0 at t.rs:28:5: 28:7
        goto -> bb6;                     // scope 0 at t.rs:29:2: 29:2
    }

    bb6: {
        return;                          // scope 0 at t.rs:29:2: 29:2
    }
}

Because these two switch instructions have different otherwise blocks.

So, my guess is, even though EarlyOtherwiseBranch is unsound, today the compiler FE is still unable to produce MIR that could trigger that.

[tmiasko]

even though EarlyOtherwiseBranch is unsound, today the compiler FE is still unable to produce MIR that could trigger that.

If one considers shape of MIR immediately after it has been built, that is probably true. One thing to keep mind: this transformation runs on MIR that have been transformed already. In fact, the whole optimization pipeline might have run an arbitrary number of times as a result of inlining. Transformations shouldn't rely on any invariants that aren't preserved by all transformations.

EarlyOtherwiseBranch only applies to basic blocks of which the last statement is a discriminant

The place_of_adt_discr_read might be unrelated to a local used in switch, since this is not checked anywhere. Last time out of curiosity I added an assertion checking that, enabled the transform by default and bootstraped rustc, the assertion did fail.

[wecing]

The place_of_adt_discr_read might be unrelated to a local used in switch, since this is not checked anywhere.

Ah! I misunderstood your point #1, but now I get it. You are right.

I think we could resolve your #1, #2, and #4 by adding the following restrictions:

1. The discriminant read in the second block must be the only statement in there;
2. The switch must use the value of the discriminant read.

---

But I found that the same effect of EarlyOtherwiseBranch could be achieved through a much easier change; we could provide a hint for LLVM optimizer to compare the discriminant values before entering the switch, e.g.:

enum E {
    A(i32),
    B(i32),
    C(i32),
    D(i32),
}

#[no_mangle]
fn add1_match(x: &E, y: &E) -> i32 {
    if (mem::discriminant(x) != mem::discriminant(y)) {
        return 99992;
    }

    match (x,y) {
        (&E::A(x), &E::A(y)) => x+y,
        (&E::B(x), &E::B(y)) => x+y,
        (&E::C(x), &E::C(y)) => x+y,
        (&E::D(x), &E::D(y)) => x+y,
        _ => 99992,
    }
}

With this hint, rustc --emit-asm -C opt-level=3 t.rs (unfortunately, opt-level=2 isn't enough) produces:

        .section        .text.add1_match,"ax",@progbits
        .globl  add1_match
        .p2align        4, 0x90
        .type   add1_match,@function
add1_match:
        .cfi_startproc
        movl    (%rdi), %ecx
        movl    $99992, %eax
        cmpl    (%rsi), %ecx
        jne     .LBB7_3
        movl    %ecx, %eax
        leaq    .LJTI7_0(%rip), %rcx
        movslq  (%rcx,%rax,4), %rax
        addq    %rcx, %rax
        jmpq    *%rax
.LBB7_2:
        movl    4(%rsi), %eax
        addl    4(%rdi), %eax
.LBB7_3:
        retq
.Lfunc_end7:
        .size   add1_match, .Lfunc_end7-add1_match
        .cfi_endproc
        .section        .rodata.add1_match,"a",@progbits
        .p2align        2
.LJTI7_0:
        .long   .LBB7_2-.LJTI7_0
        .long   .LBB7_2-.LJTI7_0
        .long   .LBB7_2-.LJTI7_0
        .long   .LBB7_2-.LJTI7_0

IMO this is much more elegant than EarlyOtherwiseBranch; it should also be pretty easy to implement, so I will give it a try.