Codegen struggles to infer information about a value computed before branching on related info

[Apersoma]

Explanation:

/// basically what the function looks like in core
fn strict_sub(x: u128, y: u128) -> u128 {
      let (diff, over) = x.overflowing_sub(y);
      if over { panic!() } else { diff }
}
/// this is better for some reason
fn strict_sub_1liner(x: u128, y: u128) -> u128 {
      if x <= y { panic!() } else { x - y }
}

/// Performs the Lucas Lehmer Test to see if `2^exp - 1` is prime.
/// Panics if `2^exp` does not fit in a u64.
pub fn lucas_lehmer_u64(exp: u8) -> bool {
    let n64: u64 = 1u64.strict_shl(exp as u32) - 1;
    let n: u128 = n64 as u128;
    
    let mut k: u128 = 4;
    
    // swapping out `strict_sub` here for `strict_sub_1line` allows the compiler to
    // figure out that the `strict_add` later won't panic.
    // For context, that wont panic because `k` is always less then `n` at the end and
    // `n` can fit in a `u64` so if you square it it won't be too large to then add a `n_m2`
    // if that also fits in a `u64` (basically the idea behind the `carrying_mul`  and 
    // `carrying_mul_add` function). 
    // If `n_m2` were to underflow, it won't fit anymore and the `strict_add` can overflow.
    let n_m2 = strict_sub(n, 2);
    
    let mut i = 2;
    
    while i < exp {
        k = (k*k).strict_add(n_m2);
        unsafe {
            k = (k & n).unchecked_add(k >> exp);
            k = (k & n).unchecked_add(k >> exp);
        }
        k = core::hint::select_unpredictable(k == n, 0, k);
        // above 3 statements perform `k mod n` so `k` is less then `n`
        unsafe { core::hint::assert_unchecked(k < n) }
        i += 1;
    }
    k == 0
}

godbolt link: https://godbolt.org/z/h4Gcn4Pvz

[Apersoma]

@rustbot label: +C-codegen

[rustbot]

Error: Unknown labels: C-codegen

Please file an issue on GitHub at triagebot if there's a problem with this bot, or reach out on #triagebot on Zulip.

[Apersoma]

@rustbot label: +A-codegen

[Apersoma]

@rustbot label: +I-Slow +C-Optimization

[Apersoma]

This can be bandaided by changing the implementation in core to match closer to what I have but that's obviously not the best solution.

[Apersoma]

Now that I think about the results of some of the testing I did, I think it can do the optimization when x-y is performed conditionally only when x >= y and can't do the optimization if it is not computed conditionally.

[Apersoma]

this probably happens with the other operations but I haven't checked

[Apersoma]

Was trying to figure out why rust wasn't figuring out that uv wasn't zero in the a sippet of similar the following and realized its the same issue.

pub fn foo(u: u32, v: u32) -> Option<NonZero<u32>>{
    let uv = u | v;
    if u == 0 || v == 0 {
        uv
    } else {
        1 << uv.trailing_zeros()
    }
}