Add Iterator::map_while

src/libcore/iter/adapters/mod.rs

@@ -1613,6 +1613,103 @@ impl<I: Iterator, P> Iterator for TakeWhile<I, P>
 impl<I, P> FusedIterator for TakeWhile<I, P>
     where I: FusedIterator, P: FnMut(&I::Item) -> bool {}
 
+/// An iterator that only transforms elements as long as the transformation succeeds.
+///
+/// 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", issue = "0")]
+#[derive(Clone)]
+pub struct MapWhile<I, F> {
+    iter: I,
+    flag: bool,
+    f: F,
+}
+
+impl<I, F> MapWhile<I, F> {
+    pub(super) fn new(iter: I, f: F) -> Self {
+        Self {
+            iter,
+            flag: false,
+            f,
+        }
+    }
+}
+
+#[unstable(feature = "iter_map_while", issue = "0")]
+impl<I: fmt::Debug, F> fmt::Debug for MapWhile<I, F> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("MapWhile")
+            .field("iter", &self.iter)
+            .field("flag", &self.flag)
+            .finish()
+    }
+}
+
+#[unstable(feature = "iter_map_while", issue = "0")]
+impl<I: Iterator, B, F> Iterator for MapWhile<I, F>
+where
+    F: FnMut(I::Item) -> Option<B>,
+{
+    type Item = B;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        self.find(|_| true)
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if self.flag {
+            (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 map<'a, T, B, Acc, R: Try<Ok = Acc>>(
+            flag: &'a mut bool,
+            f: &'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 f(x) {
+                None => {
+                    *flag = true;
+                    LoopState::Break(Try::from_ok(acc))
+                }
+                Some(x) => LoopState::from_try(fold(acc, x)),
+            }
+        }
+
+        if self.flag {
+            Try::from_ok(init)
+        } else {
+            let flag = &mut self.flag;
+            let f = &mut self.f;
+            self.iter.try_fold(init, map(flag, f, fold)).into_try()
+        }
+    }
+}
+
+#[unstable(feature = "iter_map_while", issue = "0")]
+impl<I, B, F> FusedIterator for MapWhile<I, F>
+where
+    I: FusedIterator,
+    F: FnMut(I::Item) -> Option<B>,
+{
+}
+
 /// An iterator that skips over `n` elements of `iter`.
 ///
 /// This `struct` is created by the [`skip`] method on [`Iterator`]. See its

src/libcore/iter/mod.rs

@@ -348,7 +348,7 @@ pub use self::traits::TrustedLen;
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::adapters::{Rev, Cycle, Chain, Zip, Map, Filter, FilterMap, Enumerate};
 #[stable(feature = "rust1", since = "1.0.0")]
-pub use self::adapters::{Peekable, SkipWhile, TakeWhile, Skip, Take, Scan, FlatMap};
+pub use self::adapters::{Peekable, SkipWhile, TakeWhile, MapWhile, Skip, Take, Scan, FlatMap};
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::adapters::{Fuse, Inspect};
 #[stable(feature = "iter_cloned", since = "1.1.0")]

src/libcore/iter/traits/iterator.rs

@@ -1,11 +1,13 @@
+// ignore-tidy-filelength
+
 use crate::cmp::{self, Ordering};
 use crate::ops::{Add, Try};
 
 use super::super::LoopState;
 use super::super::{Chain, Cycle, Copied, Cloned, Enumerate, Filter, FilterMap, Fuse};
 use super::super::{Flatten, FlatMap};
 use super::super::{Inspect, Map, Peekable, Scan, Skip, SkipWhile, StepBy, Take, TakeWhile, Rev};
-use super::super::{Zip, Sum, Product, FromIterator};
+use super::super::{Zip, Sum, Product, FromIterator, MapWhile};
 
 fn _assert_is_object_safe(_: &dyn Iterator<Item=()>) {}
 
@@ -973,6 +975,57 @@ pub trait Iterator {
         TakeWhile::new(self, predicate)
     }
 
+    /// Creates an iterator that transforms elements until a transformation fails.
+    ///
+    /// `map_while` creates an iterator that calls the given closure on each
+    /// element. If the closure returns [`Some(element)`][`Some`], then that
+    /// element is returned. If the closure returns [`None`], then iteration stops.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(iter_map_while)]
+    ///
+    /// let iter = (0..).map_while(|x| std::char::from_digit(x, 16));
+    /// assert!(iter.eq("0123456789abcdef".chars()));
+    /// ```
+    ///
+    /// Note that `map_while` consumes the first element of the iterator for
+    /// which the transformation fails:
+    ///
+    /// ```
+    /// #![feature(iter_map_while)]
+    ///
+    /// use std::convert::TryFrom;
+    ///
+    /// let xs: &[i32] = &[2, 1, 0, -1, -2];
+    /// let mut iter = xs.iter();
+    ///
+    /// // the first three elements are collected, and the fourth element is discarded
+    /// let vec: Vec<_> = iter
+    ///     .by_ref()
+    ///     .map_while(|&x| u32::try_from(x).ok())
+    ///     .collect();
+    ///
+    /// assert_eq!(vec, vec![2, 1, 0]);
+    /// assert_eq!(iter.next(), Some(&-2));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    ///
+    /// [`Some`]: ../../std/option/enum.Option.html#variant.Some
+    /// [`None`]: ../../std/option/enum.Option.html#variant.None
+    #[inline]
+    #[unstable(feature = "iter_map_while", issue = "0")]
+    fn map_while<B, F>(self, f: F) -> MapWhile<Self, F>
+    where
+        Self: Sized,
+        F: FnMut(Self::Item) -> Option<B>,
+    {
+        MapWhile::new(self, f)
+    }
+
     /// 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

@@ -872,6 +872,16 @@ fn test_iterator_take_while() {
     assert_eq!(i, ys.len());
 }
 
+#[test]
+fn test_iterator_map_while() {
+    let xs = [3, 2, 1, 0, 1, 2, 3];
+    let mut iter = xs.iter().map_while(|&x| 6_i32.checked_div(x));
+    assert_eq!(iter.next(), Some(2));
+    assert_eq!(iter.next(), Some(3));
+    assert_eq!(iter.next(), Some(6));
+    assert_eq!(iter.next(), None);
+}
+
 #[test]
 fn test_iterator_skip_while() {
     let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19];
@@ -1479,6 +1489,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(Some).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));
@@ -1493,6 +1504,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(Some).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)));
@@ -2664,6 +2676,23 @@ fn test_take_while_folds() {
     assert_eq!(iter.next(), Some(20));
 }
 
+#[test]
+fn test_map_while_folds() {
+    let mut iter = "defghi".chars().map_while(|c| c.to_digit(16));
+    let mut next = || iter.try_fold((), |(), x| Err(x)).err();
+    assert_eq!(next(), Some(13));
+    assert_eq!(next(), Some(14));
+    assert_eq!(next(), Some(15));
+    assert_eq!(next(), None);
+
+    let mut iter = "abcdefghi".chars().map_while(|c| c.to_digit(16));
+    let f = |x: u8, y: u32| x.checked_mul(3)?.checked_add(y as u8);
+    assert_eq!(iter.try_fold(0, f), None);
+    assert_eq!(iter.clone().next(), Some(14));
+    assert_eq!(iter.try_fold(0, f), Some(57));
+    assert_eq!(iter.next(), None);
+}
+
 #[test]
 fn test_skip_try_folds() {
     let f = &|acc, x| i32::checked_add(2*acc, x);

src/libcore/tests/lib.rs

@@ -35,6 +35,7 @@
 #![feature(iter_is_partitioned)]
 #![feature(iter_order_by)]
 #![feature(cmp_min_max_by)]
+#![feature(iter_map_while)]
 
 extern crate test;