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

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