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