1use crate::fmt;
2use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedFused};
3use crate::num::NonZero;
4use crate::ops::Try;
5use core::array;
6use core::mem::{ManuallyDrop, MaybeUninit};
7use 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)]
19pub struct Filter<I, P> {
20 // Used for `SplitWhitespace` and `SplitAsciiWhitespace` `as_str` methods
21 pub(crate) iter: I,
22 predicate: P,
23}
24impl<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")]
31impl<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
37fn 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
44fn 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")]
52impl<I: Iterator, P> Iterator for Filter<I, P>
53where
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")]
162impl<I: DoubleEndedIterator, P> DoubleEndedIterator for Filter<I, P>
163where
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")]
191impl<I: FusedIterator, P> FusedIterator for Filter<I, P> where P: FnMut(&I::Item) -> bool {}
192
193#[unstable(issue = "none", feature = "trusted_fused")]
194unsafe impl<I: TrustedFused, F> TrustedFused for Filter<I, F> {}
195
196#[unstable(issue = "none", feature = "inplace_iteration")]
197unsafe impl<P, I> SourceIter for Filter<I, P>
198where
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")]
211unsafe 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