diff --git a/src/libcore/iter/adapters/mod.rs b/src/libcore/iter/adapters/mod.rs index 5787b9174edab..6ef4a49434002 100644 --- a/src/libcore/iter/adapters/mod.rs +++ b/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 { + iter: I, + finished: bool, + predicate: P, +} + +impl MapWhile { + pub(super) fn new(iter: I, predicate: P) -> MapWhile { + MapWhile { iter, finished: false, predicate } + } +} + +#[unstable(feature = "iter_map_while", reason = "recently added", issue = "none")] +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.finished).finish() + } +} + +#[unstable(feature = "iter_map_while", reason = "recently added", issue = "none")] +impl Iterator for MapWhile +where + P: FnMut(I::Item) -> Option, +{ + type Item = B; + + #[inline] + fn next(&mut self) -> Option { + 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) { + 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(&mut self, init: Acc, fold: Fold) -> R + where + Self: Sized, + Fold: FnMut(Acc, Self::Item) -> R, + R: Try, + { + fn check<'a, B, T, Acc, R: Try>( + flag: &'a mut bool, + p: &'a mut impl FnMut(T) -> Option, + mut fold: impl FnMut(Acc, B) -> R + 'a, + ) -> impl FnMut(Acc, T) -> LoopState + '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 FusedIterator for TakeWhile where diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index 0d5af3986fbff..c9f5e77d14b3e 100644 --- a/src/libcore/iter/mod.rs +++ b/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")] diff --git a/src/libcore/iter/traits/iterator.rs b/src/libcore/iter/traits/iterator.rs index 21a569867b178..6bb92099ced51 100644 --- a/src/libcore/iter/traits/iterator.rs +++ b/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) {} @@ -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 = iter.by_ref() + /// .map_while(|n| u32::try_from(*n).ok()) + /// .collect(); + /// + /// assert_eq!(result, &[1, 2]); + /// + /// let result: Vec = 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(self, predicate: P) -> MapWhile + where + Self: Sized, + P: FnMut(Self::Item) -> Option, + { + 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. diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index 8b8dc941534ee..bd3218ec27f32 100644 --- a/src/libcore/tests/iter.rs +++ b/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))); diff --git a/src/libcore/tests/lib.rs b/src/libcore/tests/lib.rs index 567c840b01ade..8fd19ef67fccf 100644 --- a/src/libcore/tests/lib.rs +++ b/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)]