1 | use crate::cmp; |
2 | use crate::iter::{ |
3 | adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedFused, TrustedLen, |
4 | TrustedRandomAccess, |
5 | }; |
6 | use crate::num::NonZero; |
7 | use crate::ops::{ControlFlow, Try}; |
8 | |
9 | /// An iterator that only iterates over the first `n` iterations of `iter`. |
10 | /// |
11 | /// This `struct` is created by the [`take`] method on [`Iterator`]. See its |
12 | /// documentation for more. |
13 | /// |
14 | /// [`take`]: Iterator::take |
15 | /// [`Iterator`]: trait.Iterator.html |
16 | #[derive (Clone, Debug)] |
17 | #[must_use = "iterators are lazy and do nothing unless consumed" ] |
18 | #[stable (feature = "rust1" , since = "1.0.0" )] |
19 | pub struct Take<I> { |
20 | iter: I, |
21 | n: usize, |
22 | } |
23 | |
24 | impl<I> Take<I> { |
25 | pub(in crate::iter) fn new(iter: I, n: usize) -> Take<I> { |
26 | Take { iter, n } |
27 | } |
28 | } |
29 | |
30 | #[stable (feature = "rust1" , since = "1.0.0" )] |
31 | impl<I> Iterator for Take<I> |
32 | where |
33 | I: Iterator, |
34 | { |
35 | type Item = <I as Iterator>::Item; |
36 | |
37 | #[inline ] |
38 | fn next(&mut self) -> Option<<I as Iterator>::Item> { |
39 | if self.n != 0 { |
40 | self.n -= 1; |
41 | self.iter.next() |
42 | } else { |
43 | None |
44 | } |
45 | } |
46 | |
47 | #[inline ] |
48 | fn nth(&mut self, n: usize) -> Option<I::Item> { |
49 | if self.n > n { |
50 | self.n -= n + 1; |
51 | self.iter.nth(n) |
52 | } else { |
53 | if self.n > 0 { |
54 | self.iter.nth(self.n - 1); |
55 | self.n = 0; |
56 | } |
57 | None |
58 | } |
59 | } |
60 | |
61 | #[inline ] |
62 | fn size_hint(&self) -> (usize, Option<usize>) { |
63 | if self.n == 0 { |
64 | return (0, Some(0)); |
65 | } |
66 | |
67 | let (lower, upper) = self.iter.size_hint(); |
68 | |
69 | let lower = cmp::min(lower, self.n); |
70 | |
71 | let upper = match upper { |
72 | Some(x) if x < self.n => Some(x), |
73 | _ => Some(self.n), |
74 | }; |
75 | |
76 | (lower, upper) |
77 | } |
78 | |
79 | #[inline ] |
80 | fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
81 | where |
82 | Fold: FnMut(Acc, Self::Item) -> R, |
83 | R: Try<Output = Acc>, |
84 | { |
85 | fn check<'a, T, Acc, R: Try<Output = Acc>>( |
86 | n: &'a mut usize, |
87 | mut fold: impl FnMut(Acc, T) -> R + 'a, |
88 | ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> + 'a { |
89 | move |acc, x| { |
90 | *n -= 1; |
91 | let r = fold(acc, x); |
92 | if *n == 0 { ControlFlow::Break(r) } else { ControlFlow::from_try(r) } |
93 | } |
94 | } |
95 | |
96 | if self.n == 0 { |
97 | try { init } |
98 | } else { |
99 | let n = &mut self.n; |
100 | self.iter.try_fold(init, check(n, fold)).into_try() |
101 | } |
102 | } |
103 | |
104 | #[inline ] |
105 | fn fold<B, F>(self, init: B, f: F) -> B |
106 | where |
107 | Self: Sized, |
108 | F: FnMut(B, Self::Item) -> B, |
109 | { |
110 | Self::spec_fold(self, init, f) |
111 | } |
112 | |
113 | #[inline ] |
114 | fn for_each<F: FnMut(Self::Item)>(self, f: F) { |
115 | Self::spec_for_each(self, f) |
116 | } |
117 | |
118 | #[inline ] |
119 | #[rustc_inherit_overflow_checks ] |
120 | fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
121 | let min = self.n.min(n); |
122 | let rem = match self.iter.advance_by(min) { |
123 | Ok(()) => 0, |
124 | Err(rem) => rem.get(), |
125 | }; |
126 | let advanced = min - rem; |
127 | self.n -= advanced; |
128 | NonZero::new(n - advanced).map_or(Ok(()), Err) |
129 | } |
130 | } |
131 | |
132 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
133 | unsafe impl<I> SourceIter for Take<I> |
134 | where |
135 | I: SourceIter, |
136 | { |
137 | type Source = I::Source; |
138 | |
139 | #[inline ] |
140 | unsafe fn as_inner(&mut self) -> &mut I::Source { |
141 | // SAFETY: unsafe function forwarding to unsafe function with the same requirements |
142 | unsafe { SourceIter::as_inner(&mut self.iter) } |
143 | } |
144 | } |
145 | |
146 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
147 | unsafe impl<I: InPlaceIterable> InPlaceIterable for Take<I> { |
148 | const EXPAND_BY: Option<NonZero<usize>> = I::EXPAND_BY; |
149 | const MERGE_BY: Option<NonZero<usize>> = I::MERGE_BY; |
150 | } |
151 | |
152 | #[stable (feature = "double_ended_take_iterator" , since = "1.38.0" )] |
153 | impl<I> DoubleEndedIterator for Take<I> |
154 | where |
155 | I: DoubleEndedIterator + ExactSizeIterator, |
156 | { |
157 | #[inline ] |
158 | fn next_back(&mut self) -> Option<Self::Item> { |
159 | if self.n == 0 { |
160 | None |
161 | } else { |
162 | let n = self.n; |
163 | self.n -= 1; |
164 | self.iter.nth_back(self.iter.len().saturating_sub(n)) |
165 | } |
166 | } |
167 | |
168 | #[inline ] |
169 | fn nth_back(&mut self, n: usize) -> Option<Self::Item> { |
170 | let len = self.iter.len(); |
171 | if self.n > n { |
172 | let m = len.saturating_sub(self.n) + n; |
173 | self.n -= n + 1; |
174 | self.iter.nth_back(m) |
175 | } else { |
176 | if len > 0 { |
177 | self.iter.nth_back(len - 1); |
178 | } |
179 | None |
180 | } |
181 | } |
182 | |
183 | #[inline ] |
184 | fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
185 | where |
186 | Self: Sized, |
187 | Fold: FnMut(Acc, Self::Item) -> R, |
188 | R: Try<Output = Acc>, |
189 | { |
190 | if self.n == 0 { |
191 | try { init } |
192 | } else { |
193 | let len = self.iter.len(); |
194 | if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() { |
195 | try { init } |
196 | } else { |
197 | self.iter.try_rfold(init, fold) |
198 | } |
199 | } |
200 | } |
201 | |
202 | #[inline ] |
203 | fn rfold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc |
204 | where |
205 | Self: Sized, |
206 | Fold: FnMut(Acc, Self::Item) -> Acc, |
207 | { |
208 | if self.n == 0 { |
209 | init |
210 | } else { |
211 | let len = self.iter.len(); |
212 | if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() { |
213 | init |
214 | } else { |
215 | self.iter.rfold(init, fold) |
216 | } |
217 | } |
218 | } |
219 | |
220 | #[inline ] |
221 | #[rustc_inherit_overflow_checks ] |
222 | fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
223 | // The amount by which the inner iterator needs to be shortened for it to be |
224 | // at most as long as the take() amount. |
225 | let trim_inner = self.iter.len().saturating_sub(self.n); |
226 | // The amount we need to advance inner to fulfill the caller's request. |
227 | // take(), advance_by() and len() all can be at most usize, so we don't have to worry |
228 | // about having to advance more than usize::MAX here. |
229 | let advance_by = trim_inner.saturating_add(n); |
230 | |
231 | let remainder = match self.iter.advance_back_by(advance_by) { |
232 | Ok(()) => 0, |
233 | Err(rem) => rem.get(), |
234 | }; |
235 | let advanced_by_inner = advance_by - remainder; |
236 | let advanced_by = advanced_by_inner - trim_inner; |
237 | self.n -= advanced_by; |
238 | NonZero::new(n - advanced_by).map_or(Ok(()), Err) |
239 | } |
240 | } |
241 | |
242 | #[stable (feature = "rust1" , since = "1.0.0" )] |
243 | impl<I> ExactSizeIterator for Take<I> where I: ExactSizeIterator {} |
244 | |
245 | #[stable (feature = "fused" , since = "1.26.0" )] |
246 | impl<I> FusedIterator for Take<I> where I: FusedIterator {} |
247 | |
248 | #[unstable (issue = "none" , feature = "trusted_fused" )] |
249 | unsafe impl<I: TrustedFused> TrustedFused for Take<I> {} |
250 | |
251 | #[unstable (feature = "trusted_len" , issue = "37572" )] |
252 | unsafe impl<I: TrustedLen> TrustedLen for Take<I> {} |
253 | |
254 | trait SpecTake: Iterator { |
255 | fn spec_fold<B, F>(self, init: B, f: F) -> B |
256 | where |
257 | Self: Sized, |
258 | F: FnMut(B, Self::Item) -> B; |
259 | |
260 | fn spec_for_each<F: FnMut(Self::Item)>(self, f: F); |
261 | } |
262 | |
263 | impl<I: Iterator> SpecTake for Take<I> { |
264 | #[inline ] |
265 | default fn spec_fold<B, F>(mut self, init: B, f: F) -> B |
266 | where |
267 | Self: Sized, |
268 | F: FnMut(B, Self::Item) -> B, |
269 | { |
270 | use crate::ops::NeverShortCircuit; |
271 | self.try_fold(init, NeverShortCircuit::wrap_mut_2(f)).0 |
272 | } |
273 | |
274 | #[inline ] |
275 | default fn spec_for_each<F: FnMut(Self::Item)>(mut self, f: F) { |
276 | // The default implementation would use a unit accumulator, so we can |
277 | // avoid a stateful closure by folding over the remaining number |
278 | // of items we wish to return instead. |
279 | fn check<'a, Item>( |
280 | mut action: impl FnMut(Item) + 'a, |
281 | ) -> impl FnMut(usize, Item) -> Option<usize> + 'a { |
282 | move |more, x| { |
283 | action(x); |
284 | more.checked_sub(1) |
285 | } |
286 | } |
287 | |
288 | let remaining = self.n; |
289 | if remaining > 0 { |
290 | self.iter.try_fold(remaining - 1, check(f)); |
291 | } |
292 | } |
293 | } |
294 | |
295 | impl<I: Iterator + TrustedRandomAccess> SpecTake for Take<I> { |
296 | #[inline ] |
297 | fn spec_fold<B, F>(mut self, init: B, mut f: F) -> B |
298 | where |
299 | Self: Sized, |
300 | F: FnMut(B, Self::Item) -> B, |
301 | { |
302 | let mut acc = init; |
303 | let end = self.n.min(self.iter.size()); |
304 | for i in 0..end { |
305 | // SAFETY: i < end <= self.iter.size() and we discard the iterator at the end |
306 | let val = unsafe { self.iter.__iterator_get_unchecked(i) }; |
307 | acc = f(acc, val); |
308 | } |
309 | acc |
310 | } |
311 | |
312 | #[inline ] |
313 | fn spec_for_each<F: FnMut(Self::Item)>(mut self, mut f: F) { |
314 | let end = self.n.min(self.iter.size()); |
315 | for i in 0..end { |
316 | // SAFETY: i < end <= self.iter.size() and we discard the iterator at the end |
317 | let val = unsafe { self.iter.__iterator_get_unchecked(i) }; |
318 | f(val); |
319 | } |
320 | } |
321 | } |
322 | |