1 | use crate::alloc::{Allocator, Global}; |
2 | use core::ptr; |
3 | use core::slice; |
4 | |
5 | use super::Vec; |
6 | |
7 | /// An iterator which uses a closure to determine if an element should be removed. |
8 | /// |
9 | /// This struct is created by [`Vec::extract_if`]. |
10 | /// See its documentation for more. |
11 | /// |
12 | /// # Example |
13 | /// |
14 | /// ``` |
15 | /// #![feature(extract_if)] |
16 | /// |
17 | /// let mut v = vec![0, 1, 2]; |
18 | /// let iter: std::vec::ExtractIf<'_, _, _> = v.extract_if(|x| *x % 2 == 0); |
19 | /// ``` |
20 | #[unstable (feature = "extract_if" , reason = "recently added" , issue = "43244" )] |
21 | #[derive (Debug)] |
22 | #[must_use = "iterators are lazy and do nothing unless consumed" ] |
23 | pub struct ExtractIf< |
24 | 'a, |
25 | T, |
26 | F, |
27 | #[unstable (feature = "allocator_api" , issue = "32838" )] A: Allocator = Global, |
28 | > where |
29 | F: FnMut(&mut T) -> bool, |
30 | { |
31 | pub(super) vec: &'a mut Vec<T, A>, |
32 | /// The index of the item that will be inspected by the next call to `next`. |
33 | pub(super) idx: usize, |
34 | /// The number of items that have been drained (removed) thus far. |
35 | pub(super) del: usize, |
36 | /// The original length of `vec` prior to draining. |
37 | pub(super) old_len: usize, |
38 | /// The filter test predicate. |
39 | pub(super) pred: F, |
40 | } |
41 | |
42 | impl<T, F, A: Allocator> ExtractIf<'_, T, F, A> |
43 | where |
44 | F: FnMut(&mut T) -> bool, |
45 | { |
46 | /// Returns a reference to the underlying allocator. |
47 | #[unstable (feature = "allocator_api" , issue = "32838" )] |
48 | #[inline ] |
49 | pub fn allocator(&self) -> &A { |
50 | self.vec.allocator() |
51 | } |
52 | } |
53 | |
54 | #[unstable (feature = "extract_if" , reason = "recently added" , issue = "43244" )] |
55 | impl<T, F, A: Allocator> Iterator for ExtractIf<'_, T, F, A> |
56 | where |
57 | F: FnMut(&mut T) -> bool, |
58 | { |
59 | type Item = T; |
60 | |
61 | fn next(&mut self) -> Option<T> { |
62 | unsafe { |
63 | while self.idx < self.old_len { |
64 | let i = self.idx; |
65 | let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len); |
66 | let drained = (self.pred)(&mut v[i]); |
67 | // Update the index *after* the predicate is called. If the index |
68 | // is updated prior and the predicate panics, the element at this |
69 | // index would be leaked. |
70 | self.idx += 1; |
71 | if drained { |
72 | self.del += 1; |
73 | return Some(ptr::read(&v[i])); |
74 | } else if self.del > 0 { |
75 | let del = self.del; |
76 | let src: *const T = &v[i]; |
77 | let dst: *mut T = &mut v[i - del]; |
78 | ptr::copy_nonoverlapping(src, dst, 1); |
79 | } |
80 | } |
81 | None |
82 | } |
83 | } |
84 | |
85 | fn size_hint(&self) -> (usize, Option<usize>) { |
86 | (0, Some(self.old_len - self.idx)) |
87 | } |
88 | } |
89 | |
90 | #[unstable (feature = "extract_if" , reason = "recently added" , issue = "43244" )] |
91 | impl<T, F, A: Allocator> Drop for ExtractIf<'_, T, F, A> |
92 | where |
93 | F: FnMut(&mut T) -> bool, |
94 | { |
95 | fn drop(&mut self) { |
96 | unsafe { |
97 | if self.idx < self.old_len && self.del > 0 { |
98 | // This is a pretty messed up state, and there isn't really an |
99 | // obviously right thing to do. We don't want to keep trying |
100 | // to execute `pred`, so we just backshift all the unprocessed |
101 | // elements and tell the vec that they still exist. The backshift |
102 | // is required to prevent a double-drop of the last successfully |
103 | // drained item prior to a panic in the predicate. |
104 | let ptr: *mut T = self.vec.as_mut_ptr(); |
105 | let src: *mut T = ptr.add(self.idx); |
106 | let dst: *mut T = src.sub(self.del); |
107 | let tail_len: usize = self.old_len - self.idx; |
108 | src.copy_to(dest:dst, count:tail_len); |
109 | } |
110 | self.vec.set_len(self.old_len - self.del); |
111 | } |
112 | } |
113 | } |
114 | |