Regression: iter_mut().nth().next() retains unwrap panic path and extra branch vs split_at_mut in optimized code (Rust 1.82+)

[camc314]

Code

I tried this code:

#![crate_type = "lib"]
use std::hint::unreachable_unchecked;
use std::collections::HashMap;
type Map = HashMap<String, Vec<u32>>;
pub struct Stack {
    stack: Vec<Map>,
    depth: usize,
}
// ============================================================================
// BEFORE: Using iterator with nth() and next()
// - Generates unwrap panic path even with assert_unchecked hints
// - More instructions, conditional branch
// ============================================================================
#[no_mangle]
pub fn use_iter(s: &mut Stack) -> (&mut Map, &mut Map) {
    // Assert invariants
    if s.depth == 0 {
        unsafe { unreachable_unchecked() };
    }
    if s.stack.len() <= s.depth {
        unsafe { unreachable_unchecked() };
    }
    
    let mut iter = s.stack.iter_mut();
    let parent = iter.nth(s.depth - 1).unwrap();
    let current = iter.next().unwrap();
    (current, parent)
}
// ============================================================================
// AFTER: Using split_at_mut
// - No panic path, fully inlined
// - Fewer instructions, no branches
// ============================================================================
#[no_mangle]
pub fn use_split(s: &mut Stack) -> (&mut Map, &mut Map) {
    // Assert invariants
    if s.depth == 0 {
        unsafe { unreachable_unchecked() };
    }
    if s.stack.len() <= s.depth {
        unsafe { unreachable_unchecked() };
    }
    
    let (head, tail) = s.stack.split_at_mut(s.depth);
    let parent = &mut head[s.depth - 1];
    let current = &mut tail[0];
    (current, parent)
}

I expected to see this happen:

use_iter:
        mov     rcx, qword ptr [rdi + 8]
        mov     rax, qword ptr [rdi + 24]
        lea     rdx, [rax + 2*rax]
        shl     rdx, 4
        lea     rax, [rcx + rdx]
        add     rdx, rcx
        add     rdx, -48
        ret

use_split:
        mov     rcx, qword ptr [rdi + 8]
        mov     rax, qword ptr [rdi + 24]
        lea     rdx, [rax + 2*rax]
        shl     rdx, 4
        lea     rax, [rcx + rdx]
        add     rdx, rcx
        add     rdx, -48
        ret

Instead, this happened:

use_iter:
        mov     rax, qword ptr [rdi + 24]
        movabs  rcx, 1152921504606846975
        and     rcx, rax
        cmp     rcx, qword ptr [rdi + 16]
        je      .LBB0_2
        lea     rcx, [rax + 2*rax]
        shl     rcx, 4
        mov     rdx, qword ptr [rdi + 8]
        lea     rax, [rdx + rcx]
        add     rdx, rcx
        add     rdx, -48
        ret
.LBB0_2:
        push    rax
        lea     rdi, [rip + .Lanon.621aa17d7ebd209918a0c2e574629b83.1]
        call    qword ptr [rip + core::option::unwrap_failed::he1a8284b5a1e2496@GOTPCREL]

use_split:
        mov     rcx, qword ptr [rdi + 8]
        mov     rax, qword ptr [rdi + 24]
        lea     rdx, [rax + 2*rax]
        shl     rdx, 4
        lea     rax, [rcx + rdx]
        add     rdx, rcx
        add     rdx, -48
        ret

.Lanon.621aa17d7ebd209918a0c2e574629b83.0:
        .asciz  "/app/example.rs"

.Lanon.621aa17d7ebd209918a0c2e574629b83.1:
        .quad   .Lanon.621aa17d7ebd209918a0c2e574629b83.0
        .asciz  "\017\000\000\000\000\000\000\000 \000\000\000\037\000\000"

Version it worked on

It most recently worked on: Rust 1.81
It broke in 1.82

[camc314]

take this with a grain of salt, but I got claude got to bisect, and it has decded that #126991 appears to cause this regression

Again, take this with a grain of salt, but I made a repro here: https://github.com/camc314/rust-issue-150235, and the LLVM is outputted by rust appears to be the same in 1.81, and 1.82, pointing to the llvm v19 upgrade causing this issue

[moxian]

re 👎 : I tried bisecting without claude (but with the regular cargo-bisect-rustc), and it does point at LLVM 19 update #127513 (or indeed at #126991, but that doesn't sound right since the issue still repros with -Zmir-opt-level=0 and since it has been demonstrated above that it can be isolated to just LLVM opt alone). More specifically

Regression in nightly-2024-08-01. PRs in range:

 - #128083 (Bump bootstrap compiler to new beta) by Mark-Simulacrum
 - #128413 (Rollup of 7 pull requests) by matthiaskrgr
   - #128357 (Detect non-lifetime binder params shadowing item params)
   - #128367 (CI: rfl: build the generated doctests and documentation)
   - #128376 (Mark `Parser::eat`/`check` methods as `#[must_use]`)
   - #128379 (the output in stderr expects panic-unwind)
   - #128380 (make `///` doc comments compatible with naked functions)
   - #128382 (cargo-miri: better error when we seem to run inside bootstrap but something is wrong)
   - #128398 (tidy: Fix quote in error message)
 - #128075 (Migrate `rlib-format-packed-bundled-libs-2`, `native-link-modifier-whole-archive` and `no-builtins-attribute` `run-make` tests to rmake) by Oneirical
 - #128420 (Update cargo) by weihanglo
 - #127513 (Update to LLVM 19) by nikic
 - #128435 (Rollup of 9 pull requests) by matthiaskrgr
   - #126454 (bump-stage0: use IndexMap for determinism)
   - #127681 (derive(SmartPointer): rewrite bounds in where and generic bounds)
   - #127830 (When an archive fails to build, print the path)
   - #128151 (Structured suggestion for `extern crate foo` when `foo` isn't resolved in import)
   - #128387 (More detailed note to deprecate ONCE_INIT)
   - #128388 (Match LLVM ABI in `extern "C"` functions for `f128` on Windows)
   - #128402 (Attribute checking simplifications)
   - #128412 (Remove `crate_level_only` from `ELIDED_LIFETIMES_IN_PATHS`)
   - #128430 (Use a separate pattern type for `rustc_pattern_analysis` diagnostics )
 - #126991 (Accelerate GVN a little) by cjgillot

godbolt link for the llvm side of things (far from minimal): https://godbolt.org/z/K7Ez4s8Ej
@rustbot label: +A-LLVM -C-bug +C-optimization +I-slow -regression-untriaged +regression-from-stable-to-stable

[scottmcm]

Note that if you want two contiguous things from a slice, I highly suggest you do something more like

head[s.depth - 1..].first_chunk::<2>().unwrap()

instead. That's far more direct about what's happening, and thus much more likely to go well in the optimizer.

I don't know that rust is likely to really do much about this; if you're using unreachable_unchecked like that to get things later to optimize out, those kinds of optimizations are almost always LLVM-side. https://github.com/llvm/llvm-project/issues/new

[apiraino]

Assigning priority (discussion on Zulip).

@rustbot label -I-prioritize +P-low