Improve bounds check for function that always return in-bounds index
Opened this issue · 7 comments
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
The solution is to add refinement typing in Rust:
https://antelang.org/docs/language/#refinement-types
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.
Well, refinement typing is a more general solution, you can use for much more than this case of binary_search :)
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.
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.