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

   1 use crate::fmt;
   2 use crate::iter::{adapters::SourceIter, FusedIterator, InPlaceIterable};
   3 use crate::ops::{ControlFlow, Try};
   4
   5 /// An iterator that only accepts elements while `predicate` returns `true`.
   6 ///
   7 /// This `struct` is created by the [`take_while`] method on [`Iterator`]. See its
   8 /// documentation for more.
   9 ///
  10 /// [`take_while`]: Iterator::take_while
  11 /// [`Iterator`]: trait.Iterator.html
  12 #[must_use = "iterators are lazy and do nothing unless consumed"]
  13 #[stable(feature = "rust1", since = "1.0.0")]
  14 #[derive(Clone)]
  15 pub struct TakeWhile<I, P> {
  16     iter: I,
  17     flag: bool,
  18     predicate: P,
  19 }
  20
  21 impl<I, P> TakeWhile<I, P> {
  22     pub(in crate::iter) fn new(iter: I, predicate: P) -> TakeWhile<I, P> {
  23         TakeWhile { iter, flag: false, predicate }
  24     }
  25 }
  26
  27 #[stable(feature = "core_impl_debug", since = "1.9.0")]
  28 impl<I: fmt::Debug, P> fmt::Debug for TakeWhile<I, P> {
  29     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  30         f.debug_struct("TakeWhile").field("iter", &self.iter).field("flag", &self.flag).finish()
  31     }
  32 }
  33
  34 #[stable(feature = "rust1", since = "1.0.0")]
  35 impl<I: Iterator, P> Iterator for TakeWhile<I, P>
  36 where
  37     P: FnMut(&I::Item) -> bool,
  38 {
  39     type Item = I::Item;
  40
  41     #[inline]
  42     fn next(&mut self) -> Option<I::Item> {
  43         if self.flag {
  44             None
  45         } else {
  46             let x = self.iter.next()?;
  47             if (self.predicate)(&x) {
  48                 Some(x)
  49             } else {
  50                 self.flag = true;
  51                 None
  52             }
  53         }
  54     }
  55
  56     #[inline]
  57     fn size_hint(&self) -> (usize, Option<usize>) {
  58         if self.flag {
  59             (0, Some(0))
  60         } else {
  61             let (_, upper) = self.iter.size_hint();
  62             (0, upper) // can't know a lower bound, due to the predicate
  63         }
  64     }
  65
  66     #[inline]
  67     fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
  68     where
  69         Self: Sized,
  70         Fold: FnMut(Acc, Self::Item) -> R,
  71         R: Try<Ok = Acc>,
  72     {
  73         fn check<'a, T, Acc, R: Try<Ok = Acc>>(
  74             flag: &'a mut bool,
  75             p: &'a mut impl FnMut(&T) -> bool,
  76             mut fold: impl FnMut(Acc, T) -> R + 'a,
  77         ) -> impl FnMut(Acc, T) -> ControlFlow<R, Acc> + 'a {
  78             move |acc, x| {
  79                 if p(&x) {
  80                     ControlFlow::from_try(fold(acc, x))
  81                 } else {
  82                     *flag = true;
  83                     ControlFlow::Break(try { acc })
  84                 }
  85             }
  86         }
  87
  88         if self.flag {
  89             try { init }
  90         } else {
  91             let flag = &mut self.flag;
  92             let p = &mut self.predicate;
  93             self.iter.try_fold(init, check(flag, p, fold)).into_try()
  94         }
  95     }
  96
  97     #[inline]
  98     fn fold<Acc, Fold>(mut self, init: Acc, fold: Fold) -> Acc
  99     where
 100         Self: Sized,
 101         Fold: FnMut(Acc, Self::Item) -> Acc,
 102     {
 103         #[inline]
 104         fn ok<B, T>(mut f: impl FnMut(B, T) -> B) -> impl FnMut(B, T) -> Result<B, !> {
 105             move |acc, x| Ok(f(acc, x))
 106         }
 107
 108         self.try_fold(init, ok(fold)).unwrap()
 109     }
 110 }
 111
 112 #[stable(feature = "fused", since = "1.26.0")]
 113 impl<I, P> FusedIterator for TakeWhile<I, P>
 114 where
 115     I: FusedIterator,
 116     P: FnMut(&I::Item) -> bool,
 117 {
 118 }
 119
 120 #[unstable(issue = "none", feature = "inplace_iteration")]
 121 unsafe impl<S: Iterator, P, I: Iterator> SourceIter for TakeWhile<I, P>
 122 where
 123     P: FnMut(&I::Item) -> bool,
 124     I: SourceIter<Source = S>,
 125 {
 126     type Source = S;
 127
 128     #[inline]
 129     unsafe fn as_inner(&mut self) -> &mut S {
 130         // SAFETY: unsafe function forwarding to unsafe function with the same requirements
 131         unsafe { SourceIter::as_inner(&mut self.iter) }
 132     }
 133 }
 134
 135 #[unstable(issue = "none", feature = "inplace_iteration")]
 136 unsafe impl<I: InPlaceIterable, F> InPlaceIterable for TakeWhile<I, F> where
 137     F: FnMut(&I::Item) -> bool
 138 {
 139 }