1use crate::cmp;
2use crate::iter::{
3 adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedFused, TrustedLen,
4 TrustedRandomAccess,
5};
6use crate::num::NonZero;
7use 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")]
19pub struct Take<I> {
20 iter: I,
21 n: usize,
22}
23
24impl<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")]
31impl<I> Iterator for Take<I>
32where
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")]
133unsafe impl<I> SourceIter for Take<I>
134where
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")]
147unsafe 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")]
153impl<I> DoubleEndedIterator for Take<I>
154where
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")]
243impl<I> ExactSizeIterator for Take<I> where I: ExactSizeIterator {}
244
245#[stable(feature = "fused", since = "1.26.0")]
246impl<I> FusedIterator for Take<I> where I: FusedIterator {}
247
248#[unstable(issue = "none", feature = "trusted_fused")]
249unsafe impl<I: TrustedFused> TrustedFused for Take<I> {}
250
251#[unstable(feature = "trusted_len", issue = "37572")]
252unsafe impl<I: TrustedLen> TrustedLen for Take<I> {}
253
254trait 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
263impl<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
295impl<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