Vec::reserve(n) followed by n calls to push should only check capacity once

[jhorstmann]

I tried this code in compiler explorer (https://rust.godbolt.org/z/W35vT4osM):

pub fn push1(v: &mut Vec<u32>, x: u32) {
    v.reserve(1);
    v.push(x);
}

pub fn push2(v: &mut Vec<u32>, x: u32) {
    v.reserve(2);
    v.push(x);
    v.push(x);
}

I expected the generated assembly to only check the capacity once, and reallocate if needed, and then write the value. Instead, for each push there is another check of remaining capacity and potential call to reserve_for_push.

Adding explicit assume calls leads to much shorter code with only a single call to do_reserve_and_handle.

pub fn push2_assume(v: &mut Vec<u32>, x: u32) {
    v.reserve(2);

    unsafe {
        core::intrinsics::assume(v.len() < v.capacity());
    }
    v.push(x);

    unsafe {
        core::intrinsics::assume(v.len() < v.capacity());
    }
    v.push(x);
}

This seems similar to #82801 but that issue seems specific to the initial vector capacity and code for that case looks ok now.

Meta

rustc --version --verbose:

rustc 1.67.0-nightly (e0098a5cc 2022-11-29)
binary: rustc
commit-hash: e0098a5cc3a87d857e597af824d0ce1ed1ad85e0
commit-date: 2022-11-29
host: x86_64-unknown-linux-gnu
release: 1.67.0-nightly
LLVM version: 15.0.4
Compiler returned: 0

[SUPERCILEX]

Just ran into this too. Here's the generalization:

pub fn insert(v: &mut Vec<usize>, n: usize) {
    v.reserve(n);
    for i in 0..n {
        v.push(i);
    }
}

pub fn create(n: usize) -> Vec<usize> {
    let mut v = Vec::with_capacity(n);
    for i in 0..n {
        v.push(i);
    }
    v
}

[SUPERCILEX]

Fixed with this:

pub fn insert(v: &mut Vec<usize>, n: usize) {
    v.reserve(n);
    for i in 0..n {
        if v.len() >= v.capacity() {
            unsafe { std::hint::unreachable_unchecked() }
        }
        v.push(i);
    }
}

pub fn create(n: usize) -> Vec<usize> {
    let mut v = Vec::with_capacity(n);
    for i in 0..n {
        if v.len() >= v.capacity() {
            unsafe { std::hint::unreachable_unchecked() }
        }
        v.push(i);
    }
    v
}

So should be handled by #82801

[c410-f3r]

An example that eliminates branches: https://godbolt.org/z/TTEv73o44

Change v.allocate(4); to any number less than 4 to see the compiler insert guards accordingly.

For whatever reason, the trick is in the use of unchecked_add in allocate as hinted at #121300 (comment).

[cynecx]

I'd like to leave some more observations I've noticed while playing around with this: https://gist.github.com/cynecx/24e3a4ae8756fdf9257a6cbf50005e08

tl;dr: A single assume in Vec's reserve method should've been enough to actually optimize out all further bound checks. But the big caveat here is that llvm isn't actually able to apply these optimizations with its default pipeline of passes. The constraint elimination pass must be run several times interleaved with other optimizations. This could be a limitation in the constraint elimination pass itself