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 crate::alloc::{Allocator, Global}; |
8 | |
9 | use super::VecDeque; |
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" )] |
18 | pub 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 | deque: NonNull<VecDeque<T, A>>, |
26 | // drain_start is stored in deque.len |
27 | drain_len: usize, |
28 | // index into the logical array, not the physical one (always lies in [0..deque.len)) |
29 | idx: usize, |
30 | // number of elements after the drain range |
31 | tail_len: usize, |
32 | remaining: usize, |
33 | // Needed to make Drain covariant over T |
34 | _marker: PhantomData<&'a T>, |
35 | } |
36 | |
37 | impl<'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 | unsafe fn as_slices(&self) -> (*mut [T], *mut [T]) { |
58 | unsafe { |
59 | let deque = self.deque.as_ref(); |
60 | |
61 | // We know that `self.idx + self.remaining <= deque.len <= usize::MAX`, so this won't overflow. |
62 | let logical_remaining_range = self.idx..self.idx + self.remaining; |
63 | |
64 | // SAFETY: `logical_remaining_range` represents the |
65 | // range into the logical buffer of elements that |
66 | // haven't been drained yet, so they're all initialized, |
67 | // and `slice::range(start..end, end) == start..end`, |
68 | // so the preconditions for `slice_ranges` are met. |
69 | let (a_range, b_range) = |
70 | deque.slice_ranges(logical_remaining_range.clone(), logical_remaining_range.end); |
71 | (deque.buffer_range(a_range), deque.buffer_range(b_range)) |
72 | } |
73 | } |
74 | } |
75 | |
76 | #[stable (feature = "collection_debug" , since = "1.17.0" )] |
77 | impl<T: fmt::Debug, A: Allocator> fmt::Debug for Drain<'_, T, A> { |
78 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
79 | f&mut DebugTuple<'_, '_>.debug_tuple(name:"Drain" ) |
80 | .field(&self.drain_len) |
81 | .field(&self.idx) |
82 | .field(&self.tail_len) |
83 | .field(&self.remaining) |
84 | .finish() |
85 | } |
86 | } |
87 | |
88 | #[stable (feature = "drain" , since = "1.6.0" )] |
89 | unsafe impl<T: Sync, A: Allocator + Sync> Sync for Drain<'_, T, A> {} |
90 | #[stable (feature = "drain" , since = "1.6.0" )] |
91 | unsafe impl<T: Send, A: Allocator + Send> Send for Drain<'_, T, A> {} |
92 | |
93 | #[stable (feature = "drain" , since = "1.6.0" )] |
94 | impl<T, A: Allocator> Drop for Drain<'_, T, A> { |
95 | fn drop(&mut self) { |
96 | struct DropGuard<'r, 'a, T, A: Allocator>(&'r mut Drain<'a, T, A>); |
97 | |
98 | impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> { |
99 | fn drop(&mut self) { |
100 | if self.0.remaining != 0 { |
101 | unsafe { |
102 | // SAFETY: We just checked that `self.remaining != 0`. |
103 | let (front, back) = self.0.as_slices(); |
104 | ptr::drop_in_place(front); |
105 | ptr::drop_in_place(back); |
106 | } |
107 | } |
108 | |
109 | let source_deque = unsafe { self.0.deque.as_mut() }; |
110 | |
111 | let drain_start = source_deque.len(); |
112 | 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; |
116 | |
117 | if T::IS_ZST { |
118 | // no need to copy around any memory if T is a ZST |
119 | source_deque.len = orig_len - drain_len; |
120 | return; |
121 | } |
122 | |
123 | let head_len = drain_start; |
124 | let tail_len = self.0.tail_len; |
125 | |
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; |
147 | } 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; |
154 | } |
155 | }, |
156 | } |
157 | } |
158 | } |
159 | |
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); |
171 | } |
172 | } |
173 | |
174 | // Dropping `guard` handles moving the remaining elements into place. |
175 | } |
176 | } |
177 | |
178 | #[stable (feature = "drain" , since = "1.6.0" )] |
179 | impl<T, A: Allocator> Iterator for Drain<'_, T, A> { |
180 | type Item = T; |
181 | |
182 | #[inline ] |
183 | fn next(&mut self) -> Option<T> { |
184 | if self.remaining == 0 { |
185 | return None; |
186 | } |
187 | let wrapped_idx: usize = unsafe { self.deque.as_ref().to_physical_idx(self.idx) }; |
188 | self.idx += 1; |
189 | self.remaining -= 1; |
190 | Some(unsafe { self.deque.as_mut().buffer_read(off:wrapped_idx) }) |
191 | } |
192 | |
193 | #[inline ] |
194 | fn size_hint(&self) -> (usize, Option<usize>) { |
195 | let len: usize = self.remaining; |
196 | (len, Some(len)) |
197 | } |
198 | } |
199 | |
200 | #[stable (feature = "drain" , since = "1.6.0" )] |
201 | impl<T, A: Allocator> DoubleEndedIterator for Drain<'_, T, A> { |
202 | #[inline ] |
203 | fn next_back(&mut self) -> Option<T> { |
204 | if self.remaining == 0 { |
205 | return None; |
206 | } |
207 | self.remaining -= 1; |
208 | let wrapped_idx: usize = unsafe { self.deque.as_ref().to_physical_idx(self.idx + self.remaining) }; |
209 | Some(unsafe { self.deque.as_mut().buffer_read(off:wrapped_idx) }) |
210 | } |
211 | } |
212 | |
213 | #[stable (feature = "drain" , since = "1.6.0" )] |
214 | impl<T, A: Allocator> ExactSizeIterator for Drain<'_, T, A> {} |
215 | |
216 | #[stable (feature = "fused" , since = "1.26.0" )] |
217 | impl<T, A: Allocator> FusedIterator for Drain<'_, T, A> {} |
218 | |