1 | use crate::fmt; |
2 | use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedFused}; |
3 | use crate::num::NonZero; |
4 | use crate::ops::Try; |
5 | use core::array; |
6 | use core::mem::{ManuallyDrop, MaybeUninit}; |
7 | use core::ops::ControlFlow; |
8 | |
9 | /// An iterator that filters the elements of `iter` with `predicate`. |
10 | /// |
11 | /// This `struct` is created by the [`filter`] method on [`Iterator`]. See its |
12 | /// documentation for more. |
13 | /// |
14 | /// [`filter`]: Iterator::filter |
15 | /// [`Iterator`]: trait.Iterator.html |
16 | #[must_use = "iterators are lazy and do nothing unless consumed" ] |
17 | #[stable (feature = "rust1" , since = "1.0.0" )] |
18 | #[derive (Clone)] |
19 | pub struct Filter<I, P> { |
20 | // Used for `SplitWhitespace` and `SplitAsciiWhitespace` `as_str` methods |
21 | pub(crate) iter: I, |
22 | predicate: P, |
23 | } |
24 | impl<I, P> Filter<I, P> { |
25 | pub(in crate::iter) fn new(iter: I, predicate: P) -> Filter<I, P> { |
26 | Filter { iter, predicate } |
27 | } |
28 | } |
29 | |
30 | #[stable (feature = "core_impl_debug" , since = "1.9.0" )] |
31 | impl<I: fmt::Debug, P> fmt::Debug for Filter<I, P> { |
32 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
33 | f.debug_struct("Filter" ).field(name:"iter" , &self.iter).finish() |
34 | } |
35 | } |
36 | |
37 | fn filter_fold<T, Acc>( |
38 | mut predicate: impl FnMut(&T) -> bool, |
39 | mut fold: impl FnMut(Acc, T) -> Acc, |
40 | ) -> impl FnMut(Acc, T) -> Acc { |
41 | move |acc: Acc, item: T| if predicate(&item) { fold(acc, item) } else { acc } |
42 | } |
43 | |
44 | fn filter_try_fold<'a, T, Acc, R: Try<Output = Acc>>( |
45 | predicate: &'a mut impl FnMut(&T) -> bool, |
46 | mut fold: impl FnMut(Acc, T) -> R + 'a, |
47 | ) -> impl FnMut(Acc, T) -> R + 'a { |
48 | move |acc: Acc, item: T| if predicate(&item) { fold(acc, item) } else { try { acc } } |
49 | } |
50 | |
51 | #[stable (feature = "rust1" , since = "1.0.0" )] |
52 | impl<I: Iterator, P> Iterator for Filter<I, P> |
53 | where |
54 | P: FnMut(&I::Item) -> bool, |
55 | { |
56 | type Item = I::Item; |
57 | |
58 | #[inline ] |
59 | fn next(&mut self) -> Option<I::Item> { |
60 | self.iter.find(&mut self.predicate) |
61 | } |
62 | |
63 | #[inline ] |
64 | fn next_chunk<const N: usize>( |
65 | &mut self, |
66 | ) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>> { |
67 | let mut array: [MaybeUninit<Self::Item>; N] = MaybeUninit::uninit_array(); |
68 | |
69 | struct Guard<'a, T> { |
70 | array: &'a mut [MaybeUninit<T>], |
71 | initialized: usize, |
72 | } |
73 | |
74 | impl<T> Drop for Guard<'_, T> { |
75 | #[inline ] |
76 | fn drop(&mut self) { |
77 | if const { crate::mem::needs_drop::<T>() } { |
78 | // SAFETY: self.initialized is always <= N, which also is the length of the array. |
79 | unsafe { |
80 | core::ptr::drop_in_place(MaybeUninit::slice_assume_init_mut( |
81 | self.array.get_unchecked_mut(..self.initialized), |
82 | )); |
83 | } |
84 | } |
85 | } |
86 | } |
87 | |
88 | let mut guard = Guard { array: &mut array, initialized: 0 }; |
89 | |
90 | let result = self.iter.try_for_each(|element| { |
91 | let idx = guard.initialized; |
92 | guard.initialized = idx + (self.predicate)(&element) as usize; |
93 | |
94 | // SAFETY: Loop conditions ensure the index is in bounds. |
95 | unsafe { guard.array.get_unchecked_mut(idx) }.write(element); |
96 | |
97 | if guard.initialized < N { ControlFlow::Continue(()) } else { ControlFlow::Break(()) } |
98 | }); |
99 | |
100 | let guard = ManuallyDrop::new(guard); |
101 | |
102 | match result { |
103 | ControlFlow::Break(()) => { |
104 | // SAFETY: The loop above is only explicitly broken when the array has been fully initialized |
105 | Ok(unsafe { MaybeUninit::array_assume_init(array) }) |
106 | } |
107 | ControlFlow::Continue(()) => { |
108 | let initialized = guard.initialized; |
109 | // SAFETY: The range is in bounds since the loop breaks when reaching N elements. |
110 | Err(unsafe { array::IntoIter::new_unchecked(array, 0..initialized) }) |
111 | } |
112 | } |
113 | } |
114 | |
115 | #[inline ] |
116 | fn size_hint(&self) -> (usize, Option<usize>) { |
117 | let (_, upper) = self.iter.size_hint(); |
118 | (0, upper) // can't know a lower bound, due to the predicate |
119 | } |
120 | |
121 | // this special case allows the compiler to make `.filter(_).count()` |
122 | // branchless. Barring perfect branch prediction (which is unattainable in |
123 | // the general case), this will be much faster in >90% of cases (containing |
124 | // virtually all real workloads) and only a tiny bit slower in the rest. |
125 | // |
126 | // Having this specialization thus allows us to write `.filter(p).count()` |
127 | // where we would otherwise write `.map(|x| p(x) as usize).sum()`, which is |
128 | // less readable and also less backwards-compatible to Rust before 1.10. |
129 | // |
130 | // Using the branchless version will also simplify the LLVM byte code, thus |
131 | // leaving more budget for LLVM optimizations. |
132 | #[inline ] |
133 | fn count(self) -> usize { |
134 | #[inline ] |
135 | fn to_usize<T>(mut predicate: impl FnMut(&T) -> bool) -> impl FnMut(T) -> usize { |
136 | move |x| predicate(&x) as usize |
137 | } |
138 | |
139 | self.iter.map(to_usize(self.predicate)).sum() |
140 | } |
141 | |
142 | #[inline ] |
143 | fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
144 | where |
145 | Self: Sized, |
146 | Fold: FnMut(Acc, Self::Item) -> R, |
147 | R: Try<Output = Acc>, |
148 | { |
149 | self.iter.try_fold(init, filter_try_fold(&mut self.predicate, fold)) |
150 | } |
151 | |
152 | #[inline ] |
153 | fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc |
154 | where |
155 | Fold: FnMut(Acc, Self::Item) -> Acc, |
156 | { |
157 | self.iter.fold(init, filter_fold(self.predicate, fold)) |
158 | } |
159 | } |
160 | |
161 | #[stable (feature = "rust1" , since = "1.0.0" )] |
162 | impl<I: DoubleEndedIterator, P> DoubleEndedIterator for Filter<I, P> |
163 | where |
164 | P: FnMut(&I::Item) -> bool, |
165 | { |
166 | #[inline ] |
167 | fn next_back(&mut self) -> Option<I::Item> { |
168 | self.iter.rfind(&mut self.predicate) |
169 | } |
170 | |
171 | #[inline ] |
172 | fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
173 | where |
174 | Self: Sized, |
175 | Fold: FnMut(Acc, Self::Item) -> R, |
176 | R: Try<Output = Acc>, |
177 | { |
178 | self.iter.try_rfold(init, f:filter_try_fold(&mut self.predicate, fold)) |
179 | } |
180 | |
181 | #[inline ] |
182 | fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc |
183 | where |
184 | Fold: FnMut(Acc, Self::Item) -> Acc, |
185 | { |
186 | self.iter.rfold(init, f:filter_fold(self.predicate, fold)) |
187 | } |
188 | } |
189 | |
190 | #[stable (feature = "fused" , since = "1.26.0" )] |
191 | impl<I: FusedIterator, P> FusedIterator for Filter<I, P> where P: FnMut(&I::Item) -> bool {} |
192 | |
193 | #[unstable (issue = "none" , feature = "trusted_fused" )] |
194 | unsafe impl<I: TrustedFused, F> TrustedFused for Filter<I, F> {} |
195 | |
196 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
197 | unsafe impl<P, I> SourceIter for Filter<I, P> |
198 | where |
199 | I: SourceIter, |
200 | { |
201 | type Source = I::Source; |
202 | |
203 | #[inline ] |
204 | unsafe fn as_inner(&mut self) -> &mut I::Source { |
205 | // SAFETY: unsafe function forwarding to unsafe function with the same requirements |
206 | unsafe { SourceIter::as_inner(&mut self.iter) } |
207 | } |
208 | } |
209 | |
210 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
211 | unsafe impl<I: InPlaceIterable, P> InPlaceIterable for Filter<I, P> { |
212 | const EXPAND_BY: Option<NonZero<usize>> = I::EXPAND_BY; |
213 | const MERGE_BY: Option<NonZero<usize>> = I::MERGE_BY; |
214 | } |
215 | |