Skip to content

Commit 218fee2

Browse files
author
Lukas Markeffsky
committed
address review comments
1 parent a77d6b9 commit 218fee2

File tree

1 file changed

+99
-44
lines changed
  • library/alloc/src/collections/vec_deque

1 file changed

+99
-44
lines changed

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

Lines changed: 99 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,22 @@ 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> {
99115
#[inline]
100116
fn drop(&mut self) {
@@ -118,63 +134,102 @@ impl<T, A: Allocator> Drop for Drain<'_, T, A> {
118134
return;
119135
}
120136

121-
let head_len = source_deque.len();
122-
let tail_len = new_len - head_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
139+
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 . . ]
123187

124188
// When draining at the front (`.drain(..n)`) or at the back (`.drain(n..)`),
125-
// we don't need to copy any data.
126-
// Outlining this function helps LLVM to eliminate the copies in these cases.
189+
// we don't need to copy any data. The number of elements copied would be 0.
127190
if head_len != 0 && tail_len != 0 {
128-
copy_data(source_deque, drain_len, head_len, tail_len);
191+
join_head_and_tail(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<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;
210+
} else {
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);
218+
}
219+
}
129220
}
130221

131222
if new_len == 0 {
223+
// Special case: If the entire dequeue was drained, reset the head back to 0,
224+
// like `.clear()` does.
132225
source_deque.head = 0;
133226
} else if head_len < tail_len {
227+
// If we moved the head above, then we need to adjust the head index here.
134228
source_deque.head = source_deque.to_physical_idx(drain_len);
135229
}
136230
source_deque.len = new_len;
137-
138-
#[cold]
139-
fn copy_data<T, A: Allocator>(
140-
source_deque: &mut VecDeque<T, A>,
141-
drain_len: usize,
142-
head_len: usize,
143-
tail_len: usize,
144-
) {
145-
let (src, dst, len);
146-
if head_len < tail_len {
147-
src = source_deque.head;
148-
dst = source_deque.to_physical_idx(drain_len);
149-
len = head_len;
150-
} else {
151-
src = source_deque.to_physical_idx(head_len + drain_len);
152-
dst = source_deque.to_physical_idx(head_len);
153-
len = tail_len;
154-
};
155-
156-
unsafe {
157-
source_deque.wrap_copy(src, dst, len);
158-
}
159-
}
160-
}
161-
}
162-
163-
let guard = DropGuard(self);
164-
if mem::needs_drop::<T>() && guard.0.remaining != 0 {
165-
unsafe {
166-
// SAFETY: We just checked that `self.remaining != 0`.
167-
let (front, back) = guard.0.as_slices();
168-
// since idx is a logical index, we don't need to worry about wrapping.
169-
guard.0.idx += front.len();
170-
guard.0.remaining -= front.len();
171-
ptr::drop_in_place(front);
172-
guard.0.remaining = 0;
173-
ptr::drop_in_place(back);
174231
}
175232
}
176-
177-
// Dropping `guard` handles moving the remaining elements into place.
178233
}
179234
}
180235

0 commit comments

Comments
 (0)