bump_scope/owned_slice/
extract_if.rs

1use core::ptr::NonNull;
2
3use crate::{BumpBox, polyfill::non_null};
4
5/// An iterator which uses a closure to determine if an element should be removed.
6///
7/// This struct is created by the `extract_if` method on
8/// [`BumpBox`](BumpBox::extract_if),
9/// [`FixedBumpVec`](crate::FixedBumpVec::extract_if),
10/// [`BumpVec`](crate::BumpVec::extract_if) and
11/// [`MutBumpVec`](crate::MutBumpVec::extract_if).
12///
13/// See their documentation for more.
14///
15/// # Example
16///
17/// ```
18/// use bump_scope::{Bump, owned_slice::ExtractIf};
19/// let bump: Bump = Bump::new();
20///
21/// let mut v = bump.alloc_slice_copy(&[0, 1, 2]);
22/// let iter: ExtractIf<'_, _, _> = v.extract_if(|x| *x % 2 == 0);
23/// # _ = iter;
24/// ```
25#[derive(Debug)]
26#[must_use = "iterators are lazy and do nothing unless consumed"]
27pub struct ExtractIf<'a, T, F>
28where
29    F: FnMut(&mut T) -> bool,
30{
31    ptr: &'a mut NonNull<[T]>,
32    index: usize,
33    drained_count: usize,
34    original_len: usize,
35    filter: F,
36}
37
38impl<'a, T, F> ExtractIf<'a, T, F>
39where
40    F: FnMut(&mut T) -> bool,
41{
42    pub(crate) fn new<'a2>(boxed: &'a mut BumpBox<'a2, [T]>, filter: F) -> Self {
43        // When the ExtractIf is first created, it shortens the length of
44        // the source boxed slice to make sure no uninitialized or moved-from elements
45        // are accessible at all if the ExtractIf's destructor never gets to run.
46        //
47        // The 'a2 lifetime is shortened to 'a even though &'b mut BumpBox<'a2> is invariant
48        // over 'a2. We are careful not to expose any api where that could cause issues.
49
50        let ptr = unsafe { boxed.mut_ptr() };
51        let len = ptr.len();
52
53        non_null::set_len(ptr, 0);
54
55        Self {
56            ptr,
57            index: 0,
58            drained_count: 0,
59            original_len: len,
60            filter,
61        }
62    }
63}
64
65impl<T, F> Iterator for ExtractIf<'_, T, F>
66where
67    F: FnMut(&mut T) -> bool,
68{
69    type Item = T;
70
71    fn next(&mut self) -> Option<T> {
72        unsafe {
73            while self.index < self.original_len {
74                let start_ptr = non_null::as_non_null_ptr(*self.ptr);
75                let mut value_ptr = start_ptr.add(self.index);
76
77                let drained = (self.filter)(value_ptr.as_mut());
78
79                // Update the index *after* the predicate is called. If the index
80                // is updated prior and the predicate panics, the element at this
81                // index would be leaked.
82                self.index += 1;
83
84                if drained {
85                    self.drained_count += 1;
86                    return Some(value_ptr.as_ptr().read());
87                } else if self.drained_count > 0 {
88                    let src = value_ptr;
89                    let dst = value_ptr.sub(self.drained_count);
90                    src.copy_to_nonoverlapping(dst, 1);
91                }
92            }
93            None
94        }
95    }
96
97    #[inline]
98    fn size_hint(&self) -> (usize, Option<usize>) {
99        (0, Some(self.original_len - self.index))
100    }
101}
102
103impl<T, F> Drop for ExtractIf<'_, T, F>
104where
105    F: FnMut(&mut T) -> bool,
106{
107    fn drop(&mut self) {
108        unsafe {
109            if self.index < self.original_len && self.drained_count > 0 {
110                // This is a pretty messed up state, and there isn't really an
111                // obviously right thing to do. We don't want to keep trying
112                // to execute `pred`, so we just backshift all the unprocessed
113                // elements and tell the vec that they still exist. The backshift
114                // is required to prevent a double-drop of the last successfully
115                // drained item prior to a panic in the predicate.
116                let ptr = non_null::as_non_null_ptr(*self.ptr);
117                let src = ptr.add(self.index);
118                let dst = src.sub(self.drained_count);
119                let tail_len = self.original_len - self.index;
120                src.copy_to(dst, tail_len);
121            }
122
123            non_null::set_len(self.ptr, self.original_len - self.drained_count);
124        }
125    }
126}
bump_scope/owned_slice/extract_if.rs

bump_scope/owned_slice/
extract_if.rs
