alloc/vec/
drain.rs

1use core::fmt;
2use core::iter::{FusedIterator, TrustedLen};
3use core::mem::{self, ManuallyDrop, SizedTypeProperties};
4use core::ptr::{self, NonNull};
5use core::slice::{self};
6
7use super::Vec;
8use crate::alloc::{Allocator, Global};
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")]
22pub 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")]
37impl<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("Drain").field(&self.iter.as_slice()).finish()
40    }
41}
42
43impl<'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")]
140impl<'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")]
147unsafe impl<T: Sync, A: Sync + Allocator> Sync for Drain<'_, T, A> {}
148#[stable(feature = "drain", since = "1.6.0")]
149unsafe impl<T: Send, A: Send + Allocator> Send for Drain<'_, T, A> {}
150
151#[stable(feature = "drain", since = "1.6.0")]
152impl<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| unsafe { ptr::read(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")]
166impl<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| unsafe { ptr::read(elt as *const _) })
170    }
171}
172
173#[stable(feature = "drain", since = "1.6.0")]
174impl<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.offset_from_unsigned(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")]
243impl<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")]
250unsafe impl<T, A: Allocator> TrustedLen for Drain<'_, T, A> {}
251
252#[stable(feature = "fused", since = "1.26.0")]
253impl<T, A: Allocator> FusedIterator for Drain<'_, T, A> {}
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy