1 use crate::alloc::{Allocator, Global};
3 use core::slice::{self};
7 /// An iterator which uses a closure to determine if an element should be removed.
9 /// This struct is created by [`Vec::drain_filter`].
10 /// See its documentation for more.
15 /// #![feature(drain_filter)]
17 /// let mut v = vec![0, 1, 2];
18 /// let iter: std::vec::DrainFilter<_, _> = v.drain_filter(|x| *x % 2 == 0);
20 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
22 pub struct DrainFilter<
26 #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
28 F: FnMut(&mut T) -> bool,
30 pub(super) vec: &'a mut Vec<T, A>,
31 /// The index of the item that will be inspected by the next call to `next`.
32 pub(super) idx: usize,
33 /// The number of items that have been drained (removed) thus far.
34 pub(super) del: usize,
35 /// The original length of `vec` prior to draining.
36 pub(super) old_len: usize,
37 /// The filter test predicate.
39 /// A flag that indicates a panic has occurred in the filter test predicate.
40 /// This is used as a hint in the drop implementation to prevent consumption
41 /// of the remainder of the `DrainFilter`. Any unprocessed items will be
42 /// backshifted in the `vec`, but no further items will be dropped or
43 /// tested by the filter predicate.
44 pub(super) panic_flag: bool,
47 impl<T, F, A: Allocator> DrainFilter<'_, T, F, A>
49 F: FnMut(&mut T) -> bool,
51 /// Returns a reference to the underlying allocator.
52 #[unstable(feature = "allocator_api", issue = "32838")]
54 pub fn allocator(&self) -> &A {
59 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
60 impl<T, F, A: Allocator> Iterator for DrainFilter<'_, T, F, A>
62 F: FnMut(&mut T) -> bool,
66 fn next(&mut self) -> Option<T> {
68 while self.idx < self.old_len {
70 let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
71 self.panic_flag = true;
72 let drained = (self.pred)(&mut v[i]);
73 self.panic_flag = false;
74 // Update the index *after* the predicate is called. If the index
75 // is updated prior and the predicate panics, the element at this
76 // index would be leaked.
80 return Some(ptr::read(&v[i]));
81 } else if self.del > 0 {
83 let src: *const T = &v[i];
84 let dst: *mut T = &mut v[i - del];
85 ptr::copy_nonoverlapping(src, dst, 1);
92 fn size_hint(&self) -> (usize, Option<usize>) {
93 (0, Some(self.old_len - self.idx))
97 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
98 impl<T, F, A: Allocator> Drop for DrainFilter<'_, T, F, A>
100 F: FnMut(&mut T) -> bool,
103 struct BackshiftOnDrop<'a, 'b, T, F, A: Allocator>
105 F: FnMut(&mut T) -> bool,
107 drain: &'b mut DrainFilter<'a, T, F, A>,
110 impl<'a, 'b, T, F, A: Allocator> Drop for BackshiftOnDrop<'a, 'b, T, F, A>
112 F: FnMut(&mut T) -> bool,
116 if self.drain.idx < self.drain.old_len && self.drain.del > 0 {
117 // This is a pretty messed up state, and there isn't really an
118 // obviously right thing to do. We don't want to keep trying
119 // to execute `pred`, so we just backshift all the unprocessed
120 // elements and tell the vec that they still exist. The backshift
121 // is required to prevent a double-drop of the last successfully
122 // drained item prior to a panic in the predicate.
123 let ptr = self.drain.vec.as_mut_ptr();
124 let src = ptr.add(self.drain.idx);
125 let dst = src.sub(self.drain.del);
126 let tail_len = self.drain.old_len - self.drain.idx;
127 src.copy_to(dst, tail_len);
129 self.drain.vec.set_len(self.drain.old_len - self.drain.del);
134 let backshift = BackshiftOnDrop { drain: self };
136 // Attempt to consume any remaining elements if the filter predicate
137 // has not yet panicked. We'll backshift any remaining elements
138 // whether we've already panicked or if the consumption here panics.
139 if !backshift.drain.panic_flag {
140 backshift.drain.for_each(drop);