Improve bounds check for function that always return in-bounds index

[Pzixel]

Consider following function:

pub fn foo(mut a: Vec<u64>) -> Vec<u64> {
    let val = 15;
    if let Err(idx) = a.binary_search(&val) {
        a.insert(idx, val);
    }
    a
}

On current rust 1.61.0 you can expect following output:

...
.LBB3_12:
        mov     rdi, rbx
        mov     rsi, r12
        call    qword ptr [rip + alloc::vec::Vec<T,A>::insert::assert_failed@GOTPCREL]
        ud2
...

It would be nice to have some attribute or other way of saying that binary search is always returning a valid insert index and bounds check should be eliminated. You can get an expected output in current rustc versions via:

pub fn foo(mut a: Vec<u64>) -> Vec<u64> {
    let val = 15;
    if let Err(idx) = a.binary_search(&val) {
        if idx > a.len() {
            unsafe { std::hint::unreachable_unchecked() };
        }
        a.insert(idx, val);
    }
    a
}

Since performing a search and then inserting is one of main binary_search use cases it might be worthwhile to implement such an optimization.
See godbolt link for a whole example

[leonardo-m]

The solution is to add refinement typing in Rust:
https://antelang.org/docs/language/#refinement-types

[Pzixel]

Well while refinement types are great I think rust may just go with what is already supported. For example if we replace Err with Ok then it can see that index is no bigger than a.len() and removes the assert_failed. It just about implementing a bit more sophisticated heuristics.

[leonardo-m]

Well, refinement typing is a more general solution, you can use for much more than this case of binary_search :)

[the8472]

Have you benchmarked this? Inserting in the middle of a vec incurs O(n) cost for the shifting, that likely dwarves the cost of the bounds check so the optimization may not be worth it.

[Pzixel]

Well right, but inserting in the end of a vec incurs no costs for the shifting. I can also imagine this might come with some unexpected benefits in other functions and/or scenarios.

[scottmcm]

I've got a PR open that might improve this: #102535

(It'll definitely help for some slice methods, but I'm unsure about Vec ones, since the bounds checking is more complicated when Deref is involved.)

[clubby789]

@rustbot label +A-llvm +I-slow
Looks like the call to insert::assert_failed is still there