1 | use crate::alloc::{Allocator, Global}; |
2 | use core::fmt; |
3 | use core::iter::{FusedIterator, TrustedLen}; |
4 | use core::mem::{self, ManuallyDrop, SizedTypeProperties}; |
5 | use core::ptr::{self, NonNull}; |
6 | use core::slice::{self}; |
7 | |
8 | use super::Vec; |
9 | |
10 | /// A draining iterator for `Vec<T>`. |
11 | /// |
12 | /// This `struct` is created by [`Vec::drain`]. |
13 | /// See its documentation for more. |
14 | /// |
15 | /// # Example |
16 | /// |
17 | /// ``` |
18 | /// let mut v = vec![0, 1, 2]; |
19 | /// let iter: std::vec::Drain<'_, _> = v.drain(..); |
20 | /// ``` |
21 | #[stable (feature = "drain" , since = "1.6.0" )] |
22 | pub struct Drain< |
23 | 'a, |
24 | T: 'a, |
25 | #[unstable (feature = "allocator_api" , issue = "32838" )] A: Allocator + 'a = Global, |
26 | > { |
27 | /// Index of tail to preserve |
28 | pub(super) tail_start: usize, |
29 | /// Length of tail |
30 | pub(super) tail_len: usize, |
31 | /// Current remaining range to remove |
32 | pub(super) iter: slice::Iter<'a, T>, |
33 | pub(super) vec: NonNull<Vec<T, A>>, |
34 | } |
35 | |
36 | #[stable (feature = "collection_debug" , since = "1.17.0" )] |
37 | impl<T: fmt::Debug, A: Allocator> fmt::Debug for Drain<'_, T, A> { |
38 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
39 | f.debug_tuple(name:"Drain" ).field(&self.iter.as_slice()).finish() |
40 | } |
41 | } |
42 | |
43 | impl<'a, T, A: Allocator> Drain<'a, T, A> { |
44 | /// Returns the remaining items of this iterator as a slice. |
45 | /// |
46 | /// # Examples |
47 | /// |
48 | /// ``` |
49 | /// let mut vec = vec!['a' , 'b' , 'c' ]; |
50 | /// let mut drain = vec.drain(..); |
51 | /// assert_eq!(drain.as_slice(), &['a' , 'b' , 'c' ]); |
52 | /// let _ = drain.next().unwrap(); |
53 | /// assert_eq!(drain.as_slice(), &['b' , 'c' ]); |
54 | /// ``` |
55 | #[must_use ] |
56 | #[stable (feature = "vec_drain_as_slice" , since = "1.46.0" )] |
57 | pub fn as_slice(&self) -> &[T] { |
58 | self.iter.as_slice() |
59 | } |
60 | |
61 | /// Returns a reference to the underlying allocator. |
62 | #[unstable (feature = "allocator_api" , issue = "32838" )] |
63 | #[must_use ] |
64 | #[inline ] |
65 | pub fn allocator(&self) -> &A { |
66 | unsafe { self.vec.as_ref().allocator() } |
67 | } |
68 | |
69 | /// Keep unyielded elements in the source `Vec`. |
70 | /// |
71 | /// # Examples |
72 | /// |
73 | /// ``` |
74 | /// #![feature(drain_keep_rest)] |
75 | /// |
76 | /// let mut vec = vec!['a' , 'b' , 'c' ]; |
77 | /// let mut drain = vec.drain(..); |
78 | /// |
79 | /// assert_eq!(drain.next().unwrap(), 'a' ); |
80 | /// |
81 | /// // This call keeps 'b' and 'c' in the vec. |
82 | /// drain.keep_rest(); |
83 | /// |
84 | /// // If we wouldn't call `keep_rest()`, |
85 | /// // `vec` would be empty. |
86 | /// assert_eq!(vec, ['b' , 'c' ]); |
87 | /// ``` |
88 | #[unstable (feature = "drain_keep_rest" , issue = "101122" )] |
89 | pub fn keep_rest(self) { |
90 | // At this moment layout looks like this: |
91 | // |
92 | // [head] [yielded by next] [unyielded] [yielded by next_back] [tail] |
93 | // ^-- start \_________/-- unyielded_len \____/-- self.tail_len |
94 | // ^-- unyielded_ptr ^-- tail |
95 | // |
96 | // Normally `Drop` impl would drop [unyielded] and then move [tail] to the `start`. |
97 | // Here we want to |
98 | // 1. Move [unyielded] to `start` |
99 | // 2. Move [tail] to a new start at `start + len(unyielded)` |
100 | // 3. Update length of the original vec to `len(head) + len(unyielded) + len(tail)` |
101 | // a. In case of ZST, this is the only thing we want to do |
102 | // 4. Do *not* drop self, as everything is put in a consistent state already, there is nothing to do |
103 | let mut this = ManuallyDrop::new(self); |
104 | |
105 | unsafe { |
106 | let source_vec = this.vec.as_mut(); |
107 | |
108 | let start = source_vec.len(); |
109 | let tail = this.tail_start; |
110 | |
111 | let unyielded_len = this.iter.len(); |
112 | let unyielded_ptr = this.iter.as_slice().as_ptr(); |
113 | |
114 | // ZSTs have no identity, so we don't need to move them around. |
115 | if !T::IS_ZST { |
116 | let start_ptr = source_vec.as_mut_ptr().add(start); |
117 | |
118 | // memmove back unyielded elements |
119 | if unyielded_ptr != start_ptr { |
120 | let src = unyielded_ptr; |
121 | let dst = start_ptr; |
122 | |
123 | ptr::copy(src, dst, unyielded_len); |
124 | } |
125 | |
126 | // memmove back untouched tail |
127 | if tail != (start + unyielded_len) { |
128 | let src = source_vec.as_ptr().add(tail); |
129 | let dst = start_ptr.add(unyielded_len); |
130 | ptr::copy(src, dst, this.tail_len); |
131 | } |
132 | } |
133 | |
134 | source_vec.set_len(start + unyielded_len + this.tail_len); |
135 | } |
136 | } |
137 | } |
138 | |
139 | #[stable (feature = "vec_drain_as_slice" , since = "1.46.0" )] |
140 | impl<'a, T, A: Allocator> AsRef<[T]> for Drain<'a, T, A> { |
141 | fn as_ref(&self) -> &[T] { |
142 | self.as_slice() |
143 | } |
144 | } |
145 | |
146 | #[stable (feature = "drain" , since = "1.6.0" )] |
147 | unsafe impl<T: Sync, A: Sync + Allocator> Sync for Drain<'_, T, A> {} |
148 | #[stable (feature = "drain" , since = "1.6.0" )] |
149 | unsafe impl<T: Send, A: Send + Allocator> Send for Drain<'_, T, A> {} |
150 | |
151 | #[stable (feature = "drain" , since = "1.6.0" )] |
152 | impl<T, A: Allocator> Iterator for Drain<'_, T, A> { |
153 | type Item = T; |
154 | |
155 | #[inline ] |
156 | fn next(&mut self) -> Option<T> { |
157 | self.iter.next().map(|elt: &T| unsafe { ptr::read(src:elt as *const _) }) |
158 | } |
159 | |
160 | fn size_hint(&self) -> (usize, Option<usize>) { |
161 | self.iter.size_hint() |
162 | } |
163 | } |
164 | |
165 | #[stable (feature = "drain" , since = "1.6.0" )] |
166 | impl<T, A: Allocator> DoubleEndedIterator for Drain<'_, T, A> { |
167 | #[inline ] |
168 | fn next_back(&mut self) -> Option<T> { |
169 | self.iter.next_back().map(|elt: &T| unsafe { ptr::read(src:elt as *const _) }) |
170 | } |
171 | } |
172 | |
173 | #[stable (feature = "drain" , since = "1.6.0" )] |
174 | impl<T, A: Allocator> Drop for Drain<'_, T, A> { |
175 | fn drop(&mut self) { |
176 | /// Moves back the un-`Drain`ed elements to restore the original `Vec`. |
177 | struct DropGuard<'r, 'a, T, A: Allocator>(&'r mut Drain<'a, T, A>); |
178 | |
179 | impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> { |
180 | fn drop(&mut self) { |
181 | if self.0.tail_len > 0 { |
182 | unsafe { |
183 | let source_vec = self.0.vec.as_mut(); |
184 | // memmove back untouched tail, update to new length |
185 | let start = source_vec.len(); |
186 | let tail = self.0.tail_start; |
187 | if tail != start { |
188 | let src = source_vec.as_ptr().add(tail); |
189 | let dst = source_vec.as_mut_ptr().add(start); |
190 | ptr::copy(src, dst, self.0.tail_len); |
191 | } |
192 | source_vec.set_len(start + self.0.tail_len); |
193 | } |
194 | } |
195 | } |
196 | } |
197 | |
198 | let iter = mem::take(&mut self.iter); |
199 | let drop_len = iter.len(); |
200 | |
201 | let mut vec = self.vec; |
202 | |
203 | if T::IS_ZST { |
204 | // ZSTs have no identity, so we don't need to move them around, we only need to drop the correct amount. |
205 | // this can be achieved by manipulating the Vec length instead of moving values out from `iter`. |
206 | unsafe { |
207 | let vec = vec.as_mut(); |
208 | let old_len = vec.len(); |
209 | vec.set_len(old_len + drop_len + self.tail_len); |
210 | vec.truncate(old_len + self.tail_len); |
211 | } |
212 | |
213 | return; |
214 | } |
215 | |
216 | // ensure elements are moved back into their appropriate places, even when drop_in_place panics |
217 | let _guard = DropGuard(self); |
218 | |
219 | if drop_len == 0 { |
220 | return; |
221 | } |
222 | |
223 | // as_slice() must only be called when iter.len() is > 0 because |
224 | // it also gets touched by vec::Splice which may turn it into a dangling pointer |
225 | // which would make it and the vec pointer point to different allocations which would |
226 | // lead to invalid pointer arithmetic below. |
227 | let drop_ptr = iter.as_slice().as_ptr(); |
228 | |
229 | unsafe { |
230 | // drop_ptr comes from a slice::Iter which only gives us a &[T] but for drop_in_place |
231 | // a pointer with mutable provenance is necessary. Therefore we must reconstruct |
232 | // it from the original vec but also avoid creating a &mut to the front since that could |
233 | // invalidate raw pointers to it which some unsafe code might rely on. |
234 | let vec_ptr = vec.as_mut().as_mut_ptr(); |
235 | let drop_offset = drop_ptr.sub_ptr(vec_ptr); |
236 | let to_drop = ptr::slice_from_raw_parts_mut(vec_ptr.add(drop_offset), drop_len); |
237 | ptr::drop_in_place(to_drop); |
238 | } |
239 | } |
240 | } |
241 | |
242 | #[stable (feature = "drain" , since = "1.6.0" )] |
243 | impl<T, A: Allocator> ExactSizeIterator for Drain<'_, T, A> { |
244 | fn is_empty(&self) -> bool { |
245 | self.iter.is_empty() |
246 | } |
247 | } |
248 | |
249 | #[unstable (feature = "trusted_len" , issue = "37572" )] |
250 | unsafe impl<T, A: Allocator> TrustedLen for Drain<'_, T, A> {} |
251 | |
252 | #[stable (feature = "fused" , since = "1.26.0" )] |
253 | impl<T, A: Allocator> FusedIterator for Drain<'_, T, A> {} |
254 | |