Add `Iterator::map_while` method and corresponding `MapWhile` adapter

This commit is contained in:
Waffle 2020-01-24 14:49:34 +03:00
parent b181835a6b
commit 1aff08010d
5 changed files with 195 additions and 1 deletions

View File

@ -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

View File

@ -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")]

View File

@ -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.

View File

@ -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)));

View File

@ -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)]