alloc/collections/vec_deque/drain.rs
1use core::iter::FusedIterator;
2use core::marker::PhantomData;
3use core::mem::{self, DropGuard, SizedTypeProperties};
4use core::ptr::NonNull;
5use core::{fmt, ptr};
6
7use super::VecDeque;
8use super::index::WrappedIndex;
9use crate::alloc::{Allocator, Global};
10
11/// A draining iterator over the elements of a `VecDeque`.
12///
13/// This `struct` is created by the [`drain`] method on [`VecDeque`]. See its
14/// documentation for more.
15///
16/// [`drain`]: VecDeque::drain
17#[stable(feature = "drain", since = "1.6.0")]
18pub struct Drain<
19 'a,
20 T: 'a,
21 #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
22> {
23 // We can't just use a &mut VecDeque<T, A>, as that would make Drain invariant over T
24 // and we want it to be covariant instead
25 pub(super) deque: NonNull<VecDeque<T, A>>,
26 // drain_start is stored in deque.len
27 pub(super) drain_len: usize,
28 // index into the logical array, not the physical one (always lies in [0..deque.len))
29 pub(super) idx: usize,
30 // number of elements after the drained range
31 pub(super) tail_len: usize,
32 pub(super) remaining: usize,
33 // Needed to make Drain covariant over T
34 _marker: PhantomData<&'a T>,
35}
36
37impl<'a, T, A: Allocator> Drain<'a, T, A> {
38 pub(super) unsafe fn new(
39 deque: &'a mut VecDeque<T, A>,
40 drain_start: usize,
41 drain_len: usize,
42 ) -> Self {
43 let orig_len = mem::replace(&mut deque.len, drain_start);
44 let tail_len = orig_len - drain_start - drain_len;
45 Drain {
46 deque: NonNull::from(deque),
47 drain_len,
48 idx: drain_start,
49 tail_len,
50 remaining: drain_len,
51 _marker: PhantomData,
52 }
53 }
54
55 // Only returns pointers to the slices, as that's all we need
56 // to drop them. May only be called if `self.remaining != 0`.
57 pub(super) unsafe fn as_slices(&self) -> (*mut [T], *mut [T]) {
58 // ignore-tidy-undocumented-unsafe
59 unsafe {
60 let deque = self.deque.as_ref();
61
62 // We know that `self.idx + self.remaining <= deque.len <= usize::MAX`, so this won't overflow.
63 let logical_remaining_range = self.idx..self.idx + self.remaining;
64
65 // SAFETY: `logical_remaining_range` represents the
66 // range into the logical buffer of elements that
67 // haven't been drained yet, so they're all initialized,
68 // and `slice::range(start..end, end) == start..end`,
69 // so the preconditions for `slice_ranges` are met.
70 let (a_range, b_range) =
71 deque.slice_ranges(logical_remaining_range.clone(), logical_remaining_range.end);
72 (deque.buffer_range(a_range), deque.buffer_range(b_range))
73 }
74 }
75}
76
77#[stable(feature = "collection_debug", since = "1.17.0")]
78impl<T: fmt::Debug, A: Allocator> fmt::Debug for Drain<'_, T, A> {
79 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
80 f.debug_tuple("Drain")
81 .field(&self.drain_len)
82 .field(&self.idx)
83 .field(&self.tail_len)
84 .field(&self.remaining)
85 .finish()
86 }
87}
88
89#[stable(feature = "drain", since = "1.6.0")]
90unsafe impl<T: Sync, A: Allocator + Sync> Sync for Drain<'_, T, A> {}
91#[stable(feature = "drain", since = "1.6.0")]
92unsafe impl<T: Send, A: Allocator + Send> Send for Drain<'_, T, A> {}
93
94#[stable(feature = "drain", since = "1.6.0")]
95impl<T, A: Allocator> Drop for Drain<'_, T, A> {
96 fn drop(&mut self) {
97 // Dropping `guard` handles moving the remaining elements into place.
98 let mut guard = DropGuard::new(self, |drain| {
99 if mem::needs_drop::<T>() && drain.remaining != 0 {
100 // SAFETY: We just checked that `self.remaining != 0`.
101 unsafe {
102 let (front, back) = drain.as_slices();
103 ptr::drop_in_place(front);
104 ptr::drop_in_place(back);
105 }
106 }
107
108 // ignore-tidy-undocumented-unsafe
109 let source_deque = unsafe { drain.deque.as_mut() };
110
111 let drain_len = drain.drain_len;
112 let head_len = source_deque.len; // #elements in front of the drain
113 let tail_len = drain.tail_len; // #elements behind the drain
114 let new_len = head_len + tail_len;
115
116 if T::IS_ZST {
117 // no need to copy around any memory if T is a ZST
118 source_deque.len = new_len;
119 return;
120 }
121
122 // Next, we will fill the hole left by the drain with as few writes as possible.
123 // The code below handles the following control flow and reduces the amount of
124 // branches under the assumption that `head_len == 0 || tail_len == 0`, i.e.
125 // draining at the front or at the back of the dequeue is especially common.
126 //
127 // H = "head index" = `deque.head`
128 // h = elements in front of the drain
129 // d = elements in the drain
130 // t = elements behind the drain
131 //
132 // Note that the buffer may wrap at any point and the wrapping is handled by
133 // `wrap_copy` and `to_physical_idx`.
134 //
135 // Case 1: if `head_len == 0 && tail_len == 0`
136 // Everything was drained, reset the head index back to 0.
137 // H
138 // [ . . . . . d d d d . . . . . ]
139 // H
140 // [ . . . . . . . . . . . . . . ]
141 //
142 // Case 2: else if `tail_len == 0`
143 // Don't move data or the head index.
144 // H
145 // [ . . . h h h h d d d d . . . ]
146 // H
147 // [ . . . h h h h . . . . . . . ]
148 //
149 // Case 3: else if `head_len == 0`
150 // Don't move data, but move the head index.
151 // H
152 // [ . . . d d d d t t t t . . . ]
153 // H
154 // [ . . . . . . . t t t t . . . ]
155 //
156 // Case 4: else if `tail_len <= head_len`
157 // Move data, but not the head index.
158 // H
159 // [ . . h h h h d d d d t t . . ]
160 // H
161 // [ . . h h h h t t . . . . . . ]
162 //
163 // Case 5: else
164 // Move data and the head index.
165 // H
166 // [ . . h h d d d d t t t t . . ]
167 // H
168 // [ . . . . . . h h t t t t . . ]
169
170 // When draining at the front (`.drain(..n)`) or at the back (`.drain(n..)`),
171 // we don't need to copy any data. The number of elements copied would be 0.
172 if head_len != 0 && tail_len != 0 {
173 join_head_and_tail_wrapping(source_deque, drain_len, head_len, tail_len);
174 // Marking this function as cold helps LLVM to eliminate it entirely if
175 // this branch is never taken.
176 // We use `#[cold]` instead of `#[inline(never)]`, because inlining this
177 // function into the general case (`.drain(n..m)`) is fine.
178 // See `tests/codegen-llvm/vecdeque-drain.rs` for a test.
179 #[cold]
180 fn join_head_and_tail_wrapping<T, A: Allocator>(
181 source_deque: &mut VecDeque<T, A>,
182 drain_len: usize,
183 head_len: usize,
184 tail_len: usize,
185 ) {
186 // Pick whether to move the head or the tail here.
187 let (src, dst, len);
188 if head_len < tail_len {
189 src = source_deque.head;
190 dst = source_deque.to_wrapped_index(drain_len);
191 len = head_len;
192 } else {
193 src = source_deque.to_wrapped_index(head_len + drain_len);
194 dst = source_deque.to_wrapped_index(head_len);
195 len = tail_len;
196 };
197
198 // ignore-tidy-undocumented-unsafe
199 unsafe {
200 source_deque.wrap_copy(src, dst, len);
201 }
202 }
203 }
204
205 if new_len == 0 {
206 // Special case: If the entire deque was drained, reset the head back to 0,
207 // like `.clear()` does.
208 source_deque.head = WrappedIndex::zero();
209 } else if head_len < tail_len {
210 // If we moved the head above, then we need to adjust the head index here.
211 source_deque.head = source_deque.to_wrapped_index(drain_len);
212 }
213 source_deque.len = new_len;
214 });
215
216 if mem::needs_drop::<T>() && guard.remaining != 0 {
217 // SAFETY: We just checked that `self.remaining != 0`.
218 let (front, back) = unsafe { guard.as_slices() };
219 // since idx is a logical index, we don't need to worry about wrapping.
220 guard.idx += front.len();
221 guard.remaining -= front.len();
222 // SAFETY: This can't have been dropped before since
223 // `idx` & `remaining` track what's been dropped.
224 unsafe { ptr::drop_in_place(front) };
225 guard.remaining = 0;
226 // SAFETY: Ditto.
227 unsafe { ptr::drop_in_place(back) };
228 }
229 }
230}
231
232#[stable(feature = "drain", since = "1.6.0")]
233impl<T, A: Allocator> Iterator for Drain<'_, T, A> {
234 type Item = T;
235
236 #[inline]
237 fn next(&mut self) -> Option<T> {
238 if self.remaining == 0 {
239 return None;
240 }
241 // ignore-tidy-undocumented-unsafe
242 let wrapped_idx = unsafe { self.deque.as_ref().to_wrapped_index(self.idx) };
243 self.idx += 1;
244 self.remaining -= 1;
245 // ignore-tidy-undocumented-unsafe
246 Some(unsafe { self.deque.as_mut().buffer_read(wrapped_idx) })
247 }
248
249 #[inline]
250 fn size_hint(&self) -> (usize, Option<usize>) {
251 let len = self.remaining;
252 (len, Some(len))
253 }
254}
255
256#[stable(feature = "drain", since = "1.6.0")]
257impl<T, A: Allocator> DoubleEndedIterator for Drain<'_, T, A> {
258 #[inline]
259 fn next_back(&mut self) -> Option<T> {
260 if self.remaining == 0 {
261 return None;
262 }
263 self.remaining -= 1;
264 let wrapped_idx =
265 // ignore-tidy-undocumented-unsafe
266 unsafe { self.deque.as_ref().to_wrapped_index(self.idx + self.remaining) };
267 // ignore-tidy-undocumented-unsafe
268 Some(unsafe { self.deque.as_mut().buffer_read(wrapped_idx) })
269 }
270}
271
272#[stable(feature = "drain", since = "1.6.0")]
273impl<T, A: Allocator> ExactSizeIterator for Drain<'_, T, A> {}
274
275#[stable(feature = "fused", since = "1.26.0")]
276impl<T, A: Allocator> FusedIterator for Drain<'_, T, A> {}