library/core/src/iter/adapters/skip.rs

   1 use crate::intrinsics::unlikely;
   2 use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable};
   3 use crate::ops::{ControlFlow, Try};
   4
   5 /// An iterator that skips over `n` elements of `iter`.
   6 ///
   7 /// This `struct` is created by the [`skip`] method on [`Iterator`]. See its
   8 /// documentation for more.
   9 ///
  10 /// [`skip`]: Iterator::skip
  11 /// [`Iterator`]: trait.Iterator.html
  12 #[derive(Clone, Debug)]
  13 #[must_use = "iterators are lazy and do nothing unless consumed"]
  14 #[stable(feature = "rust1", since = "1.0.0")]
  15 pub struct Skip<I> {
  16     iter: I,
  17     n: usize,
  18 }
  19
  20 impl<I> Skip<I> {
  21     pub(in crate::iter) fn new(iter: I, n: usize) -> Skip<I> {
  22         Skip { iter, n }
  23     }
  24 }
  25
  26 #[stable(feature = "rust1", since = "1.0.0")]
  27 impl<I> Iterator for Skip<I>
  28 where
  29     I: Iterator,
  30 {
  31     type Item = <I as Iterator>::Item;
  32
  33     #[inline]
  34     fn next(&mut self) -> Option<I::Item> {
  35         if unlikely(self.n > 0) {
  36             self.iter.nth(crate::mem::take(&mut self.n))
  37         } else {
  38             self.iter.next()
  39         }
  40     }
  41
  42     #[inline]
  43     fn nth(&mut self, n: usize) -> Option<I::Item> {
  44         if self.n > 0 {
  45             let skip: usize = crate::mem::take(&mut self.n);
  46             // Checked add to handle overflow case.
  47             let n = match skip.checked_add(n) {
  48                 Some(nth) => nth,
  49                 None => {
  50                     // In case of overflow, load skip value, before loading `n`.
  51                     // Because the amount of elements to iterate is beyond `usize::MAX`, this
  52                     // is split into two `nth` calls where the `skip` `nth` call is discarded.
  53                     self.iter.nth(skip - 1)?;
  54                     n
  55                 }
  56             };
  57             // Load nth element including skip.
  58             self.iter.nth(n)
  59         } else {
  60             self.iter.nth(n)
  61         }
  62     }
  63
  64     #[inline]
  65     fn count(mut self) -> usize {
  66         if self.n > 0 {
  67             // nth(n) skips n+1
  68             if self.iter.nth(self.n - 1).is_none() {
  69                 return 0;
  70             }
  71         }
  72         self.iter.count()
  73     }
  74
  75     #[inline]
  76     fn last(mut self) -> Option<I::Item> {
  77         if self.n > 0 {
  78             // nth(n) skips n+1
  79             self.iter.nth(self.n - 1)?;
  80         }
  81         self.iter.last()
  82     }
  83
  84     #[inline]
  85     fn size_hint(&self) -> (usize, Option<usize>) {
  86         let (lower, upper) = self.iter.size_hint();
  87
  88         let lower = lower.saturating_sub(self.n);
  89         let upper = match upper {
  90             Some(x) => Some(x.saturating_sub(self.n)),
  91             None => None,
  92         };
  93
  94         (lower, upper)
  95     }
  96
  97     #[inline]
  98     fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
  99     where
 100         Self: Sized,
 101         Fold: FnMut(Acc, Self::Item) -> R,
 102         R: Try<Output = Acc>,
 103     {
 104         let n = self.n;
 105         self.n = 0;
 106         if n > 0 {
 107             // nth(n) skips n+1
 108             if self.iter.nth(n - 1).is_none() {
 109                 return try { init };
 110             }
 111         }
 112         self.iter.try_fold(init, fold)
 113     }
 114
 115     #[inline]
 116     fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
 117     where
 118         Fold: FnMut(Acc, Self::Item) -> Acc,
 119     {
 120         if self.n > 0 {
 121             // nth(n) skips n+1
 122             if self.iter.nth(self.n - 1).is_none() {
 123                 return init;
 124             }
 125         }
 126         self.iter.fold(init, fold)
 127     }
 128
 129     #[inline]
 130     #[rustc_inherit_overflow_checks]
 131     fn advance_by(&mut self, n: usize) -> Result<(), usize> {
 132         let mut rem = n;
 133         let step_one = self.n.saturating_add(rem);
 134
 135         match self.iter.advance_by(step_one) {
 136             Ok(_) => {
 137                 rem -= step_one - self.n;
 138                 self.n = 0;
 139             }
 140             Err(advanced) => {
 141                 let advanced_without_skip = advanced.saturating_sub(self.n);
 142                 self.n = self.n.saturating_sub(advanced);
 143                 return if n == 0 { Ok(()) } else { Err(advanced_without_skip) };
 144             }
 145         }
 146
 147         // step_one calculation may have saturated
 148         if unlikely(rem > 0) {
 149             return match self.iter.advance_by(rem) {
 150                 ret @ Ok(_) => ret,
 151                 Err(advanced) => {
 152                     rem -= advanced;
 153                     Err(n - rem)
 154                 }
 155             };
 156         }
 157
 158         Ok(())
 159     }
 160 }
 161
 162 #[stable(feature = "rust1", since = "1.0.0")]
 163 impl<I> ExactSizeIterator for Skip<I> where I: ExactSizeIterator {}
 164
 165 #[stable(feature = "double_ended_skip_iterator", since = "1.9.0")]
 166 impl<I> DoubleEndedIterator for Skip<I>
 167 where
 168     I: DoubleEndedIterator + ExactSizeIterator,
 169 {
 170     fn next_back(&mut self) -> Option<Self::Item> {
 171         if self.len() > 0 { self.iter.next_back() } else { None }
 172     }
 173
 174     #[inline]
 175     fn nth_back(&mut self, n: usize) -> Option<I::Item> {
 176         let len = self.len();
 177         if n < len {
 178             self.iter.nth_back(n)
 179         } else {
 180             if len > 0 {
 181                 // consume the original iterator
 182                 self.iter.nth_back(len - 1);
 183             }
 184             None
 185         }
 186     }
 187
 188     fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
 189     where
 190         Self: Sized,
 191         Fold: FnMut(Acc, Self::Item) -> R,
 192         R: Try<Output = Acc>,
 193     {
 194         fn check<T, Acc, R: Try<Output = Acc>>(
 195             mut n: usize,
 196             mut fold: impl FnMut(Acc, T) -> R,
 197         ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> {
 198             move |acc, x| {
 199                 n -= 1;
 200                 let r = fold(acc, x);
 201                 if n == 0 { ControlFlow::Break(r) } else { ControlFlow::from_try(r) }
 202             }
 203         }
 204
 205         let n = self.len();
 206         if n == 0 { try { init } } else { self.iter.try_rfold(init, check(n, fold)).into_try() }
 207     }
 208
 209     impl_fold_via_try_fold! { rfold -> try_rfold }
 210
 211     #[inline]
 212     fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
 213         let min = crate::cmp::min(self.len(), n);
 214         return match self.iter.advance_back_by(min) {
 215             ret @ Ok(_) if n <= min => ret,
 216             Ok(_) => Err(min),
 217             _ => panic!("ExactSizeIterator contract violation"),
 218         };
 219     }
 220 }
 221
 222 #[stable(feature = "fused", since = "1.26.0")]
 223 impl<I> FusedIterator for Skip<I> where I: FusedIterator {}
 224
 225 #[unstable(issue = "none", feature = "inplace_iteration")]
 226 unsafe impl<I> SourceIter for Skip<I>
 227 where
 228     I: SourceIter,
 229 {
 230     type Source = I::Source;
 231
 232     #[inline]
 233     unsafe fn as_inner(&mut self) -> &mut I::Source {
 234         // SAFETY: unsafe function forwarding to unsafe function with the same requirements
 235         unsafe { SourceIter::as_inner(&mut self.iter) }
 236     }
 237 }
 238
 239 #[unstable(issue = "none", feature = "inplace_iteration")]
 240 unsafe impl<I: InPlaceIterable> InPlaceIterable for Skip<I> {}