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

   1 use crate::iter::adapters::{zip::try_get_unchecked, SourceIter, TrustedRandomAccess};
   2 use crate::iter::{FusedIterator, InPlaceIterable, TrustedLen};
   3 use crate::ops::{Add, AddAssign, Try};
   4
   5 /// An iterator that yields the current count and the element during iteration.
   6 ///
   7 /// This `struct` is created by the [`enumerate`] method on [`Iterator`]. See its
   8 /// documentation for more.
   9 ///
  10 /// [`enumerate`]: Iterator::enumerate
  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 Enumerate<I> {
  16     iter: I,
  17     count: usize,
  18 }
  19 impl<I> Enumerate<I> {
  20     pub(in crate::iter) fn new(iter: I) -> Enumerate<I> {
  21         Enumerate { iter, count: 0 }
  22     }
  23 }
  24
  25 #[stable(feature = "rust1", since = "1.0.0")]
  26 impl<I> Iterator for Enumerate<I>
  27 where
  28     I: Iterator,
  29 {
  30     type Item = (usize, <I as Iterator>::Item);
  31
  32     /// # Overflow Behavior
  33     ///
  34     /// The method does no guarding against overflows, so enumerating more than
  35     /// `usize::MAX` elements either produces the wrong result or panics. If
  36     /// debug assertions are enabled, a panic is guaranteed.
  37     ///
  38     /// # Panics
  39     ///
  40     /// Might panic if the index of the element overflows a `usize`.
  41     #[inline]
  42     fn next(&mut self) -> Option<(usize, <I as Iterator>::Item)> {
  43         let a = self.iter.next()?;
  44         let i = self.count;
  45         // Possible undefined overflow.
  46         AddAssign::add_assign(&mut self.count, 1);
  47         Some((i, a))
  48     }
  49
  50     #[inline]
  51     fn size_hint(&self) -> (usize, Option<usize>) {
  52         self.iter.size_hint()
  53     }
  54
  55     #[inline]
  56     fn nth(&mut self, n: usize) -> Option<(usize, I::Item)> {
  57         let a = self.iter.nth(n)?;
  58         // Possible undefined overflow.
  59         let i = Add::add(self.count, n);
  60         self.count = Add::add(i, 1);
  61         Some((i, a))
  62     }
  63
  64     #[inline]
  65     fn count(self) -> usize {
  66         self.iter.count()
  67     }
  68
  69     #[inline]
  70     fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
  71     where
  72         Self: Sized,
  73         Fold: FnMut(Acc, Self::Item) -> R,
  74         R: Try<Ok = Acc>,
  75     {
  76         #[inline]
  77         fn enumerate<'a, T, Acc, R>(
  78             count: &'a mut usize,
  79             mut fold: impl FnMut(Acc, (usize, T)) -> R + 'a,
  80         ) -> impl FnMut(Acc, T) -> R + 'a {
  81             move |acc, item| {
  82                 let acc = fold(acc, (*count, item));
  83                 // Possible undefined overflow.
  84                 AddAssign::add_assign(count, 1);
  85                 acc
  86             }
  87         }
  88
  89         self.iter.try_fold(init, enumerate(&mut self.count, fold))
  90     }
  91
  92     #[inline]
  93     fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
  94     where
  95         Fold: FnMut(Acc, Self::Item) -> Acc,
  96     {
  97         #[inline]
  98         fn enumerate<T, Acc>(
  99             mut count: usize,
 100             mut fold: impl FnMut(Acc, (usize, T)) -> Acc,
 101         ) -> impl FnMut(Acc, T) -> Acc {
 102             move |acc, item| {
 103                 let acc = fold(acc, (count, item));
 104                 // Possible undefined overflow.
 105                 AddAssign::add_assign(&mut count, 1);
 106                 acc
 107             }
 108         }
 109
 110         self.iter.fold(init, enumerate(self.count, fold))
 111     }
 112
 113     unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> <Self as Iterator>::Item
 114     where
 115         Self: TrustedRandomAccess,
 116     {
 117         // SAFETY: the caller must uphold the contract for
 118         // `Iterator::__iterator_get_unchecked`.
 119         let value = unsafe { try_get_unchecked(&mut self.iter, idx) };
 120         (Add::add(self.count, idx), value)
 121     }
 122 }
 123
 124 #[stable(feature = "rust1", since = "1.0.0")]
 125 impl<I> DoubleEndedIterator for Enumerate<I>
 126 where
 127     I: ExactSizeIterator + DoubleEndedIterator,
 128 {
 129     #[inline]
 130     fn next_back(&mut self) -> Option<(usize, <I as Iterator>::Item)> {
 131         let a = self.iter.next_back()?;
 132         let len = self.iter.len();
 133         // Can safely add, `ExactSizeIterator` promises that the number of
 134         // elements fits into a `usize`.
 135         Some((self.count + len, a))
 136     }
 137
 138     #[inline]
 139     fn nth_back(&mut self, n: usize) -> Option<(usize, <I as Iterator>::Item)> {
 140         let a = self.iter.nth_back(n)?;
 141         let len = self.iter.len();
 142         // Can safely add, `ExactSizeIterator` promises that the number of
 143         // elements fits into a `usize`.
 144         Some((self.count + len, a))
 145     }
 146
 147     #[inline]
 148     fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
 149     where
 150         Self: Sized,
 151         Fold: FnMut(Acc, Self::Item) -> R,
 152         R: Try<Ok = Acc>,
 153     {
 154         // Can safely add and subtract the count, as `ExactSizeIterator` promises
 155         // that the number of elements fits into a `usize`.
 156         fn enumerate<T, Acc, R>(
 157             mut count: usize,
 158             mut fold: impl FnMut(Acc, (usize, T)) -> R,
 159         ) -> impl FnMut(Acc, T) -> R {
 160             move |acc, item| {
 161                 count -= 1;
 162                 fold(acc, (count, item))
 163             }
 164         }
 165
 166         let count = self.count + self.iter.len();
 167         self.iter.try_rfold(init, enumerate(count, fold))
 168     }
 169
 170     #[inline]
 171     fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
 172     where
 173         Fold: FnMut(Acc, Self::Item) -> Acc,
 174     {
 175         // Can safely add and subtract the count, as `ExactSizeIterator` promises
 176         // that the number of elements fits into a `usize`.
 177         fn enumerate<T, Acc>(
 178             mut count: usize,
 179             mut fold: impl FnMut(Acc, (usize, T)) -> Acc,
 180         ) -> impl FnMut(Acc, T) -> Acc {
 181             move |acc, item| {
 182                 count -= 1;
 183                 fold(acc, (count, item))
 184             }
 185         }
 186
 187         let count = self.count + self.iter.len();
 188         self.iter.rfold(init, enumerate(count, fold))
 189     }
 190 }
 191
 192 #[stable(feature = "rust1", since = "1.0.0")]
 193 impl<I> ExactSizeIterator for Enumerate<I>
 194 where
 195     I: ExactSizeIterator,
 196 {
 197     fn len(&self) -> usize {
 198         self.iter.len()
 199     }
 200
 201     fn is_empty(&self) -> bool {
 202         self.iter.is_empty()
 203     }
 204 }
 205
 206 #[doc(hidden)]
 207 #[unstable(feature = "trusted_random_access", issue = "none")]
 208 unsafe impl<I> TrustedRandomAccess for Enumerate<I>
 209 where
 210     I: TrustedRandomAccess,
 211 {
 212     fn may_have_side_effect() -> bool {
 213         I::may_have_side_effect()
 214     }
 215 }
 216
 217 #[stable(feature = "fused", since = "1.26.0")]
 218 impl<I> FusedIterator for Enumerate<I> where I: FusedIterator {}
 219
 220 #[unstable(feature = "trusted_len", issue = "37572")]
 221 unsafe impl<I> TrustedLen for Enumerate<I> where I: TrustedLen {}
 222
 223 #[unstable(issue = "none", feature = "inplace_iteration")]
 224 unsafe impl<S: Iterator, I: Iterator> SourceIter for Enumerate<I>
 225 where
 226     I: SourceIter<Source = S>,
 227 {
 228     type Source = S;
 229
 230     #[inline]
 231     unsafe fn as_inner(&mut self) -> &mut S {
 232         // SAFETY: unsafe function forwarding to unsafe function with the same requirements
 233         unsafe { SourceIter::as_inner(&mut self.iter) }
 234     }
 235 }
 236
 237 #[unstable(issue = "none", feature = "inplace_iteration")]
 238 unsafe impl<I: InPlaceIterable> InPlaceIterable for Enumerate<I> {}