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