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