The compiler looses vector capacity information

[MSxDOS]

In this example:

pub fn test() -> Vec<u32> {
    let mut vec = Vec::with_capacity(3);
    let mut i = 0;
    while i < 3 {
        vec.push(32);
        i += 1;
    }
    vec
}

the compiler produces sub-optimal assembly.
https://rust.godbolt.org/z/1T3Knodn8

[nikic]

Don't know what role the MIR opts play here, but on the LLVM side this looks like a phase ordering issue. The final IR has:

  %10 = bitcast i64* %4 to <2 x i64>*
  store <2 x i64> <i64 3, i64 2>, <2 x i64>* %10, align 8
  %_6.idx.val.i.2 = load i64, i64* %4, align 8
  %_3.i.2 = icmp eq i64 %_6.idx.val.i.2, 2

The load can be load store forwarded to 3 and everything else folds from there. A sequence of -gvn -simplifycfg -sroa does it.

[nikic]

Took a while to root cause this, but should be fixed by llvm/llvm-project@3fedaf2. The issue was that GVN did not fully fold the function because available loads were materialized suboptimally, leading to unnecessary poisoning of MDA cache entries.

https://reviews.llvm.org/D98114 would also fix fix this case, though it's arguably more of a workaround.

[nikic]

I'm not completely sure what I'm supposed to be looking for here, but I think this has not been fixed by the LLVM 13 upgrade. It looks like the IR substantially changed in 1.53 and the goal now is presumably to eliminate the branch that contains the do_reserve_and_handle call.

[MSxDOS]

It wasn't fixed by LLVM 13, tested locally as the compiler on godbolt is outdated.

It fails to see that self.len == self.buf.capacity() inside push can't be true in the first place.
Adding

if vec.len() == vec.capacity() {
    unsafe { std::hint::unreachable_unchecked() }
}

in the beginning of the loop removes the reserve branch entirely.

[nikic]

Principally, a combination of loop load PRE and SCCP should be able to handle this. However, loop load PRE will currently fail for two reasons: First, the clobbering instruction (do_reserve_and_handle) is an invoke terminator, which would need edge splitting support in loop load PRE. Second, we don't have a guarantee that the call won't free the vec. The second issue is the more problematic one.

[MSxDOS]

The example now produces good assembly if the loop gets unrolled (possibly improved by #91352), no change otherwise.

[SUPERCILEX]

Probably going to be fixed by #114370

[kornelski]

A very similar case generates 3 separate calls to reserve_for_push:

pub fn test(vec: &mut Vec<u32>) {
    vec.reserve(3);
    let mut i = 0;
    while i < 3 {
        vec.push(32);
        i += 1;
    }
}

https://rust.godbolt.org/z/cr7zf6vqq

[krtab]

Unfortunately, I don't think that neither this issue nor #105156 will be fixed by either:

• PR Hint optimizer about reserved capacity #116568
• PR Add invariant to Vec::pop that len < cap if pop successful #114370

The issue is actually pretty hard to fix. One needs the compiler to deduce that given an assume that after the reserve the, capacity - len >= 3 (easy), there is at each point capacity > len.

For now, I haven't even got LLVM to deduce that capacity > len for the first push (which should come from the fact that 3 > 0 and no overflow occurs).

The diff I have tried variations of:

diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs
index 35015238e6e..acf18bb6dc9 100644
--- a/library/alloc/src/vec/mod.rs
+++ b/library/alloc/src/vec/mod.rs
@@ -907,6 +907,10 @@ pub fn capacity(&self) -> usize {
     #[stable(feature = "rust1", since = "1.0.0")]
     pub fn reserve(&mut self, additional: usize) {
         self.buf.reserve(self.len, additional);
+        unsafe {
+            core::intrinsics::assume(self.buf.capacity() >= self.len.unchecked_add(additional));
+            // core::intrinsics::assume(additional == 0 || self.buf.capacity() > self.len);
+        }
     }

     /// Reserves the minimum capacity for at least `additional` more elements to
@@ -1878,6 +1882,7 @@ fn drop(&mut self) {
     pub fn push(&mut self, value: T) {
         // This will panic or abort if we would allocate > isize::MAX bytes
         // or if the length increment would overflow for zero-sized types.
+        let old_delta = unsafe {self.buf.capacity().unchecked_sub(self.len)};
         if self.len == self.buf.capacity() {
             self.buf.reserve_for_push(self.len);
         }
@@ -1885,6 +1890,8 @@ pub fn push(&mut self, value: T) {
             let end = self.as_mut_ptr().add(self.len);
             ptr::write(end, value);
             self.len += 1;
+            let new_delta = self.buf.capacity().unchecked_sub(self.len);
+            core::intrinsics::assume(old_delta == 0 || old_delta == new_delta.unchecked_add(1));
         }
     }

Commenting the commented assert allow to remove the first push but not subsequent ones, so I'm not sure it is really worth a PR...

[krtab]

llvm/llvm-project#75524 has been merged today.

It makes it so that LLVM understands that "x + y <= z" implies "x <= z" when y is non-negative. In our case it means that "len + additional <= capacity" implies "len <= capacity".

Once rustc's LLVM is updated to take it into account this is likely to change what rustc can optimize or not regarding this issue.

[c410-f3r]

Any news?