alloc/collections/vec_deque/drain.rs
1use core::iter::FusedIterator;
2use core::marker::PhantomData;
3use core::mem::{self, 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 struct DropGuard<'r, 'a, T, A: Allocator>(&'r mut Drain<'a, T, A>);
98
99 let guard = DropGuard(self);
100
101 if mem::needs_drop::<T>() && guard.0.remaining != 0 {
102 // SAFETY: We just checked that `self.remaining != 0`.
103 let (front, back) = unsafe { 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 // SAFETY: This can't have been dropped before since
108 // `idx` & `remaining` track what's been dropped.
109 unsafe { ptr::drop_in_place(front) };
110 guard.0.remaining = 0;
111 // SAFETY: Ditto.
112 unsafe { ptr::drop_in_place(back) };
113 }
114
115 // Dropping `guard` handles moving the remaining elements into place.
116 impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> {
117 #[inline]
118 fn drop(&mut self) {
119 if mem::needs_drop::<T>() && self.0.remaining != 0 {
120 // SAFETY: We just checked that `self.remaining != 0`.
121 unsafe {
122 let (front, back) = self.0.as_slices();
123 ptr::drop_in_place(front);
124 ptr::drop_in_place(back);
125 }
126 }
127
128 // ignore-tidy-undocumented-unsafe
129 let source_deque = unsafe { self.0.deque.as_mut() };
130
131 let drain_len = self.0.drain_len;
132 let head_len = source_deque.len; // #elements in front of the drain
133 let tail_len = self.0.tail_len; // #elements behind the drain
134 let new_len = head_len + tail_len;
135
136 if T::IS_ZST {
137 // no need to copy around any memory if T is a ZST
138 source_deque.len = new_len;
139 return;
140 }
141
142 // Next, we will fill the hole left by the drain with as few writes as possible.
143 // The code below handles the following control flow and reduces the amount of
144 // branches under the assumption that `head_len == 0 || tail_len == 0`, i.e.
145 // draining at the front or at the back of the dequeue is especially common.
146 //
147 // H = "head index" = `deque.head`
148 // h = elements in front of the drain
149 // d = elements in the drain
150 // t = elements behind the drain
151 //
152 // Note that the buffer may wrap at any point and the wrapping is handled by
153 // `wrap_copy` and `to_physical_idx`.
154 //
155 // Case 1: if `head_len == 0 && tail_len == 0`
156 // Everything was drained, reset the head index back to 0.
157 // H
158 // [ . . . . . d d d d . . . . . ]
159 // H
160 // [ . . . . . . . . . . . . . . ]
161 //
162 // Case 2: else if `tail_len == 0`
163 // Don't move data or the head index.
164 // H
165 // [ . . . h h h h d d d d . . . ]
166 // H
167 // [ . . . h h h h . . . . . . . ]
168 //
169 // Case 3: else if `head_len == 0`
170 // Don't move data, but move the head index.
171 // H
172 // [ . . . d d d d t t t t . . . ]
173 // H
174 // [ . . . . . . . t t t t . . . ]
175 //
176 // Case 4: else if `tail_len <= head_len`
177 // Move data, but not the head index.
178 // H
179 // [ . . h h h h d d d d t t . . ]
180 // H
181 // [ . . h h h h t t . . . . . . ]
182 //
183 // Case 5: else
184 // Move data and the head index.
185 // H
186 // [ . . h h d d d d t t t t . . ]
187 // H
188 // [ . . . . . . h h t t t t . . ]
189
190 // When draining at the front (`.drain(..n)`) or at the back (`.drain(n..)`),
191 // we don't need to copy any data. The number of elements copied would be 0.
192 if head_len != 0 && tail_len != 0 {
193 join_head_and_tail_wrapping(source_deque, drain_len, head_len, tail_len);
194 // Marking this function as cold helps LLVM to eliminate it entirely if
195 // this branch is never taken.
196 // We use `#[cold]` instead of `#[inline(never)]`, because inlining this
197 // function into the general case (`.drain(n..m)`) is fine.
198 // See `tests/codegen-llvm/vecdeque-drain.rs` for a test.
199 #[cold]
200 fn join_head_and_tail_wrapping<T, A: Allocator>(
201 source_deque: &mut VecDeque<T, A>,
202 drain_len: usize,
203 head_len: usize,
204 tail_len: usize,
205 ) {
206 // Pick whether to move the head or the tail here.
207 let (src, dst, len);
208 if head_len < tail_len {
209 src = source_deque.head;
210 dst = source_deque.to_wrapped_index(drain_len);
211 len = head_len;
212 } else {
213 src = source_deque.to_wrapped_index(head_len + drain_len);
214 dst = source_deque.to_wrapped_index(head_len);
215 len = tail_len;
216 };
217
218 // ignore-tidy-undocumented-unsafe
219 unsafe {
220 source_deque.wrap_copy(src, dst, len);
221 }
222 }
223 }
224
225 if new_len == 0 {
226 // Special case: If the entire deque was drained, reset the head back to 0,
227 // like `.clear()` does.
228 source_deque.head = WrappedIndex::zero();
229 } else if head_len < tail_len {
230 // If we moved the head above, then we need to adjust the head index here.
231 source_deque.head = source_deque.to_wrapped_index(drain_len);
232 }
233 source_deque.len = new_len;
234 }
235 }
236 }
237}
238
239#[stable(feature = "drain", since = "1.6.0")]
240impl<T, A: Allocator> Iterator for Drain<'_, T, A> {
241 type Item = T;
242
243 #[inline]
244 fn next(&mut self) -> Option<T> {
245 if self.remaining == 0 {
246 return None;
247 }
248 // ignore-tidy-undocumented-unsafe
249 let wrapped_idx = unsafe { self.deque.as_ref().to_wrapped_index(self.idx) };
250 self.idx += 1;
251 self.remaining -= 1;
252 // ignore-tidy-undocumented-unsafe
253 Some(unsafe { self.deque.as_mut().buffer_read(wrapped_idx) })
254 }
255
256 #[inline]
257 fn size_hint(&self) -> (usize, Option<usize>) {
258 let len = self.remaining;
259 (len, Some(len))
260 }
261}
262
263#[stable(feature = "drain", since = "1.6.0")]
264impl<T, A: Allocator> DoubleEndedIterator for Drain<'_, T, A> {
265 #[inline]
266 fn next_back(&mut self) -> Option<T> {
267 if self.remaining == 0 {
268 return None;
269 }
270 self.remaining -= 1;
271 let wrapped_idx =
272 // ignore-tidy-undocumented-unsafe
273 unsafe { self.deque.as_ref().to_wrapped_index(self.idx + self.remaining) };
274 // ignore-tidy-undocumented-unsafe
275 Some(unsafe { self.deque.as_mut().buffer_read(wrapped_idx) })
276 }
277}
278
279#[stable(feature = "drain", since = "1.6.0")]
280impl<T, A: Allocator> ExactSizeIterator for Drain<'_, T, A> {}
281
282#[stable(feature = "fused", since = "1.26.0")]
283impl<T, A: Allocator> FusedIterator for Drain<'_, T, A> {}