alloc/vec/
extract_if.rs

1use core::ops::{Range, RangeBounds};
2use core::{fmt, ptr, slice};
3
4use super::Vec;
5use crate::alloc::{Allocator, Global};
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/// let mut v = vec![0, 1, 2];
16/// let iter: std::vec::ExtractIf<'_, _, _> = v.extract_if(.., |x| *x % 2 == 0);
17/// ```
18#[stable(feature = "extract_if", since = "1.87.0")]
19#[must_use = "iterators are lazy and do nothing unless consumed"]
20pub struct ExtractIf<
21    'a,
22    T,
23    F,
24    #[unstable(feature = "allocator_api", issue = "32838")] A: Allocator = Global,
25> {
26    vec: &'a mut Vec<T, A>,
27    /// The index of the item that will be inspected by the next call to `next`.
28    idx: usize,
29    /// Elements at and beyond this point will be retained. Must be equal or smaller than `old_len`.
30    end: usize,
31    /// The number of items that have been drained (removed) thus far.
32    del: usize,
33    /// The original length of `vec` prior to draining.
34    old_len: usize,
35    /// The filter test predicate.
36    pred: F,
37}
38
39impl<'a, T, F, A: Allocator> ExtractIf<'a, T, F, A> {
40    pub(super) fn new<R: RangeBounds<usize>>(vec: &'a mut Vec<T, A>, pred: F, range: R) -> Self {
41        let old_len = vec.len();
42        let Range { start, end } = slice::range(range, ..old_len);
43
44        // Guard against the vec getting leaked (leak amplification)
45        unsafe {
46            vec.set_len(0);
47        }
48        ExtractIf { vec, idx: start, del: 0, end, old_len, pred }
49    }
50
51    /// Returns a reference to the underlying allocator.
52    #[unstable(feature = "allocator_api", issue = "32838")]
53    #[inline]
54    pub fn allocator(&self) -> &A {
55        self.vec.allocator()
56    }
57}
58
59#[stable(feature = "extract_if", since = "1.87.0")]
60impl<T, F, A: Allocator> Iterator for ExtractIf<'_, T, F, A>
61where
62    F: FnMut(&mut T) -> bool,
63{
64    type Item = T;
65
66    fn next(&mut self) -> Option<T> {
67        while self.idx < self.end {
68            let i = self.idx;
69            // SAFETY:
70            //  We know that `i < self.end` from the if guard and that `self.end <= self.old_len` from
71            //  the validity of `Self`. Therefore `i` points to an element within `vec`.
72            //
73            //  Additionally, the i-th element is valid because each element is visited at most once
74            //  and it is the first time we access vec[i].
75            //
76            //  Note: we can't use `vec.get_unchecked_mut(i)` here since the precondition for that
77            //  function is that i < vec.len(), but we've set vec's length to zero.
78            let cur = unsafe { &mut *self.vec.as_mut_ptr().add(i) };
79            let drained = (self.pred)(cur);
80            // Update the index *after* the predicate is called. If the index
81            // is updated prior and the predicate panics, the element at this
82            // index would be leaked.
83            self.idx += 1;
84            if drained {
85                self.del += 1;
86                // SAFETY: We never touch this element again after returning it.
87                return Some(unsafe { ptr::read(cur) });
88            } else if self.del > 0 {
89                // SAFETY: `self.del` > 0, so the hole slot must not overlap with current element.
90                // We use copy for move, and never touch this element again.
91                unsafe {
92                    let hole_slot = self.vec.as_mut_ptr().add(i - self.del);
93                    ptr::copy_nonoverlapping(cur, hole_slot, 1);
94                }
95            }
96        }
97        None
98    }
99
100    fn size_hint(&self) -> (usize, Option<usize>) {
101        (0, Some(self.end - self.idx))
102    }
103}
104
105#[stable(feature = "extract_if", since = "1.87.0")]
106impl<T, F, A: Allocator> Drop for ExtractIf<'_, T, F, A> {
107    fn drop(&mut self) {
108        if self.del > 0 {
109            // SAFETY: Trailing unchecked items must be valid since we never touch them.
110            unsafe {
111                ptr::copy(
112                    self.vec.as_ptr().add(self.idx),
113                    self.vec.as_mut_ptr().add(self.idx - self.del),
114                    self.old_len - self.idx,
115                );
116            }
117        }
118        // SAFETY: After filling holes, all items are in contiguous memory.
119        unsafe {
120            self.vec.set_len(self.old_len - self.del);
121        }
122    }
123}
124
125#[stable(feature = "extract_if", since = "1.87.0")]
126impl<T, F, A> fmt::Debug for ExtractIf<'_, T, F, A>
127where
128    T: fmt::Debug,
129    A: Allocator,
130{
131    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
132        let peek = if self.idx < self.end { self.vec.get(self.idx) } else { None };
133        f.debug_struct("ExtractIf").field("peek", &peek).finish_non_exhaustive()
134    }
135}
alloc/vec/extract_if.rs

alloc/vec/
extract_if.rs
