Optimize Vec push by preventing address escapes

[gerben-stavenga]

View all comments

This PR changes Vec push and growth so the vector header can remain local to the caller and LLVM can keep its pointer, capacity, and length in SSA values/registers across push loops.

Library changes

• Move the allocator from RawVecInner<A> to RawVec<T, A>. RawVecInner then contains only the allocation pointer and capacity and can be Copy.
• Pass the concrete allocator as a generic &A parameter to RawVecInner operations that need it. This does not use allocator type erasure.
• Make the first outlined grow_one path consume and return RawVecInner by value instead of taking &mut RawVecInner.
• Keep grow_amortized as the single allocation/growth implementation; it now consumes and returns RawVecInner and receives the allocator separately.
• For a non-ZST allocator, temporarily capture the allocator in an unwind-safe local before entering the outlined growth path, then restore it on normal return or unwind. This prevents the containing Vec reference from being carried into allocator calls.
• Special-case ZST allocators so the outlined growth entry does not need an allocator-reference argument.
• Make Vec::push and Vec::push_mut always inline so the fast path and vector fields remain visible in the caller.
• Preserve const_heap support with conditionally-const generic allocator calls.

MIR optimization

The library optimization helps locally owned vectors, but a function taking &mut Vec<T> still begins with a borrowed vector header. This PR therefore adds an experimental CaptureMutVec MIR optimization at MIR optimization level 3.

The pass considers a &mut Vec<T> argument when:

1. At least one direct Vec::push call is in a CFG cycle. A one-off push does not justify the transformation.
2. Every use of the argument is as the receiver of a direct inherent Vec method. Other methods such as is_empty, clear, reserve, or truncate therefore do not inhibit capture.
3. The receiver reference is consumed only by that method call; passing &mut Vec<T> to an unknown helper or otherwise exposing it rejects the candidate.
4. A method does not return a receiver-tied value that could retain the local header address, such as Vec::drain. Proving that such a value does not escape would require additional dataflow.
5. The element layout is at most 16 bytes. The initial implementation considers at most one qualifying argument per function.

When those conditions hold, the pass:

1. Moves the vector header into an owned MIR local at function entry.
2. Rewrites both shared and mutable Vec method receiver borrows to refer to that local.
3. Moves the vector header back into the original argument on every normal return and, for unwinding targets, through a shared cleanup path before unwinding. Aborting targets do not receive an invalid unwind cleanup.

This transformation is enabled by a guarantee specific to Rust: an active &mut Vec<T> provides exclusive access to the Vec value, so a legal caller cannot concurrently observe or modify its header through another reference. The conservative MIR use analysis additionally ensures that the callee neither exposes the header address nor passes it to an unknown function. An equivalent copy-in/copy-out transformation is not generally available for an aliasable mutable pointer.

This extends the same local-header optimization opportunity to borrowed-vector push loops while allowing normal direct Vec operations around and within the loop.

Benchmarks

The PR includes benchmarks covering growing and preallocated locally owned vectors, borrowed vectors, manually captured borrowed vectors, and vectors passed and returned by value.

Machine code from push_grow

The measured growing-vector benchmark uses this hot function:

#[inline(never)]
fn push_grow(n: usize) -> Vec<usize> {
    let mut v = Vec::new();
    for i in 0..n {
        v.push(i);
    }
    v
}

The following x86-64 assembly was generated from that function for the upstream base (e457a7b0d32) and the library portion of this PR (dcb2d85f141) with the same stage compiler configuration and -O -C panic=abort -C codegen-units=1.

Before, the Vec header lives at 8(%rsp), 16(%rsp), and 24(%rsp). The loop stores length on every iteration, passes a pointer to the stack-resident header to growth, then reloads capacity and pointer:

.LBB4_4:
    movq    %r13, (%rax,%r13,8)
    incq    %r13
    movq    %r13, 24(%rsp)       # store len
    cmpq    %r13, %r14
    je      .LBB4_5
.LBB4_2:
    cmpq    %rcx, %r13           # len == cap?
    jne     .LBB4_4
    movq    %r15, %rdi           # &mut stack-resident Vec header
    callq   *%r12
    movq    8(%rsp), %rcx        # reload cap
    movq    16(%rsp), %rax       # reload ptr
    jmp     .LBB4_4

After, capacity stays in %rax, pointer in %rsi, and length in %r15. Growth consumes the current length/header state by value and returns the new capacity and pointer. There is no vector-header stack traffic in the loop:

.LBB1_6:
    movq    %r15, (%rsi,%r15,8)
    incq    %r15
    cmpq    %r15, %rbx
    je      .LBB1_2
.LBB1_4:
    cmpq    %rax, %r15           # len == cap?
    jne     .LBB1_6
    movl    $8, %edx
    movl    $8, %ecx
    movq    %r15, %rdi           # len passed by value
    callq   *%r12
    movq    %rdx, %rsi           # returned ptr; cap remains in %rax
    jmp     .LBB1_6

The upstream function reserves a 32-byte stack frame for the header; the PR function only has the 8-byte alignment slot required around calls.

Disclosure: The implementation and this description were developed with assistance from OpenAI Codex and reviewed by the author.

[saethlin]

• All these functions are #[inline(always)]

Please try to only use that attribute where it is demonstrated to be better than #[inline].

Secondary benefit: grow_one_impl and other growth functions use &dyn Allocator, so they are compiled once in libstd rather than monomorphized per allocator type.

Isn't this a penalty for small custom allocators that can be inlined?

[gerben-stavenga]

• All these functions are #[inline(always)]

Please try to only use that attribute where it is demonstrated to be better than #[inline].

These are on xxx(&mut self) functions that forward to functions that take by self and return self. If not always inlined the &mut self escapes a reference, producing code that is drastically worse. The inner loop in the benchmarks with this PR is

bf330: mov %r13,(%rdx,%r13,8) ; vec[len] = len
bf334: inc %r13 ; len++
bf337: cmp %r13,%r15 ; compare with target
bf33a: je bf370 ; done if equal
bf33c: cmp %rax,%r13 ; compare with capacity
bf33f: jne bf330 ; loop back

before:

bf600: mov -0x38(%rbp),%rax ; LOAD ptr from stack
bf604: mov %r15,(%rax,%r15,8) ; vec[len] = len
bf608: inc %r15 ; len++
bf60b: mov %r15,-0x30(%rbp) ; STORE len to stack
bf60f: cmp %r15,%r14 ; compare with target
bf612: je bf630 ; done if equal
bf614: cmp -0x40(%rbp),%r15 ; LOAD capacity from stack
bf618: jne bf600 ; loop back

Secondary benefit: grow_one_impl and other growth functions use &dyn Allocator, so they are compiled once in libstd rather than monomorphized per allocator type.

Isn't this a penalty for small custom allocators that can be inlined?

I suspect there is a small penalty due to indirection (although the compiler seem to generate call reg in the direct case too). But there are also positive side effects due to code dedup. These are fallback paths so from that perspective a tiny regression isn't the worst. The point of this PR is that the existence of fallback path should not influence the compilers ability to optimize the fast path and keep that clean and tight.

The &dyn Allocator change can be changed to &Allocator at the cost of monomorphizing grow function.

[rustbot]

r? @tgross35

rustbot has assigned @tgross35.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

[tgross35]

@bors try @rust-timer queue

[rust-timer]

Awaiting bors try build completion.

@rustbot label: +S-waiting-on-perf

[rust-bors]

⌛ Trying commit ac22726 with merge b00b54d…

To cancel the try build, run the command @bors try cancel.

Workflow: https://github.com/rust-lang/rust/actions/runs/20890110771

[tgross35]

^ to reiterate that, the rule of thumb now is that any use of #[inline(always)] needs to be backed up by benchmarks and codegen showing it makes a meaningful difference over ‘#[inline]. ‘#[inline(always)]` hurts unoptimized builds and size-optimized binaries so we need to be very cautious with its use.

In general here, it would be helpful if you could put a mini version of the before and after code on godbolt so we can get the bigger picture of what’s actually happening at the different levels.

[gerben-stavenga]

^ to reiterate that, the rule of thumb now is that any use of #[inline(always)] needs to be backed up by benchmarks and codegen showing it makes a meaningful difference over ‘#[inline]. ‘#[inline(always)]` hurts unoptimized builds and size-optimized binaries so we need to be very cautious with its use.

In general here, it would be helpful if you could put a mini version of the before and after code on godbolt so we can get the bigger picture of what’s actually happening at the different levels.

https://godbolt.org/z/nrnP4T83e

shows a rather minimal version, you can see the difference in codegen the test functions

vec_push vs rf_push

[JohnCSimon]

Ping from triage: @gerben-stavenga - can you post your status on this PR? I'll be closing because of inactivity otherwise. Thanks

[rustbot]

This PR was rebased onto a different main commit. Here's a range-diff highlighting what actually changed.

Rebasing is a normal part of keeping PRs up to date, so no action is needed—this note is just to help reviewers.

[gerben-stavenga]

Ping from triage: @gerben-stavenga - can you post your status on this PR? I'll be closing because of inactivity otherwise. Thanks

Hi,

I re-synced the pr and fixed some const-eval issues.

The basic optimization is I think good. Ensuring that all relevant state is passed-in and returned by value is good and can drastically improve codegen as shown by the benchmarks.

The type-erasure is something I like. De-duplicating the same outline coldish resizing code is something I like, but willing to remove. It's all up to the maintainers. If you decide its worthwhile to pursue I'm willing to update the PR as you see fit. If not close it out.

[rustbot]

Some changes occurred to MIR optimizations

cc @rust-lang/wg-mir-opt

[tgross35]

I haven't been keeping up with reviews unfortunately

r? libs