diff --git a/src/libcore/iter/adapters/mod.rs b/src/libcore/iter/adapters/mod.rs index 3b8edc2ad6177..714c53778451b 100644 --- a/src/libcore/iter/adapters/mod.rs +++ b/src/libcore/iter/adapters/mod.rs @@ -1613,6 +1613,103 @@ impl Iterator for TakeWhile impl FusedIterator for TakeWhile 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 { + iter: I, + flag: bool, + f: F, +} + +impl MapWhile { + pub(super) fn new(iter: I, f: F) -> Self { + Self { + iter, + flag: false, + f, + } + } +} + +#[unstable(feature = "iter_map_while", issue = "0")] +impl fmt::Debug for MapWhile { + 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 Iterator for MapWhile +where + F: FnMut(I::Item) -> Option, +{ + type Item = B; + + #[inline] + fn next(&mut self) -> Option { + self.find(|_| true) + } + + #[inline] + fn size_hint(&self) -> (usize, Option) { + 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(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try, + { + fn map<'a, T, B, Acc, R: Try>( + flag: &'a mut bool, + f: &'a mut impl FnMut(T) -> Option, + mut fold: impl FnMut(Acc, B) -> R + 'a, + ) -> impl FnMut(Acc, T) -> LoopState + '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 FusedIterator for MapWhile +where + I: FusedIterator, + F: FnMut(I::Item) -> Option, +{ +} + /// An iterator that skips over `n` elements of `iter`. /// /// This `struct` is created by the [`skip`] method on [`Iterator`]. See its diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index aba8e84d58be5..cc75c2dc8c2a3 100644 --- a/src/libcore/iter/mod.rs +++ b/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")] diff --git a/src/libcore/iter/traits/iterator.rs b/src/libcore/iter/traits/iterator.rs index a272035150a15..85cfe0a01cf4c 100644 --- a/src/libcore/iter/traits/iterator.rs +++ b/src/libcore/iter/traits/iterator.rs @@ -1,3 +1,5 @@ +// ignore-tidy-filelength + use crate::cmp::{self, Ordering}; use crate::ops::{Add, Try}; @@ -5,7 +7,7 @@ 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) {} @@ -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(self, f: F) -> MapWhile + where + Self: Sized, + F: FnMut(Self::Item) -> Option, + { + 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. diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index c9096b713f20e..77735b4203b19 100644 --- a/src/libcore/tests/iter.rs +++ b/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); diff --git a/src/libcore/tests/lib.rs b/src/libcore/tests/lib.rs index 35661356028cb..2260ebb96ec6e 100644 --- a/src/libcore/tests/lib.rs +++ b/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;