Add Iterator::map_while method and corresponding MapWhile adapter

Author: WaffleLapkin

src/libcore/iter/adapters/mod.rs

@@ -1752,6 +1752,95 @@ where
     }
 }
 
+/// An iterator that only accepts elements while `predicate` returns `Some(_)`.
+///
+/// This `struct` is created by the [`map_while`] method on [`Iterator`]. See its
+/// documentation for more.
+///
+/// [`map_while`]: trait.Iterator.html#method.map_while
+/// [`Iterator`]: trait.Iterator.html
+#[must_use = "iterators are lazy and do nothing unless consumed"]
+#[unstable(feature = "iter_map_while", reason = "recently added", issue = "none")]
+#[derive(Clone)]
+pub struct MapWhile<I, P> {
+    iter: I,
+    finished: bool,
+    predicate: P,
+}
+
+impl<I, P> MapWhile<I, P> {
+    pub(super) fn new(iter: I, predicate: P) -> MapWhile<I, P> {
+        MapWhile { iter, finished: false, predicate }
+    }
+}
+
+#[unstable(feature = "iter_map_while", reason = "recently added", issue = "none")]
+impl<I: fmt::Debug, P> fmt::Debug for MapWhile<I, P> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("MapWhile").field("iter", &self.iter).field("flag", &self.finished).finish()
+    }
+}
+
+#[unstable(feature = "iter_map_while", reason = "recently added", issue = "none")]
+impl<B, I: Iterator, P> Iterator for MapWhile<I, P>
+where
+    P: FnMut(I::Item) -> Option<B>,
+{
+    type Item = B;
+
+    #[inline]
+    fn next(&mut self) -> Option<B> {
+        if self.finished {
+            None
+        } else {
+            let x = self.iter.next()?;
+            let ret = (self.predicate)(x);
+            self.finished = ret.is_none();
+            ret
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if self.finished {
+            (0, Some(0))
+        } else {
+            let (_, upper) = self.iter.size_hint();
+            (0, upper) // can't know a lower bound, due to the predicate
+        }
+    }
+
+    #[inline]
+    fn try_fold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R
+    where
+        Self: Sized,
+        Fold: FnMut(Acc, Self::Item) -> R,
+        R: Try<Ok = Acc>,
+    {
+        fn check<'a, B, T, Acc, R: Try<Ok = Acc>>(
+            flag: &'a mut bool,
+            p: &'a mut impl FnMut(T) -> Option<B>,
+            mut fold: impl FnMut(Acc, B) -> R + 'a,
+        ) -> impl FnMut(Acc, T) -> LoopState<Acc, R> + 'a {
+            move |acc, x| match p(x) {
+                Some(item) => LoopState::from_try(fold(acc, item)),
+                None => {
+                    *flag = true;
+                    LoopState::Break(Try::from_ok(acc))
+                }
+            }
+        }
+
+        if self.finished {
+            Try::from_ok(init)
+        } else {
+            let flag = &mut self.finished;
+            let p = &mut self.predicate;
+            self.iter.try_fold(init, check(flag, p, fold)).into_try()
+        }
+    }
+}
+
 #[stable(feature = "fused", since = "1.26.0")]
 impl<I, P> FusedIterator for TakeWhile<I, P>
 where

src/libcore/iter/mod.rs

@@ -351,6 +351,8 @@ pub use self::adapters::Cloned;
 pub use self::adapters::Copied;
 #[stable(feature = "iterator_flatten", since = "1.29.0")]
 pub use self::adapters::Flatten;
+#[unstable(feature = "iter_map_while", reason = "recently added", issue = "none")]
+pub use self::adapters::MapWhile;
 #[stable(feature = "iterator_step_by", since = "1.28.0")]
 pub use self::adapters::StepBy;
 #[stable(feature = "rust1", since = "1.0.0")]

src/libcore/iter/traits/iterator.rs

@@ -1,4 +1,6 @@
 // ignore-tidy-filelength
+// This file almost exclusively consists of the definition of `Iterator`. We
+// can't split that into multiple files.
 
 use crate::cmp::{self, Ordering};
 use crate::ops::{Add, Try};
@@ -7,7 +9,9 @@ use super::super::LoopState;
 use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse};
 use super::super::{FlatMap, Flatten};
 use super::super::{FromIterator, Product, Sum, Zip};
-use super::super::{Inspect, Map, Peekable, Rev, Scan, Skip, SkipWhile, StepBy, Take, TakeWhile};
+use super::super::{
+    Inspect, Map, MapWhile, Peekable, Rev, Scan, Skip, SkipWhile, StepBy, Take, TakeWhile,
+};
 
 fn _assert_is_object_safe(_: &dyn Iterator<Item = ()>) {}
 
@@ -1026,6 +1030,102 @@ pub trait Iterator {
         TakeWhile::new(self, predicate)
     }
 
+    /// Creates an iterator that both yields elements based on a predicate and maps.
+    ///
+    /// `map_while()` takes a closure as an argument. It will call this
+    /// closure on each element of the iterator, and yield elements
+    /// while it returns [`Some(_)`][`Some`].
+    ///
+    /// After [`None`] is returned, `map_while()`'s job is over, and the
+    /// rest of the elements are ignored.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(iter_map_while)]
+    /// let a = [-1i32, 4, 0, 1];
+    ///
+    /// let mut iter = a.iter().map_while(|x| 16i32.checked_div(*x));
+    ///
+    /// assert_eq!(iter.next(), Some(-16));
+    /// assert_eq!(iter.next(), Some(4));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    ///
+    /// Here's the same example, but with [`take_while`] and [`map`]:
+    ///
+    /// [`take_while`]: #method.take_while
+    /// [`map`]: #method.map
+    ///
+    /// ```
+    /// let a = [-1i32, 4, 0, 1];
+    ///
+    /// let mut iter = a.iter()
+    ///                 .map(|x| 16i32.checked_div(*x))
+    ///                 .take_while(|x| x.is_some())
+    ///                 .map(|x| x.unwrap());
+    ///
+    /// assert_eq!(iter.next(), Some(-16));
+    /// assert_eq!(iter.next(), Some(4));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    ///
+    /// Stopping after an initial [`None`]:
+    ///
+    /// ```
+    /// #![feature(iter_map_while)]
+    /// use std::convert::TryFrom;
+    ///
+    /// let a = [0, -1, 1, -2];
+    ///
+    /// let mut iter = a.iter().map_while(|x| u32::try_from(*x).ok());
+    ///
+    /// assert_eq!(iter.next(), Some(0u32));
+    ///
+    /// // We have more elements that are fit in u32, but since we already
+    /// // got a None, map_while() isn't used any more
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    ///
+    /// Because `map_while()` needs to look at the value in order to see if it
+    /// should be included or not, consuming iterators will see that it is
+    /// removed:
+    ///
+    /// ```
+    /// #![feature(iter_map_while)]
+    /// use std::convert::TryFrom;
+    ///
+    /// let a = [1, 2, -3, 4];
+    /// let mut iter = a.iter();
+    ///
+    /// let result: Vec<u32> = iter.by_ref()
+    ///                            .map_while(|n| u32::try_from(*n).ok())
+    ///                            .collect();
+    ///
+    /// assert_eq!(result, &[1, 2]);
+    ///
+    /// let result: Vec<i32> = iter.cloned().collect();
+    ///
+    /// assert_eq!(result, &[4]);
+    /// ```
+    ///
+    /// The `-3` is no longer there, because it was consumed in order to see if
+    /// the iteration should stop, but wasn't placed back into the iterator.
+    ///
+    /// [`Some`]: ../../std/option/enum.Option.html#variant.Some
+    /// [`None`]: ../../std/option/enum.Option.html#variant.None
+    #[inline]
+    #[unstable(feature = "iter_map_while", reason = "recently added", issue = "none")]
+    fn map_while<B, P>(self, predicate: P) -> MapWhile<Self, P>
+    where
+        Self: Sized,
+        P: FnMut(Self::Item) -> Option<B>,
+    {
+        MapWhile::new(self, predicate)
+    }
+
     /// Creates an iterator that skips the first `n` elements.
     ///
     /// After they have been consumed, the rest of the elements are yielded.

src/libcore/tests/iter.rs

@@ -1477,6 +1477,7 @@ fn test_iterator_size_hint() {
     assert_eq!(c.clone().take(5).size_hint(), (5, Some(5)));
     assert_eq!(c.clone().skip(5).size_hint().1, None);
     assert_eq!(c.clone().take_while(|_| false).size_hint(), (0, None));
+    assert_eq!(c.clone().map_while(|_| None::<()>).size_hint(), (0, None));
     assert_eq!(c.clone().skip_while(|_| false).size_hint(), (0, None));
     assert_eq!(c.clone().enumerate().size_hint(), (usize::MAX, None));
     assert_eq!(c.clone().chain(vi.clone().cloned()).size_hint(), (usize::MAX, None));
@@ -1491,6 +1492,7 @@ fn test_iterator_size_hint() {
     assert_eq!(vi.clone().skip(3).size_hint(), (7, Some(7)));
     assert_eq!(vi.clone().skip(12).size_hint(), (0, Some(0)));
     assert_eq!(vi.clone().take_while(|_| false).size_hint(), (0, Some(10)));
+    assert_eq!(vi.clone().map_while(|_| None::<()>).size_hint(), (0, Some(10)));
     assert_eq!(vi.clone().skip_while(|_| false).size_hint(), (0, Some(10)));
     assert_eq!(vi.clone().enumerate().size_hint(), (10, Some(10)));
     assert_eq!(vi.clone().chain(v2).size_hint(), (13, Some(13)));

src/libcore/tests/lib.rs

@@ -36,6 +36,7 @@
 #![feature(iter_is_partitioned)]
 #![feature(iter_order_by)]
 #![feature(cmp_min_max_by)]
+#![feature(iter_map_while)]
 #![feature(const_slice_from_raw_parts)]
 #![feature(const_raw_ptr_deref)]
 #![feature(never_type)]