core/iter/adapters/
filter.rs

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, TrustedLen};
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)]
21#[ferrocene::prevalidated]
22pub struct Filter<I, P> {
23    // Used for `SplitWhitespace` and `SplitAsciiWhitespace` `as_str` methods
24    pub(crate) iter: I,
25    predicate: P,
26}
27impl<I, P> Filter<I, P> {
28    #[ferrocene::prevalidated]
29    pub(in crate::iter) const fn new(iter: I, predicate: P) -> Filter<I, P> {
30        Filter { iter, predicate }
31    }
32}
33
34impl<I, P> Filter<I, P>
35where
36    I: Iterator,
37    P: FnMut(&I::Item) -> bool,
38{
39    #[inline]
40    fn next_chunk_dropless<const N: usize>(
41        &mut self,
42    ) -> Result<[I::Item; N], array::IntoIter<I::Item, N>> {
43        let mut array: [MaybeUninit<I::Item>; N] = [const { MaybeUninit::uninit() }; N];
44        let mut initialized = 0;
45
46        let result = self.iter.try_for_each(|element| {
47            let idx = initialized;
48            // branchless index update combined with unconditionally copying the value even when
49            // it is filtered reduces branching and dependencies in the loop.
50            initialized = idx + (self.predicate)(&element) as usize;
51            // SAFETY: Loop conditions ensure the index is in bounds.
52            unsafe { array.get_unchecked_mut(idx) }.write(element);
53
54            if initialized < N { ControlFlow::Continue(()) } else { ControlFlow::Break(()) }
55        });
56
57        match result {
58            ControlFlow::Break(()) => {
59                // SAFETY: The loop above is only explicitly broken when the array has been fully initialized
60                Ok(unsafe { MaybeUninit::array_assume_init(array) })
61            }
62            ControlFlow::Continue(()) => {
63                // SAFETY: The range is in bounds since the loop breaks when reaching N elements.
64                Err(unsafe { array::IntoIter::new_unchecked(array, 0..initialized) })
65            }
66        }
67    }
68}
69
70#[stable(feature = "core_impl_debug", since = "1.9.0")]
71impl<I: fmt::Debug, P> fmt::Debug for Filter<I, P> {
72    #[ferrocene::prevalidated]
73    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
74        f.debug_struct("Filter").field("iter", &self.iter).finish()
75    }
76}
77
78#[ferrocene::prevalidated]
79fn filter_fold<T, Acc>(
80    mut predicate: impl FnMut(&T) -> bool,
81    mut fold: impl FnMut(Acc, T) -> Acc,
82) -> impl FnMut(Acc, T) -> Acc {
83    move |acc, item| if predicate(&item) { fold(acc, item) } else { acc }
84}
85
86#[ferrocene::prevalidated]
87fn filter_try_fold<'a, T, Acc, R: Try<Output = Acc>>(
88    predicate: &'a mut impl FnMut(&T) -> bool,
89    mut fold: impl FnMut(Acc, T) -> R + 'a,
90) -> impl FnMut(Acc, T) -> R + 'a {
91    move |acc, item| if predicate(&item) { fold(acc, item) } else { try { acc } }
92}
93
94#[stable(feature = "rust1", since = "1.0.0")]
95impl<I: Iterator, P> Iterator for Filter<I, P>
96where
97    P: FnMut(&I::Item) -> bool,
98{
99    type Item = I::Item;
100
101    #[inline]
102    #[ferrocene::prevalidated]
103    fn next(&mut self) -> Option<I::Item> {
104        self.iter.find(&mut self.predicate)
105    }
106
107    #[inline]
108    fn next_chunk<const N: usize>(
109        &mut self,
110    ) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>> {
111        // avoid codegen for the dead branch
112        let fun = const {
113            if crate::mem::needs_drop::<I::Item>() {
114                array::iter_next_chunk::<I::Item, N>
115            } else {
116                Self::next_chunk_dropless::<N>
117            }
118        };
119
120        fun(self)
121    }
122
123    #[inline]
124    #[ferrocene::prevalidated]
125    fn size_hint(&self) -> (usize, Option<usize>) {
126        let (_, upper) = self.iter.size_hint();
127        (0, upper) // can't know a lower bound, due to the predicate
128    }
129
130    // this special case allows the compiler to make `.filter(_).count()`
131    // branchless. Barring perfect branch prediction (which is unattainable in
132    // the general case), this will be much faster in >90% of cases (containing
133    // virtually all real workloads) and only a tiny bit slower in the rest.
134    //
135    // Having this specialization thus allows us to write `.filter(p).count()`
136    // where we would otherwise write `.map(|x| p(x) as usize).sum()`, which is
137    // less readable and also less backwards-compatible to Rust before 1.10.
138    //
139    // Using the branchless version will also simplify the LLVM byte code, thus
140    // leaving more budget for LLVM optimizations.
141    #[inline]
142    #[ferrocene::prevalidated]
143    fn count(self) -> usize {
144        #[inline]
145        #[ferrocene::prevalidated]
146        fn to_usize<T>(mut predicate: impl FnMut(&T) -> bool) -> impl FnMut(T) -> usize {
147            move |x| predicate(&x) as usize
148        }
149
150        let before = self.iter.size_hint().1.unwrap_or(usize::MAX);
151        let total = self.iter.map(to_usize(self.predicate)).sum();
152        // SAFETY: `total` and `before` came from the same iterator of type `I`
153        unsafe {
154            <I as SpecAssumeCount>::assume_count_le_upper_bound(total, before);
155        }
156        total
157    }
158
159    #[inline]
160    #[ferrocene::prevalidated]
161    fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
162    where
163        Self: Sized,
164        Fold: FnMut(Acc, Self::Item) -> R,
165        R: Try<Output = Acc>,
166    {
167        self.iter.try_fold(init, filter_try_fold(&mut self.predicate, fold))
168    }
169
170    #[inline]
171    #[ferrocene::prevalidated]
172    fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
173    where
174        Fold: FnMut(Acc, Self::Item) -> Acc,
175    {
176        self.iter.fold(init, filter_fold(self.predicate, fold))
177    }
178}
179
180#[stable(feature = "rust1", since = "1.0.0")]
181impl<I: DoubleEndedIterator, P> DoubleEndedIterator for Filter<I, P>
182where
183    P: FnMut(&I::Item) -> bool,
184{
185    #[inline]
186    fn next_back(&mut self) -> Option<I::Item> {
187        self.iter.rfind(&mut self.predicate)
188    }
189
190    #[inline]
191    fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
192    where
193        Self: Sized,
194        Fold: FnMut(Acc, Self::Item) -> R,
195        R: Try<Output = Acc>,
196    {
197        self.iter.try_rfold(init, filter_try_fold(&mut self.predicate, fold))
198    }
199
200    #[inline]
201    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
202    where
203        Fold: FnMut(Acc, Self::Item) -> Acc,
204    {
205        self.iter.rfold(init, filter_fold(self.predicate, fold))
206    }
207}
208
209#[stable(feature = "fused", since = "1.26.0")]
210impl<I: FusedIterator, P> FusedIterator for Filter<I, P> where P: FnMut(&I::Item) -> bool {}
211
212#[unstable(issue = "none", feature = "trusted_fused")]
213unsafe impl<I: TrustedFused, F> TrustedFused for Filter<I, F> {}
214
215#[unstable(issue = "none", feature = "inplace_iteration")]
216unsafe impl<P, I> SourceIter for Filter<I, P>
217where
218    I: SourceIter,
219{
220    type Source = I::Source;
221
222    #[inline]
223    unsafe fn as_inner(&mut self) -> &mut I::Source {
224        // SAFETY: unsafe function forwarding to unsafe function with the same requirements
225        unsafe { SourceIter::as_inner(&mut self.iter) }
226    }
227}
228
229#[unstable(issue = "none", feature = "inplace_iteration")]
230unsafe impl<I: InPlaceIterable, P> InPlaceIterable for Filter<I, P> {
231    const EXPAND_BY: Option<NonZero<usize>> = I::EXPAND_BY;
232    const MERGE_BY: Option<NonZero<usize>> = I::MERGE_BY;
233}
234
235trait SpecAssumeCount {
236    /// # Safety
237    ///
238    /// `count` must be an number of items actually read from the iterator.
239    ///
240    /// `upper` must either:
241    /// - have come from `size_hint().1` on the iterator, or
242    /// - be `usize::MAX` which will vacuously do nothing.
243    unsafe fn assume_count_le_upper_bound(count: usize, upper: usize);
244}
245
246impl<I: Iterator> SpecAssumeCount for I {
247    #[inline]
248    #[rustc_inherit_overflow_checks]
249    #[ferrocene::prevalidated]
250    default unsafe fn assume_count_le_upper_bound(count: usize, upper: usize) {
251        // In the default we can't trust the `upper` for soundness
252        // because it came from an untrusted `size_hint`.
253
254        // In debug mode we might as well check that the size_hint wasn't too small
255        let _ = upper - count;
256    }
257}
258
259impl<I: TrustedLen> SpecAssumeCount for I {
260    #[inline]
261    #[ferrocene::prevalidated]
262    unsafe fn assume_count_le_upper_bound(count: usize, upper: usize) {
263        // SAFETY: The `upper` is trusted because it came from a `TrustedLen` iterator.
264        unsafe { crate::hint::assert_unchecked(count <= upper) }
265    }
266}
core/iter/adapters/filter.rs

core/iter/adapters/
filter.rs
