Skip to content

Commit bd6b03e

Browse files
Rollup merge of rust-lang#118264 - lukas-code:optimized-draining, r=the8472
Optimize `VecDeque::drain` for (half-)open ranges The most common use cases of `VecDeque::drain` consume either the entire queue or elements from the front or back.[^1] This PR makes these operations faster by optimizing the generated code of the destructor of the drain: * `.drain(..)` is now the same as `.clear()`. * `.drain(n..)` is now (almost[^2]) the same as `.truncate(n)`. * `.drain(..n)` is now an efficient "advance" function. This operation is not provided by a dedicated function and optimizing it is my main motivation for this PR. Previously, all of these cases generated a function call to the destructor of the `DropGuard`, emitting a lot of unused machine code as well as unnecessary branches and loads/stores of stack variables. There are no algorithmic changes in this PR, but it simplifies the code enough to allow LLVM to recognize the special cases and optimize accordingly. Most notably, it allows elimination of the rather large [`wrap_copy`] function. Some [rudimentary microbenchmarks][benches] show a performance improvement of **~3x-4x** on my machine for the special cases and roughly equal performance for the general case. Best reviewed commit by commit. [^1]: source: GitHub code search: [full range `drain(..)` = 7.5k results][full], [from front `drain(..n)` = 3.2k results][front], [from back `drain(n..)` = 1.6k results][back], [from middle `drain(n..m)` = <500 results][middle] [^2]: `.drain(0..)` and `.clear()` reset the head to 0, but `.truncate(0)` does not. [full]: https://github.com/search?type=code&q=%2FVecDeque%28.%7C%5Cn%29%2B%5C.drain%5C%280%3F%5C.%5C.%5C%29%2F+lang%3ARust [front]: https://github.com/search?type=code&q=%2FVecDeque%28.%7C%5Cn%29%2B%5C.drain%5C%280%3F%5C.%5C.%5B%5E%29%5D.*%5C%29%2F+lang%3ARust [back]: https://github.com/search?type=code&q=%2FVecDeque%28.%7C%5Cn%29%2B%5C.drain%5C%28%5B%5E0%5D.*%5C.%5C.%5C%29%2F+lang%3ARust [middle]: https://github.com/search?type=code&q=%2FVecDeque%28.%7C%5Cn%29%2B%5C.drain%5C%28%5B%5E0%5D.*%5C.%5C.%5B%5E%29%5D.*%5C%29%2F+lang%3ARust [`wrap_copy`]: https://github.com/rust-lang/rust/blob/4fd68eb47bad1c121417ac4450b2f0456150db86/library/alloc/src/collections/vec_deque/mod.rs#L262-L391 [benches]: https://gist.github.com/lukas-code/c97bd707d074c4cc31f241edbc7fd2a2 <details> <summary>generated assembly</summary> before: ```asm clear: sub rsp, 40 mov rax, qword ptr [rdi + 24] mov qword ptr [rdi + 24], 0 mov qword ptr [rsp], rdi mov qword ptr [rsp + 8], rax xorps xmm0, xmm0 movups xmmword ptr [rsp + 16], xmm0 mov qword ptr [rsp + 32], rax test rax, rax je .LBB1_2 mov rcx, qword ptr [rdi] mov rdx, qword ptr [rdi + 16] xor esi, esi cmp rdx, rcx cmovae rsi, rcx sub rdx, rsi mov rsi, rcx sub rsi, rdx lea rdi, [rdx + rax] cmp rsi, rax cmovb rdi, rcx sub rdi, rdx mov qword ptr [rsp + 16], rdi mov qword ptr [rsp + 32], 0 .LBB1_2: mov rdi, rsp call core::ptr::drop_in_place<<alloc::collections::vec_deque::drain::Drain<T,A> as core::ops::drop::Drop>::drop::DropGuard<i32,alloc::alloc::Global>> add rsp, 40 ret truncate: mov rax, qword ptr [rdi + 24] sub rax, rsi jbe .LBB2_2 sub rsp, 40 mov qword ptr [rdi + 24], rsi mov qword ptr [rsp], rdi mov qword ptr [rsp + 8], rax mov rcx, qword ptr [rdi] mov rdx, qword ptr [rdi + 16] add rdx, rsi xor edi, edi cmp rdx, rcx cmovae rdi, rcx mov qword ptr [rsp + 24], 0 sub rdx, rdi mov rdi, rcx sub rdi, rdx lea r8, [rdx + rax] cmp rdi, rax cmovb r8, rcx sub rsi, rdx add rsi, r8 mov qword ptr [rsp + 16], rsi mov qword ptr [rsp + 32], 0 mov rdi, rsp call core::ptr::drop_in_place<<alloc::collections::vec_deque::drain::Drain<T,A> as core::ops::drop::Drop>::drop::DropGuard<i32,alloc::alloc::Global>> add rsp, 40 advance: mov rcx, qword ptr [rdi + 24] mov rax, rcx sub rax, rsi jbe .LBB3_1 sub rsp, 40 mov qword ptr [rdi + 24], 0 mov qword ptr [rsp], rdi mov qword ptr [rsp + 8], rsi mov qword ptr [rsp + 16], 0 mov qword ptr [rsp + 24], rax mov qword ptr [rsp + 32], rsi test rsi, rsi je .LBB3_6 mov rax, qword ptr [rdi] mov rcx, qword ptr [rdi + 16] xor edx, edx cmp rcx, rax cmovae rdx, rax sub rcx, rdx mov rdx, rax sub rdx, rcx lea rdi, [rcx + rsi] cmp rdx, rsi cmovb rdi, rax sub rdi, rcx mov qword ptr [rsp + 16], rdi mov qword ptr [rsp + 32], 0 .LBB3_6: mov rdi, rsp call core::ptr::drop_in_place<<alloc::collections::vec_deque::drain::Drain<T,A> as core::ops::drop::Drop>::drop::DropGuard<i32,alloc::alloc::Global>> add rsp, 40 ret .LBB3_1: test rcx, rcx je .LBB3_3 mov qword ptr [rdi + 24], 0 .LBB3_3: mov qword ptr [rdi + 16], 0 ret remove: sub rsp, 40 cmp rdx, rsi jb .LBB4_5 mov rax, qword ptr [rdi + 24] mov rcx, rax sub rcx, rdx jb .LBB4_6 mov qword ptr [rdi + 24], rsi mov qword ptr [rsp], rdi sub rdx, rsi mov qword ptr [rsp + 8], rdx mov qword ptr [rsp + 16], rsi mov qword ptr [rsp + 24], rcx mov qword ptr [rsp + 32], rdx je .LBB4_4 mov rax, qword ptr [rdi] mov rcx, qword ptr [rdi + 16] add rcx, rsi xor edi, edi cmp rcx, rax cmovae rdi, rax sub rcx, rdi mov rdi, rax sub rdi, rcx lea r8, [rcx + rdx] cmp rdi, rdx cmovb r8, rax sub rsi, rcx add rsi, r8 mov qword ptr [rsp + 16], rsi mov qword ptr [rsp + 32], 0 .LBB4_4: mov rdi, rsp call core::ptr::drop_in_place<<alloc::collections::vec_deque::drain::Drain<T,A> as core::ops::drop::Drop>::drop::DropGuard<i32,alloc::alloc::Global>> add rsp, 40 ret .LBB4_5: lea rax, [rip + .L__unnamed_2] mov rdi, rsi mov rsi, rdx mov rdx, rax call qword ptr [rip + core::slice::index::slice_index_order_fail@GOTPCREL] .LBB4_6: lea rcx, [rip + .L__unnamed_2] mov rdi, rdx mov rsi, rax mov rdx, rcx call qword ptr [rip + core::slice::index::slice_end_index_len_fail@GOTPCREL] core::ptr::drop_in_place<<alloc::collections::vec_deque::drain::Drain<T,A> as core::ops::drop::Drop>::drop::DropGuard<i32,alloc::alloc::Global>>: push rbp push r15 push r14 push r13 push r12 push rbx sub rsp, 24 mov rsi, qword ptr [rdi + 32] test rsi, rsi je .LBB0_2 mov rax, qword ptr [rdi + 16] add rsi, rax jb .LBB0_45 .LBB0_2: mov r13, qword ptr [rdi] mov rbp, qword ptr [rdi + 8] mov rbx, qword ptr [r13 + 24] lea r12, [rbx + rbp] mov r15, qword ptr [rdi + 24] lea rsi, [r15 + r12] test rbx, rbx je .LBB0_10 test r15, r15 je .LBB0_42 cmp rbx, r15 jbe .LBB0_12 mov r14, qword ptr [r13] mov rax, qword ptr [r13 + 16] add r12, rax xor ecx, ecx cmp r12, r14 mov rdx, r14 cmovb rdx, rcx sub r12, rdx add rbx, rax cmp rbx, r14 cmovae rcx, r14 sub rbx, rcx mov rcx, rbx sub rcx, r12 je .LBB0_42 mov rdi, qword ptr [r13 + 8] mov rax, rcx add rax, r14 cmovae rax, rcx mov r8, r14 sub r8, r12 mov rcx, r14 sub rcx, rbx mov rdx, r15 sub rdx, r8 mov qword ptr [rsp + 16], rsi jbe .LBB0_18 cmp rax, r15 jae .LBB0_24 mov rdx, r15 sub rdx, r8 shl rdx, 2 cmp r15, rcx jbe .LBB0_30 sub r8, rcx mov qword ptr [rsp], rdi mov rax, qword ptr [rsp] lea rdi, [rax + 4*r8] mov rsi, qword ptr [rsp] mov qword ptr [rsp + 8], rcx mov r15, r8 call qword ptr [rip + memmove@GOTPCREL] sub r14, r15 mov rax, qword ptr [rsp] lea rsi, [rax + 4*r14] shl r15, 2 mov rdi, qword ptr [rsp] mov rdx, r15 call qword ptr [rip + memmove@GOTPCREL] mov rdi, qword ptr [rsp] lea rsi, [rdi + 4*r12] lea rdi, [rdi + 4*rbx] mov r15, qword ptr [rsp + 8] jmp .LBB0_36 .LBB0_10: test r15, r15 je .LBB0_17 mov rax, qword ptr [r13] sub rsi, rbp add rbp, qword ptr [r13 + 16] xor ecx, ecx cmp rbp, rax cmovae rcx, rax sub rbp, rcx mov qword ptr [r13 + 16], rbp jmp .LBB0_43 .LBB0_12: mov rdx, qword ptr [r13 + 16] mov r15, qword ptr [r13] lea rax, [rdx + rbp] xor ecx, ecx cmp rax, r15 cmovae rcx, r15 mov r12, rax sub r12, rcx mov rcx, r12 sub rcx, rdx je .LBB0_41 mov rdi, qword ptr [r13 + 8] mov rax, rcx add rax, r15 cmovae rax, rcx mov r8, r15 sub r8, rdx mov rcx, r15 sub rcx, r12 mov r14, rbx sub r14, r8 mov qword ptr [rsp + 16], rsi jbe .LBB0_21 cmp rax, rbx jae .LBB0_26 mov qword ptr [rsp], rdx mov rdx, rbx sub rdx, r8 shl rdx, 2 cmp rbx, rcx jbe .LBB0_32 sub r8, rcx mov rbx, rdi lea rdi, [rdi + 4*r8] mov rsi, rbx mov qword ptr [rsp + 8], rcx mov r14, r8 call qword ptr [rip + memmove@GOTPCREL] sub r15, r14 lea rsi, [rbx + 4*r15] shl r14, 2 mov rdi, rbx mov rdx, r14 call qword ptr [rip + memmove@GOTPCREL] mov rdi, rbx mov rax, qword ptr [rsp] lea rsi, [rbx + 4*rax] lea rdi, [rbx + 4*r12] mov rbx, qword ptr [rsp + 8] jmp .LBB0_40 .LBB0_17: xorps xmm0, xmm0 movups xmmword ptr [r13 + 16], xmm0 jmp .LBB0_44 .LBB0_18: mov r14, r15 sub r14, rcx jbe .LBB0_28 cmp rax, r15 jae .LBB0_33 lea rax, [rcx + r12] sub r15, rcx lea rsi, [rdi + 4*rax] shl r15, 2 mov r14, rdi mov rdx, r15 mov r15, rcx jmp .LBB0_31 .LBB0_21: mov r14, rbx sub r14, rcx jbe .LBB0_29 cmp rax, rbx jae .LBB0_34 lea rax, [rcx + rdx] sub rbx, rcx lea rsi, [rdi + 4*rax] shl rbx, 2 mov r14, rdi mov r15, rdx mov rdx, rbx mov rbx, rcx call qword ptr [rip + memmove@GOTPCREL] mov rdi, r14 lea rsi, [r14 + 4*r15] lea rdi, [r14 + 4*r12] jmp .LBB0_40 .LBB0_24: sub r15, rcx jbe .LBB0_35 sub rcx, r8 mov qword ptr [rsp + 8], rcx lea rsi, [rdi + 4*r12] mov r12, rdi lea rdi, [rdi + 4*rbx] lea rdx, [4*r8] mov r14, r8 call qword ptr [rip + memmove@GOTPCREL] add r14, rbx lea rdi, [r12 + 4*r14] mov rbx, qword ptr [rsp + 8] lea rdx, [4*rbx] mov rsi, r12 call qword ptr [rip + memmove@GOTPCREL] mov rdi, r12 lea rsi, [r12 + 4*rbx] jmp .LBB0_36 .LBB0_26: sub rbx, rcx jbe .LBB0_37 sub rcx, r8 lea rsi, [rdi + 4*rdx] mov r15, rdi lea rdi, [rdi + 4*r12] lea rdx, [4*r8] mov r14, rcx mov qword ptr [rsp], r8 call qword ptr [rip + memmove@GOTPCREL] add r12, qword ptr [rsp] lea rdi, [r15 + 4*r12] lea rdx, [4*r14] mov rsi, r15 call qword ptr [rip + memmove@GOTPCREL] mov rdi, r15 lea rsi, [r15 + 4*r14] jmp .LBB0_40 .LBB0_28: lea rsi, [rdi + 4*r12] lea rdi, [rdi + 4*rbx] jmp .LBB0_36 .LBB0_29: lea rsi, [rdi + 4*rdx] lea rdi, [rdi + 4*r12] jmp .LBB0_40 .LBB0_30: lea rax, [r8 + rbx] mov r14, rdi lea rdi, [rdi + 4*rax] mov rsi, r14 mov r15, r8 .LBB0_31: call qword ptr [rip + memmove@GOTPCREL] mov rdi, r14 lea rsi, [r14 + 4*r12] lea rdi, [r14 + 4*rbx] jmp .LBB0_36 .LBB0_32: lea rax, [r12 + r8] mov rbx, rdi lea rdi, [rdi + 4*rax] mov rsi, rbx mov r14, r8 call qword ptr [rip + memmove@GOTPCREL] mov rdi, rbx mov rax, qword ptr [rsp] lea rsi, [rbx + 4*rax] jmp .LBB0_38 .LBB0_33: lea rsi, [rdi + 4*r12] mov r15, rdi lea rdi, [rdi + 4*rbx] lea rdx, [4*rcx] mov rbx, rcx call qword ptr [rip + memmove@GOTPCREL] mov rdi, r15 add rbx, r12 lea rsi, [r15 + 4*rbx] mov r15, r14 jmp .LBB0_36 .LBB0_34: lea rsi, [rdi + 4*rdx] mov rbx, rdi lea rdi, [rdi + 4*r12] mov r15, rdx lea rdx, [4*rcx] mov r12, rcx call qword ptr [rip + memmove@GOTPCREL] mov rdi, rbx add r12, r15 lea rsi, [rbx + 4*r12] jmp .LBB0_39 .LBB0_35: lea rsi, [rdi + 4*r12] mov r14, rdi lea rdi, [rdi + 4*rbx] mov r12, rdx lea rdx, [4*r8] mov r15, r8 call qword ptr [rip + memmove@GOTPCREL] add r15, rbx mov rsi, r14 lea rdi, [r14 + 4*r15] mov r15, r12 .LBB0_36: shl r15, 2 mov rdx, r15 call qword ptr [rip + memmove@GOTPCREL] mov rsi, qword ptr [rsp + 16] jmp .LBB0_42 .LBB0_37: lea rsi, [rdi + 4*rdx] mov rbx, rdi lea rdi, [rdi + 4*r12] lea rdx, [4*r8] mov r15, r8 call qword ptr [rip + memmove@GOTPCREL] add r12, r15 mov rsi, rbx .LBB0_38: lea rdi, [rbx + 4*r12] .LBB0_39: mov rbx, r14 .LBB0_40: shl rbx, 2 mov rdx, rbx call qword ptr [rip + memmove@GOTPCREL] mov r15, qword ptr [r13] mov rax, qword ptr [r13 + 16] add rax, rbp mov rsi, qword ptr [rsp + 16] .LBB0_41: xor ecx, ecx cmp rax, r15 cmovae rcx, r15 sub rax, rcx mov qword ptr [r13 + 16], rax .LBB0_42: sub rsi, rbp .LBB0_43: mov qword ptr [r13 + 24], rsi .LBB0_44: add rsp, 24 pop rbx pop r12 pop r13 pop r14 pop r15 pop rbp ret .LBB0_45: lea rdx, [rip + .L__unnamed_1] mov rdi, rax call qword ptr [rip + core::slice::index::slice_index_order_fail@GOTPCREL] ``` after: ```asm clear: movups xmmword ptr [rdi + 16], xmm0 ret truncate: cmp qword ptr [rdi + 24], rsi jbe .LBB2_4 test rsi, rsi jne .LBB2_3 mov qword ptr [rdi + 16], 0 .LBB2_3: mov qword ptr [rdi + 24], rsi .LBB2_4: ret advance: mov rcx, qword ptr [rdi + 24] mov rax, rcx sub rax, rsi jbe .LBB3_1 mov rcx, qword ptr [rdi] add rsi, qword ptr [rdi + 16] xor edx, edx cmp rsi, rcx cmovae rdx, rcx sub rsi, rdx mov qword ptr [rdi + 16], rsi mov qword ptr [rdi + 24], rax ret .LBB3_1: test rcx, rcx je .LBB3_3 mov qword ptr [rdi + 24], 0 .LBB3_3: mov qword ptr [rdi + 16], 0 ret remove: push rbp push r15 push r14 push r13 push r12 push rbx push rax mov r15, rsi mov r14, rdx sub r14, rsi jb .LBB4_9 mov rbx, rdi mov r12, qword ptr [rdi + 24] mov r13, r12 sub r13, rdx jb .LBB4_10 mov qword ptr [rbx + 24], r15 mov rbp, r12 sub rbp, r14 test r15, r15 je .LBB4_4 cmp rbp, r15 jne .LBB4_11 .LBB4_4: cmp r12, r14 jne .LBB4_6 .LBB4_5: mov qword ptr [rbx + 16], 0 jmp .LBB4_8 .LBB4_11: mov rdi, rbx mov rsi, r14 mov rdx, r15 mov rcx, r13 call <<alloc::collections::vec_deque::drain::Drain<T,A> as core::ops::drop::Drop>::drop::DropGuard<T,A> as core::ops::drop::Drop>::drop::copy_data cmp r12, r14 je .LBB4_5 .LBB4_6: cmp r13, r15 jbe .LBB4_8 mov rax, qword ptr [rbx] add r14, qword ptr [rbx + 16] xor ecx, ecx cmp r14, rax cmovae rcx, rax sub r14, rcx mov qword ptr [rbx + 16], r14 .LBB4_8: mov qword ptr [rbx + 24], rbp add rsp, 8 pop rbx pop r12 pop r13 pop r14 pop r15 pop rbp ret .LBB4_9: lea rax, [rip + .L__unnamed_1] mov rdi, r15 mov rsi, rdx mov rdx, rax call qword ptr [rip + core::slice::index::slice_index_order_fail@GOTPCREL] .LBB4_10: lea rax, [rip + .L__unnamed_1] mov rdi, rdx mov rsi, r12 mov rdx, rax call qword ptr [rip + core::slice::index::slice_end_index_len_fail@GOTPCREL] <<alloc::collections::vec_deque::drain::Drain<T,A> as core::ops::drop::Drop>::drop::DropGuard<T,A> as core::ops::drop::Drop>::drop::copy_data: push rbp push r15 push r14 push r13 push r12 push rbx push rax mov r14, rsi cmp rdx, rcx jae .LBB0_1 mov r12, qword ptr [rdi] mov rax, qword ptr [rdi + 16] add r14, rax xor ecx, ecx cmp r14, r12 cmovae rcx, r12 sub r14, rcx mov r15, rdx mov r13, r14 mov r14, rax mov rcx, r13 sub rcx, r14 je .LBB0_18 .LBB0_4: mov rdi, qword ptr [rdi + 8] mov rax, rcx add rax, r12 cmovae rax, rcx mov rbx, r12 sub rbx, r14 mov rcx, r12 sub rcx, r13 mov rbp, r15 sub rbp, rbx jbe .LBB0_5 cmp rax, r15 jae .LBB0_12 mov rdx, r15 sub rdx, rbx shl rdx, 2 cmp r15, rcx jbe .LBB0_16 sub rbx, rcx mov rbp, rdi lea rdi, [rdi + 4*rbx] mov r15, qword ptr [rip + memmove@GOTPCREL] mov rsi, rbp mov qword ptr [rsp], rcx call r15 sub r12, rbx lea rsi, [4*r12] add rsi, rbp shl rbx, 2 mov rdi, rbp mov rdx, rbx call r15 mov rdi, rbp lea rsi, [4*r14] add rsi, rbp lea rdi, [4*r13] add rdi, rbp mov r15, qword ptr [rsp] jmp .LBB0_7 .LBB0_1: mov r15, rcx add r14, rdx mov r12, qword ptr [rdi] mov r13, qword ptr [rdi + 16] add r14, r13 xor eax, eax cmp r14, r12 mov rcx, r12 cmovb rcx, rax sub r14, rcx add r13, rdx cmp r13, r12 cmovae rax, r12 sub r13, rax mov rcx, r13 sub rcx, r14 jne .LBB0_4 .LBB0_18: add rsp, 8 pop rbx pop r12 pop r13 pop r14 pop r15 pop rbp ret .LBB0_5: mov rbx, r15 sub rbx, rcx jbe .LBB0_6 cmp rax, r15 jae .LBB0_9 lea rax, [rcx + r14] sub r15, rcx lea rsi, [rdi + 4*rax] shl r15, 2 mov rbx, rdi mov rdx, r15 mov r15, rcx call qword ptr [rip + memmove@GOTPCREL] mov rdi, rbx lea rsi, [rbx + 4*r14] lea rdi, [rbx + 4*r13] jmp .LBB0_7 .LBB0_12: sub r15, rcx jbe .LBB0_13 sub rcx, rbx lea rsi, [rdi + 4*r14] mov r12, rdi lea rdi, [rdi + 4*r13] lea rdx, [4*rbx] mov r14, qword ptr [rip + memmove@GOTPCREL] mov rbp, rcx call r14 add rbx, r13 lea rdi, [r12 + 4*rbx] lea rdx, [4*rbp] mov rsi, r12 call r14 mov rdi, r12 lea rsi, [r12 + 4*rbp] jmp .LBB0_7 .LBB0_6: lea rsi, [rdi + 4*r14] lea rdi, [rdi + 4*r13] jmp .LBB0_7 .LBB0_16: lea rax, [rbx + r13] mov r15, rdi lea rdi, [rdi + 4*rax] mov rsi, r15 call qword ptr [rip + memmove@GOTPCREL] mov rdi, r15 lea rsi, [r15 + 4*r14] lea rdi, [r15 + 4*r13] mov r15, rbx jmp .LBB0_7 .LBB0_9: lea rsi, [rdi + 4*r14] mov r15, rdi lea rdi, [rdi + 4*r13] lea rdx, [4*rcx] mov r12, rcx call qword ptr [rip + memmove@GOTPCREL] mov rdi, r15 add r12, r14 lea rsi, [r15 + 4*r12] mov r15, rbx jmp .LBB0_7 .LBB0_13: lea rsi, [rdi + 4*r14] mov r14, rdi lea rdi, [rdi + 4*r13] lea rdx, [4*rbx] call qword ptr [rip + memmove@GOTPCREL] add rbx, r13 mov rsi, r14 lea rdi, [r14 + 4*rbx] mov r15, rbp .LBB0_7: shl r15, 2 mov rdx, r15 add rsp, 8 pop rbx pop r12 pop r13 pop r14 pop r15 pop rbp jmp qword ptr [rip + memmove@GOTPCREL] ``` </details>
2 parents cabdf3a + 5d98977 commit bd6b03e

File tree

2 files changed

+183
-56
lines changed

2 files changed

+183
-56
lines changed

library/alloc/src/collections/vec_deque/drain.rs

+114-56
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ pub struct Drain<
2727
drain_len: usize,
2828
// index into the logical array, not the physical one (always lies in [0..deque.len))
2929
idx: usize,
30-
// number of elements after the drain range
31-
tail_len: usize,
30+
// number of elements remaining after dropping the drain
31+
new_len: usize,
3232
remaining: usize,
3333
// Needed to make Drain covariant over T
3434
_marker: PhantomData<&'a T>,
@@ -41,12 +41,12 @@ impl<'a, T, A: Allocator> Drain<'a, T, A> {
4141
drain_len: usize,
4242
) -> Self {
4343
let orig_len = mem::replace(&mut deque.len, drain_start);
44-
let tail_len = orig_len - drain_start - drain_len;
44+
let new_len = orig_len - drain_len;
4545
Drain {
4646
deque: NonNull::from(deque),
4747
drain_len,
4848
idx: drain_start,
49-
tail_len,
49+
new_len,
5050
remaining: drain_len,
5151
_marker: PhantomData,
5252
}
@@ -79,7 +79,7 @@ impl<T: fmt::Debug, A: Allocator> fmt::Debug for Drain<'_, T, A> {
7979
f.debug_tuple("Drain")
8080
.field(&self.drain_len)
8181
.field(&self.idx)
82-
.field(&self.tail_len)
82+
.field(&self.new_len)
8383
.field(&self.remaining)
8484
.finish()
8585
}
@@ -95,9 +95,26 @@ impl<T, A: Allocator> Drop for Drain<'_, T, A> {
9595
fn drop(&mut self) {
9696
struct DropGuard<'r, 'a, T, A: Allocator>(&'r mut Drain<'a, T, A>);
9797

98+
let guard = DropGuard(self);
99+
100+
if mem::needs_drop::<T>() && guard.0.remaining != 0 {
101+
unsafe {
102+
// SAFETY: We just checked that `self.remaining != 0`.
103+
let (front, back) = guard.0.as_slices();
104+
// since idx is a logical index, we don't need to worry about wrapping.
105+
guard.0.idx += front.len();
106+
guard.0.remaining -= front.len();
107+
ptr::drop_in_place(front);
108+
guard.0.remaining = 0;
109+
ptr::drop_in_place(back);
110+
}
111+
}
112+
113+
// Dropping `guard` handles moving the remaining elements into place.
98114
impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> {
115+
#[inline]
99116
fn drop(&mut self) {
100-
if self.0.remaining != 0 {
117+
if mem::needs_drop::<T>() && self.0.remaining != 0 {
101118
unsafe {
102119
// SAFETY: We just checked that `self.remaining != 0`.
103120
let (front, back) = self.0.as_slices();
@@ -108,70 +125,111 @@ impl<T, A: Allocator> Drop for Drain<'_, T, A> {
108125

109126
let source_deque = unsafe { self.0.deque.as_mut() };
110127

111-
let drain_start = source_deque.len();
112128
let drain_len = self.0.drain_len;
113-
let drain_end = drain_start + drain_len;
114-
115-
let orig_len = self.0.tail_len + drain_end;
129+
let new_len = self.0.new_len;
116130

117131
if T::IS_ZST {
118132
// no need to copy around any memory if T is a ZST
119-
source_deque.len = orig_len - drain_len;
133+
source_deque.len = new_len;
120134
return;
121135
}
122136

123-
let head_len = drain_start;
124-
let tail_len = self.0.tail_len;
137+
let head_len = source_deque.len; // #elements in front of the drain
138+
let tail_len = new_len - head_len; // #elements behind the drain
125139

126-
match (head_len, tail_len) {
127-
(0, 0) => {
128-
source_deque.head = 0;
129-
source_deque.len = 0;
130-
}
131-
(0, _) => {
132-
source_deque.head = source_deque.to_physical_idx(drain_len);
133-
source_deque.len = orig_len - drain_len;
134-
}
135-
(_, 0) => {
136-
source_deque.len = orig_len - drain_len;
137-
}
138-
_ => unsafe {
139-
if head_len <= tail_len {
140-
source_deque.wrap_copy(
141-
source_deque.head,
142-
source_deque.to_physical_idx(drain_len),
143-
head_len,
144-
);
145-
source_deque.head = source_deque.to_physical_idx(drain_len);
146-
source_deque.len = orig_len - drain_len;
140+
// Next, we will fill the hole left by the drain with as few writes as possible.
141+
// The code below handles the following control flow and reduces the amount of
142+
// branches under the assumption that `head_len == 0 || tail_len == 0`, i.e.
143+
// draining at the front or at the back of the dequeue is especially common.
144+
//
145+
// H = "head index" = `deque.head`
146+
// h = elements in front of the drain
147+
// d = elements in the drain
148+
// t = elements behind the drain
149+
//
150+
// Note that the buffer may wrap at any point and the wrapping is handled by
151+
// `wrap_copy` and `to_physical_idx`.
152+
//
153+
// Case 1: if `head_len == 0 && tail_len == 0`
154+
// Everything was drained, reset the head index back to 0.
155+
// H
156+
// [ . . . . . d d d d . . . . . ]
157+
// H
158+
// [ . . . . . . . . . . . . . . ]
159+
//
160+
// Case 2: else if `tail_len == 0`
161+
// Don't move data or the head index.
162+
// H
163+
// [ . . . h h h h d d d d . . . ]
164+
// H
165+
// [ . . . h h h h . . . . . . . ]
166+
//
167+
// Case 3: else if `head_len == 0`
168+
// Don't move data, but move the head index.
169+
// H
170+
// [ . . . d d d d t t t t . . . ]
171+
// H
172+
// [ . . . . . . . t t t t . . . ]
173+
//
174+
// Case 4: else if `tail_len <= head_len`
175+
// Move data, but not the head index.
176+
// H
177+
// [ . . h h h h d d d d t t . . ]
178+
// H
179+
// [ . . h h h h t t . . . . . . ]
180+
//
181+
// Case 5: else
182+
// Move data and the head index.
183+
// H
184+
// [ . . h h d d d d t t t t . . ]
185+
// H
186+
// [ . . . . . . h h t t t t . . ]
187+
188+
// When draining at the front (`.drain(..n)`) or at the back (`.drain(n..)`),
189+
// we don't need to copy any data. The number of elements copied would be 0.
190+
if head_len != 0 && tail_len != 0 {
191+
join_head_and_tail_wrapping(source_deque, drain_len, head_len, tail_len);
192+
// Marking this function as cold helps LLVM to eliminate it entirely if
193+
// this branch is never taken.
194+
// We use `#[cold]` instead of `#[inline(never)]`, because inlining this
195+
// function into the general case (`.drain(n..m)`) is fine.
196+
// See `tests/codegen/vecdeque-drain.rs` for a test.
197+
#[cold]
198+
fn join_head_and_tail_wrapping<T, A: Allocator>(
199+
source_deque: &mut VecDeque<T, A>,
200+
drain_len: usize,
201+
head_len: usize,
202+
tail_len: usize,
203+
) {
204+
// Pick whether to move the head or the tail here.
205+
let (src, dst, len);
206+
if head_len < tail_len {
207+
src = source_deque.head;
208+
dst = source_deque.to_physical_idx(drain_len);
209+
len = head_len;
147210
} else {
148-
source_deque.wrap_copy(
149-
source_deque.to_physical_idx(head_len + drain_len),
150-
source_deque.to_physical_idx(head_len),
151-
tail_len,
152-
);
153-
source_deque.len = orig_len - drain_len;
211+
src = source_deque.to_physical_idx(head_len + drain_len);
212+
dst = source_deque.to_physical_idx(head_len);
213+
len = tail_len;
214+
};
215+
216+
unsafe {
217+
source_deque.wrap_copy(src, dst, len);
154218
}
155-
},
219+
}
156220
}
157-
}
158-
}
159221

160-
let guard = DropGuard(self);
161-
if guard.0.remaining != 0 {
162-
unsafe {
163-
// SAFETY: We just checked that `self.remaining != 0`.
164-
let (front, back) = guard.0.as_slices();
165-
// since idx is a logical index, we don't need to worry about wrapping.
166-
guard.0.idx += front.len();
167-
guard.0.remaining -= front.len();
168-
ptr::drop_in_place(front);
169-
guard.0.remaining = 0;
170-
ptr::drop_in_place(back);
222+
if new_len == 0 {
223+
// Special case: If the entire dequeue was drained, reset the head back to 0,
224+
// like `.clear()` does.
225+
source_deque.head = 0;
226+
} else if head_len < tail_len {
227+
// If we moved the head above, then we need to adjust the head index here.
228+
source_deque.head = source_deque.to_physical_idx(drain_len);
229+
}
230+
source_deque.len = new_len;
171231
}
172232
}
173-
174-
// Dropping `guard` handles moving the remaining elements into place.
175233
}
176234
}
177235

tests/codegen/vecdeque-drain.rs

+69
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
// Check that draining at the front or back doesn't copy memory.
2+
3+
// compile-flags: -O
4+
// ignore-debug: the debug assertions get in the way
5+
6+
#![crate_type = "lib"]
7+
8+
use std::collections::VecDeque;
9+
10+
// CHECK-LABEL: @clear
11+
// CHECK-NOT: call
12+
// CHECK-NOT: br
13+
// CHECK: getelementptr inbounds
14+
// CHECK-NEXT: {{call void @llvm.memset|store}}
15+
// CHECK-NEXT: ret void
16+
#[no_mangle]
17+
pub fn clear(v: &mut VecDeque<i32>) {
18+
v.drain(..);
19+
}
20+
21+
// CHECK-LABEL: @truncate
22+
// CHECK-NOT: call
23+
// CHECK: br
24+
// CHECK-NOT: call
25+
// CHECK: br
26+
// CHECK-NOT: call
27+
// CHECK: br
28+
// CHECK-NOT: call
29+
// CHECK: br
30+
// CHECK-NOT: call
31+
// CHECK-NOT: br
32+
// CHECK: ret void
33+
#[no_mangle]
34+
pub fn truncate(v: &mut VecDeque<i32>, n: usize) {
35+
if n < v.len() {
36+
v.drain(n..);
37+
}
38+
}
39+
40+
// CHECK-LABEL: @advance
41+
// CHECK-NOT: call
42+
// CHECK: br
43+
// CHECK-NOT: call
44+
// CHECK: br
45+
// CHECK-NOT: call
46+
// CHECK: br
47+
// CHECK-NOT: call
48+
// CHECK: br
49+
// CHECK-NOT: call
50+
// CHECK: br
51+
// CHECK-NOT: call
52+
// CHECK-NOT: br
53+
// CHECK: ret void
54+
#[no_mangle]
55+
pub fn advance(v: &mut VecDeque<i32>, n: usize) {
56+
if n < v.len() {
57+
v.drain(..n);
58+
} else {
59+
v.clear();
60+
}
61+
}
62+
63+
// CHECK-LABEL: @remove
64+
// CHECK: call
65+
// CHECK: ret void
66+
#[no_mangle]
67+
pub fn remove(v: &mut VecDeque<i32>, a: usize, b: usize) {
68+
v.drain(a..b);
69+
}

0 commit comments

Comments
 (0)