Skip to main content

core/iter/adapters/
skip.rs

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