1use crate::alloc::{Allocator, Global};
2use core::ptr;
3use core::slice;
4
5use 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"]
23pub 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
42impl<T, F, A: Allocator> ExtractIf<'_, T, F, A>
43where
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")]
55impl<T, F, A: Allocator> Iterator for ExtractIf<'_, T, F, A>
56where
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")]
91impl<T, F, A: Allocator> Drop for ExtractIf<'_, T, F, A>
92where
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