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

   1 use crate::{
   2     cmp, iter::FusedIterator, iter::InPlaceIterable, iter::TrustedLen, ops::ControlFlow, ops::Try,
   3 };
   4
   5 use super::SourceIter;
   6
   7 /// An iterator that only iterates over the first `n` iterations of `iter`.
   8 ///
   9 /// This `struct` is created by the [`take`] method on [`Iterator`]. See its
  10 /// documentation for more.
  11 ///
  12 /// [`take`]: Iterator::take
  13 /// [`Iterator`]: trait.Iterator.html
  14 #[derive(Clone, Debug)]
  15 #[must_use = "iterators are lazy and do nothing unless consumed"]
  16 #[stable(feature = "rust1", since = "1.0.0")]
  17 pub struct Take<I> {
  18     iter: I,
  19     n: usize,
  20 }
  21
  22 impl<I> Take<I> {
  23     pub(in crate::iter) fn new(iter: I, n: usize) -> Take<I> {
  24         Take { iter, n }
  25     }
  26 }
  27
  28 #[stable(feature = "rust1", since = "1.0.0")]
  29 impl<I> Iterator for Take<I>
  30 where
  31     I: Iterator,
  32 {
  33     type Item = <I as Iterator>::Item;
  34
  35     #[inline]
  36     fn next(&mut self) -> Option<<I as Iterator>::Item> {
  37         if self.n != 0 {
  38             self.n -= 1;
  39             self.iter.next()
  40         } else {
  41             None
  42         }
  43     }
  44
  45     #[inline]
  46     fn nth(&mut self, n: usize) -> Option<I::Item> {
  47         if self.n > n {
  48             self.n -= n + 1;
  49             self.iter.nth(n)
  50         } else {
  51             if self.n > 0 {
  52                 self.iter.nth(self.n - 1);
  53                 self.n = 0;
  54             }
  55             None
  56         }
  57     }
  58
  59     #[inline]
  60     fn size_hint(&self) -> (usize, Option<usize>) {
  61         if self.n == 0 {
  62             return (0, Some(0));
  63         }
  64
  65         let (lower, upper) = self.iter.size_hint();
  66
  67         let lower = cmp::min(lower, self.n);
  68
  69         let upper = match upper {
  70             Some(x) if x < self.n => Some(x),
  71             _ => Some(self.n),
  72         };
  73
  74         (lower, upper)
  75     }
  76
  77     #[inline]
  78     fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
  79     where
  80         Self: Sized,
  81         Fold: FnMut(Acc, Self::Item) -> R,
  82         R: Try<Ok = Acc>,
  83     {
  84         fn check<'a, T, Acc, R: Try<Ok = Acc>>(
  85             n: &'a mut usize,
  86             mut fold: impl FnMut(Acc, T) -> R + 'a,
  87         ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> + 'a {
  88             move |acc, x| {
  89                 *n -= 1;
  90                 let r = fold(acc, x);
  91                 if *n == 0 { ControlFlow::Break(r) } else { ControlFlow::from_try(r) }
  92             }
  93         }
  94
  95         if self.n == 0 {
  96             try { init }
  97         } else {
  98             let n = &mut self.n;
  99             self.iter.try_fold(init, check(n, fold)).into_try()
 100         }
 101     }
 102
 103     #[inline]
 104     fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
 105     where
 106         Self: Sized,
 107         Fold: FnMut(Acc, Self::Item) -> Acc,
 108     {
 109         #[inline]
 110         fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
 111             move |acc, x| Ok(f(acc, x))
 112         }
 113
 114         self.try_fold(init, ok(fold)).unwrap()
 115     }
 116 }
 117
 118 #[unstable(issue = "none", feature = "inplace_iteration")]
 119 unsafe impl<S: Iterator, I: Iterator> SourceIter for Take<I>
 120 where
 121     I: SourceIter<Source = S>,
 122 {
 123     type Source = S;
 124
 125     #[inline]
 126     unsafe fn as_inner(&mut self) -> &mut S {
 127         // SAFETY: unsafe function forwarding to unsafe function with the same requirements
 128         unsafe { SourceIter::as_inner(&mut self.iter) }
 129     }
 130 }
 131
 132 #[unstable(issue = "none", feature = "inplace_iteration")]
 133 unsafe impl<I: InPlaceIterable> InPlaceIterable for Take<I> {}
 134
 135 #[stable(feature = "double_ended_take_iterator", since = "1.38.0")]
 136 impl<I> DoubleEndedIterator for Take<I>
 137 where
 138     I: DoubleEndedIterator + ExactSizeIterator,
 139 {
 140     #[inline]
 141     fn next_back(&mut self) -> Option<Self::Item> {
 142         if self.n == 0 {
 143             None
 144         } else {
 145             let n = self.n;
 146             self.n -= 1;
 147             self.iter.nth_back(self.iter.len().saturating_sub(n))
 148         }
 149     }
 150
 151     #[inline]
 152     fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
 153         let len = self.iter.len();
 154         if self.n > n {
 155             let m = len.saturating_sub(self.n) + n;
 156             self.n -= n + 1;
 157             self.iter.nth_back(m)
 158         } else {
 159             if len > 0 {
 160                 self.iter.nth_back(len - 1);
 161             }
 162             None
 163         }
 164     }
 165
 166     #[inline]
 167     fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
 168     where
 169         Self: Sized,
 170         Fold: FnMut(Acc, Self::Item) -> R,
 171         R: Try<Ok = Acc>,
 172     {
 173         if self.n == 0 {
 174             try { init }
 175         } else {
 176             let len = self.iter.len();
 177             if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() {
 178                 try { init }
 179             } else {
 180                 self.iter.try_rfold(init, fold)
 181             }
 182         }
 183     }
 184
 185     #[inline]
 186     fn rfold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
 187     where
 188         Self: Sized,
 189         Fold: FnMut(Acc, Self::Item) -> Acc,
 190     {
 191         if self.n == 0 {
 192             init
 193         } else {
 194             let len = self.iter.len();
 195             if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() {
 196                 init
 197             } else {
 198                 self.iter.rfold(init, fold)
 199             }
 200         }
 201     }
 202 }
 203
 204 #[stable(feature = "rust1", since = "1.0.0")]
 205 impl<I> ExactSizeIterator for Take<I> where I: ExactSizeIterator {}
 206
 207 #[stable(feature = "fused", since = "1.26.0")]
 208 impl<I> FusedIterator for Take<I> where I: FusedIterator {}
 209
 210 #[unstable(feature = "trusted_len", issue = "37572")]
 211 unsafe impl<I: TrustedLen> TrustedLen for Take<I> {}