1 // SPDX-License-Identifier: Apache-2.0 OR MIT
2 
3 use crate::alloc::{Allocator, Global};
4 use core::ptr::{self};
5 use core::slice::{self};
6 
7 use super::Vec;
8 
9 /// An iterator which uses a closure to determine if an element should be removed.
10 ///
11 /// This struct is created by [`Vec::drain_filter`].
12 /// See its documentation for more.
13 ///
14 /// # Example
15 ///
16 /// ```
17 /// #![feature(drain_filter)]
18 ///
19 /// let mut v = vec![0, 1, 2];
20 /// let iter: std::vec::DrainFilter<_, _> = v.drain_filter(|x| *x % 2 == 0);
21 /// ```
22 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
23 #[derive(Debug)]
24 pub struct DrainFilter<
25     'a,
26     T,
27     F,
28     #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
29 > where
30     F: FnMut(&mut T) -> bool,
31 {
32     pub(super) vec: &'a mut Vec<T, A>,
33     /// The index of the item that will be inspected by the next call to `next`.
34     pub(super) idx: usize,
35     /// The number of items that have been drained (removed) thus far.
36     pub(super) del: usize,
37     /// The original length of `vec` prior to draining.
38     pub(super) old_len: usize,
39     /// The filter test predicate.
40     pub(super) pred: F,
41     /// A flag that indicates a panic has occurred in the filter test predicate.
42     /// This is used as a hint in the drop implementation to prevent consumption
43     /// of the remainder of the `DrainFilter`. Any unprocessed items will be
44     /// backshifted in the `vec`, but no further items will be dropped or
45     /// tested by the filter predicate.
46     pub(super) panic_flag: bool,
47 }
48 
49 impl<T, F, A: Allocator> DrainFilter<'_, T, F, A>
50 where
51     F: FnMut(&mut T) -> bool,
52 {
53     /// Returns a reference to the underlying allocator.
54     #[unstable(feature = "allocator_api", issue = "32838")]
55     #[inline]
allocator(&self) -> &A56     pub fn allocator(&self) -> &A {
57         self.vec.allocator()
58     }
59 }
60 
61 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
62 impl<T, F, A: Allocator> Iterator for DrainFilter<'_, T, F, A>
63 where
64     F: FnMut(&mut T) -> bool,
65 {
66     type Item = T;
67 
next(&mut self) -> Option<T>68     fn next(&mut self) -> Option<T> {
69         unsafe {
70             while self.idx < self.old_len {
71                 let i = self.idx;
72                 let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
73                 self.panic_flag = true;
74                 let drained = (self.pred)(&mut v[i]);
75                 self.panic_flag = false;
76                 // Update the index *after* the predicate is called. If the index
77                 // is updated prior and the predicate panics, the element at this
78                 // index would be leaked.
79                 self.idx += 1;
80                 if drained {
81                     self.del += 1;
82                     return Some(ptr::read(&v[i]));
83                 } else if self.del > 0 {
84                     let del = self.del;
85                     let src: *const T = &v[i];
86                     let dst: *mut T = &mut v[i - del];
87                     ptr::copy_nonoverlapping(src, dst, 1);
88                 }
89             }
90             None
91         }
92     }
93 
size_hint(&self) -> (usize, Option<usize>)94     fn size_hint(&self) -> (usize, Option<usize>) {
95         (0, Some(self.old_len - self.idx))
96     }
97 }
98 
99 #[unstable(feature = "drain_filter", reason = "recently added", issue = "43244")]
100 impl<T, F, A: Allocator> Drop for DrainFilter<'_, T, F, A>
101 where
102     F: FnMut(&mut T) -> bool,
103 {
drop(&mut self)104     fn drop(&mut self) {
105         struct BackshiftOnDrop<'a, 'b, T, F, A: Allocator>
106         where
107             F: FnMut(&mut T) -> bool,
108         {
109             drain: &'b mut DrainFilter<'a, T, F, A>,
110         }
111 
112         impl<'a, 'b, T, F, A: Allocator> Drop for BackshiftOnDrop<'a, 'b, T, F, A>
113         where
114             F: FnMut(&mut T) -> bool,
115         {
116             fn drop(&mut self) {
117                 unsafe {
118                     if self.drain.idx < self.drain.old_len && self.drain.del > 0 {
119                         // This is a pretty messed up state, and there isn't really an
120                         // obviously right thing to do. We don't want to keep trying
121                         // to execute `pred`, so we just backshift all the unprocessed
122                         // elements and tell the vec that they still exist. The backshift
123                         // is required to prevent a double-drop of the last successfully
124                         // drained item prior to a panic in the predicate.
125                         let ptr = self.drain.vec.as_mut_ptr();
126                         let src = ptr.add(self.drain.idx);
127                         let dst = src.sub(self.drain.del);
128                         let tail_len = self.drain.old_len - self.drain.idx;
129                         src.copy_to(dst, tail_len);
130                     }
131                     self.drain.vec.set_len(self.drain.old_len - self.drain.del);
132                 }
133             }
134         }
135 
136         let backshift = BackshiftOnDrop { drain: self };
137 
138         // Attempt to consume any remaining elements if the filter predicate
139         // has not yet panicked. We'll backshift any remaining elements
140         // whether we've already panicked or if the consumption here panics.
141         if !backshift.drain.panic_flag {
142             backshift.drain.for_each(drop);
143         }
144     }
145 }
146