1use core::iter::FusedIterator;
4use core::ptr;
5use core::slice;
6
7use crate::TinyVec;
8
9pub struct ExtractIf<'a, T, const N: usize, F>
10where
11    T: 'a,
12    F: FnMut(&mut T) -> bool,
13{
14    pub (super) vec: &'a mut TinyVec<T, N>,
15    pub (super) next: usize,
16    pub (super) last: usize,
17    pub (super) original_len: usize,
18    pub (super) deleted: usize,
19    pub (super) pred: F,
20}
21
22impl<T, const N: usize, F> Iterator for ExtractIf<'_, T, N, F>
23where
24    F: FnMut(&mut T) -> bool,
25{
26    type Item = T;
27
28    fn next(&mut self) -> Option<Self::Item> {
29        unsafe {
30            while self.next < self.last {
31                let i = self.next;
32                let v = slice::from_raw_parts_mut(self.vec.as_mut_ptr(), self.original_len);
33                let drained = (self.pred)(&mut v[i]);
34
35                self.next += 1;
40
41                if drained {
42                    self.deleted += 1;
43                    return Some(ptr::read(&v[i]));
44                } else if self.deleted > 0 {
45                    let del = self.deleted;
46                    let src: *const T = &v[i];
47                    let dst: *mut T = &mut v[i - del];
48                    ptr::copy_nonoverlapping(src, dst, 1);
49                }
50            }
51            None
52        }
53    }
54}
55
56impl<T, const N: usize, F> FusedIterator for ExtractIf<'_, T, N, F>
57where
58    F: FnMut(&mut T) -> bool {}
59
60impl<T, const N: usize, F> Drop for ExtractIf<'_, T, N, F>
61where
62    F: FnMut(&mut T) -> bool,
63{
64    fn drop(&mut self) {
65        unsafe {
66            if self.next < self.original_len && self.deleted > 0 {
67                let ptr = self.vec.as_mut_ptr();
68                let src = ptr.add(self.next);
69                let dst = src.sub(self.deleted);
70                let tail_len = self.original_len - self.next;
71                src.copy_to(dst, tail_len);
72            }
73            self.vec.set_len(self.original_len - self.deleted);
74        }
75    }
76}