1 | use crate::intrinsics::unlikely; |
2 | use crate::iter::adapters::zip::try_get_unchecked; |
3 | use crate::iter::TrustedFused; |
4 | use crate::iter::{ |
5 | adapters::SourceIter, FusedIterator, InPlaceIterable, TrustedLen, TrustedRandomAccess, |
6 | TrustedRandomAccessNoCoerce, |
7 | }; |
8 | use crate::num::NonZero; |
9 | use crate::ops::{ControlFlow, Try}; |
10 | |
11 | /// An iterator that skips over `n` elements of `iter`. |
12 | /// |
13 | /// This `struct` is created by the [`skip`] method on [`Iterator`]. See its |
14 | /// documentation for more. |
15 | /// |
16 | /// [`skip`]: Iterator::skip |
17 | /// [`Iterator`]: trait.Iterator.html |
18 | #[derive (Clone, Debug)] |
19 | #[must_use = "iterators are lazy and do nothing unless consumed" ] |
20 | #[stable (feature = "rust1" , since = "1.0.0" )] |
21 | pub struct Skip<I> { |
22 | iter: I, |
23 | n: usize, |
24 | } |
25 | |
26 | impl<I> Skip<I> { |
27 | pub(in crate::iter) fn new(iter: I, n: usize) -> Skip<I> { |
28 | Skip { iter, n } |
29 | } |
30 | } |
31 | |
32 | #[stable (feature = "rust1" , since = "1.0.0" )] |
33 | impl<I> Iterator for Skip<I> |
34 | where |
35 | I: Iterator, |
36 | { |
37 | type Item = <I as Iterator>::Item; |
38 | |
39 | #[inline ] |
40 | fn next(&mut self) -> Option<I::Item> { |
41 | if unlikely(self.n > 0) { |
42 | self.iter.nth(crate::mem::take(&mut self.n)) |
43 | } else { |
44 | self.iter.next() |
45 | } |
46 | } |
47 | |
48 | #[inline ] |
49 | fn nth(&mut self, n: usize) -> Option<I::Item> { |
50 | if self.n > 0 { |
51 | let skip: usize = crate::mem::take(&mut self.n); |
52 | // Checked add to handle overflow case. |
53 | let n = match skip.checked_add(n) { |
54 | Some(nth) => nth, |
55 | None => { |
56 | // In case of overflow, load skip value, before loading `n`. |
57 | // Because the amount of elements to iterate is beyond `usize::MAX`, this |
58 | // is split into two `nth` calls where the `skip` `nth` call is discarded. |
59 | self.iter.nth(skip - 1)?; |
60 | n |
61 | } |
62 | }; |
63 | // Load nth element including skip. |
64 | self.iter.nth(n) |
65 | } else { |
66 | self.iter.nth(n) |
67 | } |
68 | } |
69 | |
70 | #[inline ] |
71 | fn count(mut self) -> usize { |
72 | if self.n > 0 { |
73 | // nth(n) skips n+1 |
74 | if self.iter.nth(self.n - 1).is_none() { |
75 | return 0; |
76 | } |
77 | } |
78 | self.iter.count() |
79 | } |
80 | |
81 | #[inline ] |
82 | fn last(mut self) -> Option<I::Item> { |
83 | if self.n > 0 { |
84 | // nth(n) skips n+1 |
85 | self.iter.nth(self.n - 1)?; |
86 | } |
87 | self.iter.last() |
88 | } |
89 | |
90 | #[inline ] |
91 | fn size_hint(&self) -> (usize, Option<usize>) { |
92 | let (lower, upper) = self.iter.size_hint(); |
93 | |
94 | let lower = lower.saturating_sub(self.n); |
95 | let upper = match upper { |
96 | Some(x) => Some(x.saturating_sub(self.n)), |
97 | None => None, |
98 | }; |
99 | |
100 | (lower, upper) |
101 | } |
102 | |
103 | #[inline ] |
104 | fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
105 | where |
106 | Self: Sized, |
107 | Fold: FnMut(Acc, Self::Item) -> R, |
108 | R: Try<Output = Acc>, |
109 | { |
110 | let n = self.n; |
111 | self.n = 0; |
112 | if n > 0 { |
113 | // nth(n) skips n+1 |
114 | if self.iter.nth(n - 1).is_none() { |
115 | return try { init }; |
116 | } |
117 | } |
118 | self.iter.try_fold(init, fold) |
119 | } |
120 | |
121 | #[inline ] |
122 | fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc |
123 | where |
124 | Fold: FnMut(Acc, Self::Item) -> Acc, |
125 | { |
126 | if self.n > 0 { |
127 | // nth(n) skips n+1 |
128 | if self.iter.nth(self.n - 1).is_none() { |
129 | return init; |
130 | } |
131 | } |
132 | self.iter.fold(init, fold) |
133 | } |
134 | |
135 | #[inline ] |
136 | #[rustc_inherit_overflow_checks ] |
137 | fn advance_by(&mut self, mut n: usize) -> Result<(), NonZero<usize>> { |
138 | let skip_inner = self.n; |
139 | let skip_and_advance = skip_inner.saturating_add(n); |
140 | |
141 | let remainder = match self.iter.advance_by(skip_and_advance) { |
142 | Ok(()) => 0, |
143 | Err(n) => n.get(), |
144 | }; |
145 | let advanced_inner = skip_and_advance - remainder; |
146 | n -= advanced_inner.saturating_sub(skip_inner); |
147 | self.n = self.n.saturating_sub(advanced_inner); |
148 | |
149 | // skip_and_advance may have saturated |
150 | if unlikely(remainder == 0 && n > 0) { |
151 | n = match self.iter.advance_by(n) { |
152 | Ok(()) => 0, |
153 | Err(n) => n.get(), |
154 | } |
155 | } |
156 | |
157 | NonZero::new(n).map_or(Ok(()), Err) |
158 | } |
159 | |
160 | #[doc (hidden)] |
161 | unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> Self::Item |
162 | where |
163 | Self: TrustedRandomAccessNoCoerce, |
164 | { |
165 | // SAFETY: the caller must uphold the contract for |
166 | // `Iterator::__iterator_get_unchecked`. |
167 | // |
168 | // Dropping the skipped prefix when index 0 is passed is safe |
169 | // since |
170 | // * the caller passing index 0 means that the inner iterator has more items than `self.n` |
171 | // * TRA contract requires that get_unchecked will only be called once |
172 | // (unless elements are copyable) |
173 | // * it does not conflict with in-place iteration since index 0 must be accessed |
174 | // before something is written into the storage used by the prefix |
175 | unsafe { |
176 | if Self::MAY_HAVE_SIDE_EFFECT && idx == 0 { |
177 | for skipped_idx in 0..self.n { |
178 | drop(try_get_unchecked(&mut self.iter, skipped_idx)); |
179 | } |
180 | } |
181 | |
182 | try_get_unchecked(&mut self.iter, idx + self.n) |
183 | } |
184 | } |
185 | } |
186 | |
187 | #[stable (feature = "rust1" , since = "1.0.0" )] |
188 | impl<I> ExactSizeIterator for Skip<I> where I: ExactSizeIterator {} |
189 | |
190 | #[stable (feature = "double_ended_skip_iterator" , since = "1.9.0" )] |
191 | impl<I> DoubleEndedIterator for Skip<I> |
192 | where |
193 | I: DoubleEndedIterator + ExactSizeIterator, |
194 | { |
195 | fn next_back(&mut self) -> Option<Self::Item> { |
196 | if self.len() > 0 { self.iter.next_back() } else { None } |
197 | } |
198 | |
199 | #[inline ] |
200 | fn nth_back(&mut self, n: usize) -> Option<I::Item> { |
201 | let len = self.len(); |
202 | if n < len { |
203 | self.iter.nth_back(n) |
204 | } else { |
205 | if len > 0 { |
206 | // consume the original iterator |
207 | self.iter.nth_back(len - 1); |
208 | } |
209 | None |
210 | } |
211 | } |
212 | |
213 | fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R |
214 | where |
215 | Self: Sized, |
216 | Fold: FnMut(Acc, Self::Item) -> R, |
217 | R: Try<Output = Acc>, |
218 | { |
219 | fn check<T, Acc, R: Try<Output = Acc>>( |
220 | mut n: usize, |
221 | mut fold: impl FnMut(Acc, T) -> R, |
222 | ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> { |
223 | move |acc, x| { |
224 | n -= 1; |
225 | let r = fold(acc, x); |
226 | if n == 0 { ControlFlow::Break(r) } else { ControlFlow::from_try(r) } |
227 | } |
228 | } |
229 | |
230 | let n = self.len(); |
231 | if n == 0 { try { init } } else { self.iter.try_rfold(init, check(n, fold)).into_try() } |
232 | } |
233 | |
234 | impl_fold_via_try_fold! { rfold -> try_rfold } |
235 | |
236 | #[inline ] |
237 | fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> { |
238 | let min = crate::cmp::min(self.len(), n); |
239 | let rem = self.iter.advance_back_by(min); |
240 | assert!(rem.is_ok(), "ExactSizeIterator contract violation" ); |
241 | NonZero::new(n - min).map_or(Ok(()), Err) |
242 | } |
243 | } |
244 | |
245 | #[stable (feature = "fused" , since = "1.26.0" )] |
246 | impl<I> FusedIterator for Skip<I> where I: FusedIterator {} |
247 | |
248 | #[unstable (issue = "none" , feature = "trusted_fused" )] |
249 | unsafe impl<I: TrustedFused> TrustedFused for Skip<I> {} |
250 | |
251 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
252 | unsafe impl<I> SourceIter for Skip<I> |
253 | where |
254 | I: SourceIter, |
255 | { |
256 | type Source = I::Source; |
257 | |
258 | #[inline ] |
259 | unsafe fn as_inner(&mut self) -> &mut I::Source { |
260 | // SAFETY: unsafe function forwarding to unsafe function with the same requirements |
261 | unsafe { SourceIter::as_inner(&mut self.iter) } |
262 | } |
263 | } |
264 | |
265 | #[unstable (issue = "none" , feature = "inplace_iteration" )] |
266 | unsafe impl<I: InPlaceIterable> InPlaceIterable for Skip<I> { |
267 | const EXPAND_BY: Option<NonZero<usize>> = I::EXPAND_BY; |
268 | const MERGE_BY: Option<NonZero<usize>> = I::MERGE_BY; |
269 | } |
270 | |
271 | #[doc (hidden)] |
272 | #[unstable (feature = "trusted_random_access" , issue = "none" )] |
273 | unsafe impl<I> TrustedRandomAccess for Skip<I> where I: TrustedRandomAccess {} |
274 | |
275 | #[doc (hidden)] |
276 | #[unstable (feature = "trusted_random_access" , issue = "none" )] |
277 | unsafe impl<I> TrustedRandomAccessNoCoerce for Skip<I> |
278 | where |
279 | I: TrustedRandomAccessNoCoerce, |
280 | { |
281 | const MAY_HAVE_SIDE_EFFECT: bool = I::MAY_HAVE_SIDE_EFFECT; |
282 | } |
283 | |
284 | // SAFETY: This adapter is shortening. TrustedLen requires the upper bound to be calculated correctly. |
285 | // These requirements can only be satisfied when the upper bound of the inner iterator's upper |
286 | // bound is never `None`. I: TrustedRandomAccess happens to provide this guarantee while |
287 | // I: TrustedLen would not. |
288 | #[unstable (feature = "trusted_len" , issue = "37572" )] |
289 | unsafe impl<I> TrustedLen for Skip<I> where I: Iterator + TrustedRandomAccess {} |
290 | |