rustc_data_structures/
thinvec.rs

Help
1//! This is a copy-paste of `Vec::extract_if` for `ThinVec`.
2//!
3//! FIXME: <https://github.com/Gankra/thin-vec/pull/66> is merged, this can be removed.
4
5use std::{ptr, slice};
6
7use thin_vec::ThinVec;
8
9/// An iterator for [`ThinVec`] which uses a closure to determine if an element should be removed.
10#[must_use = "iterators are lazy and do nothing unless consumed"]
11pub struct ExtractIf<'a, T, F> {
12    vec: &'a mut ThinVec<T>,
13    /// The index of the item that will be inspected by the next call to `next`.
14    idx: usize,
15    /// The number of items that have been drained (removed) thus far.
16    del: usize,
17    /// The original length of `vec` prior to draining.
18    old_len: usize,
19    /// The filter test predicate.
20    pred: F,
21}
22
23impl<'a, T, F> ExtractIf<'a, T, F>
24where
25    F: FnMut(&mut T) -> bool,
26{
27    pub fn new(vec: &'a mut ThinVec<T>, filter: F) -> Self {
28        let old_len = vec.len();
29
30        // Guard against us getting leaked (leak amplification)
31        unsafe {
32            vec.set_len(0);
33        }
34
35        ExtractIf { vec, idx: 0, del: 0, old_len, pred: filter }
36    }
37}
38
39impl<T, F> Iterator for ExtractIf<'_, T, F>
40where
41    F: FnMut(&mut T) -> bool,
42{
43    type Item = T;
44    fn next(&mut self) -> Option<Self::Item> {
45        unsafe {
46            while self.idx < self.old_len {
47                let i = self.idx;
48                let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.old_len);
49                let drained = (self.pred)(&mut v[i]);
50                // Update the index *after* the predicate is called. If the index
51                // is updated prior and the predicate panics, the element at this
52                // index would be leaked.
53                self.idx += 1;
54                if drained {
55                    self.del += 1;
56                    return Some(ptr::read(&v[i]));
57                } else if self.del > 0 {
58                    let del = self.del;
59                    let src: *const T = &v[i];
60                    let dst: *mut T = &mut v[i - del];
61                    ptr::copy_nonoverlapping(src, dst, 1);
62                }
63            }
64            None
65        }
66    }
67
68    fn size_hint(&self) -> (usize, Option<usize>) {
69        (0, Some(self.old_len - self.idx))
70    }
71}
72
73impl<A, F> Drop for ExtractIf<'_, A, F> {
74    fn drop(&mut self) {
75        unsafe {
76            if self.idx < self.old_len && self.del > 0 {
77                // This is a pretty messed up state, and there isn't really an
78                // obviously right thing to do. We don't want to keep trying
79                // to execute `pred`, so we just backshift all the unprocessed
80                // elements and tell the vec that they still exist. The backshift
81                // is required to prevent a double-drop of the last successfully
82                // drained item prior to a panic in the predicate.
83                let ptr = self.vec.as_mut_ptr();
84                let src = ptr.add(self.idx);
85                let dst = src.sub(self.del);
86                let tail_len = self.old_len - self.idx;
87                src.copy_to(dst, tail_len);
88            }
89            self.vec.set_len(self.old_len - self.del);
90        }
91    }
92}
rustc_data_structures/thinvec.rs

rustc_data_structures/
thinvec.rs
