1
Fork 0

add Iterator::flatten and redefine flat_map(f) in terms of map(f).flatten()

This commit is contained in:
Mazdak Farrokhzad 2018-02-20 08:26:30 +01:00
parent 27a046e933
commit 6af23f977c
4 changed files with 84 additions and 47 deletions

View file

@ -12,7 +12,7 @@ use cmp::Ordering;
use ops::Try; use ops::Try;
use super::{AlwaysOk, LoopState}; use super::{AlwaysOk, LoopState};
use super::{Chain, Cycle, Cloned, Enumerate, Filter, FilterMap, FlatMap, Fuse}; use super::{Chain, Cycle, Cloned, Enumerate, Filter, FilterMap, Flatten, FlatMap, Fuse};
use super::{Inspect, Map, Peekable, Scan, Skip, SkipWhile, StepBy, Take, TakeWhile, Rev}; use super::{Inspect, Map, Peekable, Scan, Skip, SkipWhile, StepBy, Take, TakeWhile, Rev};
use super::{Zip, Sum, Product}; use super::{Zip, Sum, Product};
use super::{ChainState, FromIterator, ZipImpl}; use super::{ChainState, FromIterator, ZipImpl};
@ -997,11 +997,15 @@ pub trait Iterator {
/// an extra layer of indirection. `flat_map()` will remove this extra layer /// an extra layer of indirection. `flat_map()` will remove this extra layer
/// on its own. /// on its own.
/// ///
/// You can think of [`flat_map(f)`][flat_map] as the equivalent of
/// [`map`]ping, and then [`flatten`]ing as in `map(f).flatten()`.
///
/// Another way of thinking about `flat_map()`: [`map`]'s closure returns /// Another way of thinking about `flat_map()`: [`map`]'s closure returns
/// one item for each element, and `flat_map()`'s closure returns an /// one item for each element, and `flat_map()`'s closure returns an
/// iterator for each element. /// iterator for each element.
/// ///
/// [`map`]: #method.map /// [`map`]: #method.map
/// [`flatten`]: #method.flatten
/// ///
/// # Examples /// # Examples
/// ///
@ -1021,7 +1025,46 @@ pub trait Iterator {
fn flat_map<U, F>(self, f: F) -> FlatMap<Self, U, F> fn flat_map<U, F>(self, f: F) -> FlatMap<Self, U, F>
where Self: Sized, U: IntoIterator, F: FnMut(Self::Item) -> U, where Self: Sized, U: IntoIterator, F: FnMut(Self::Item) -> U,
{ {
FlatMap{iter: self, f: f, frontiter: None, backiter: None } self.map(f).flatten()
}
/// Creates an iterator that flattens nested structure.
///
/// This is useful when you have an iterator of iterators or an iterator of
/// things that can be turned into iterators and you want to remove one
/// level of indirection.
///
/// # Examples
///
/// Basic usage:
///
/// ```
/// #![feature(iterator_flatten)]
///
/// let data = vec![vec![1, 2, 3, 4], vec![5, 6]];
/// let flattened = data.into_iter().flatten().collect::<Vec<u8>>();
/// assert_eq!(flattened, &[1, 2, 3, 4, 5, 6]);
/// ```
///
/// Mapping and then flattening:
///
/// ```
/// #![feature(iterator_flatten)]
///
/// let words = ["alpha", "beta", "gamma"];
///
/// // chars() returns an iterator
/// let merged: String = words.iter()
/// .map(|s| s.chars())
/// .flatten()
/// .collect();
/// assert_eq!(merged, "alphabetagamma");
/// ```
#[inline]
#[unstable(feature = "iterator_flatten", issue = "0")]
fn flatten(self) -> Flatten<Self, <Self::Item as IntoIterator>::IntoIter>
where Self: Sized, Self::Item: IntoIterator {
Flatten { iter: self, frontiter: None, backiter: None }
} }
/// Creates an iterator which ends after the first [`None`]. /// Creates an iterator which ends after the first [`None`].

View file

@ -2403,37 +2403,35 @@ impl<B, I, St, F> Iterator for Scan<I, St, F> where
/// An iterator that maps each element to an iterator, and yields the elements /// An iterator that maps each element to an iterator, and yields the elements
/// of the produced iterators. /// of the produced iterators.
/// ///
/// This `struct` is created by the [`flat_map`] method on [`Iterator`]. See its /// This `type` is created by the [`flat_map`] method on [`Iterator`]. See its
/// documentation for more. /// documentation for more.
/// ///
/// [`flat_map`]: trait.Iterator.html#method.flat_map /// [`flat_map`]: trait.Iterator.html#method.flat_map
/// [`Iterator`]: trait.Iterator.html /// [`Iterator`]: trait.Iterator.html
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"] #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
#[stable(feature = "rust1", since = "1.0.0")] #[stable(feature = "rust1", since = "1.0.0")]
#[derive(Clone)] type FlatMap<I, U, F> = Flatten<Map<I, F>, <U as IntoIterator>::IntoIter>;
pub struct FlatMap<I, U: IntoIterator, F> {
/// An iterator that flattens one level of nesting in an iterator of things
/// that can be turned into iterators.
///
/// This `struct` is created by the [`flatten`] method on [`Iterator`]. See its
/// documentation for more.
///
/// [`flatten`]: trait.Iterator.html#method.flatten
/// [`Iterator`]: trait.Iterator.html
#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
#[unstable(feature = "iterator_flatten", issue = "0")]
#[derive(Clone, Debug)]
pub struct Flatten<I, U> {
iter: I, iter: I,
f: F, frontiter: Option<U>,
frontiter: Option<U::IntoIter>, backiter: Option<U>,
backiter: Option<U::IntoIter>,
} }
#[stable(feature = "core_impl_debug", since = "1.9.0")] #[unstable(feature = "iterator_flatten", issue = "0")]
impl<I: fmt::Debug, U: IntoIterator, F> fmt::Debug for FlatMap<I, U, F> impl<I: Iterator, U: Iterator> Iterator for Flatten<I, U>
where U::IntoIter: fmt::Debug where I::Item: IntoIterator<IntoIter = U, Item = U::Item>
{
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("FlatMap")
.field("iter", &self.iter)
.field("frontiter", &self.frontiter)
.field("backiter", &self.backiter)
.finish()
}
}
#[stable(feature = "rust1", since = "1.0.0")]
impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
where F: FnMut(I::Item) -> U,
{ {
type Item = U::Item; type Item = U::Item;
@ -2441,13 +2439,11 @@ impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
fn next(&mut self) -> Option<U::Item> { fn next(&mut self) -> Option<U::Item> {
loop { loop {
if let Some(ref mut inner) = self.frontiter { if let Some(ref mut inner) = self.frontiter {
if let Some(x) = inner.by_ref().next() { if let elt@Some(_) = inner.next() { return elt }
return Some(x)
}
} }
match self.iter.next().map(&mut self.f) { match self.iter.next() {
None => return self.backiter.as_mut().and_then(|it| it.next()), None => return self.backiter.as_mut().and_then(|it| it.next()),
next => self.frontiter = next.map(IntoIterator::into_iter), Some(inner) => self.frontiter = Some(inner.into_iter()),
} }
} }
} }
@ -2473,10 +2469,9 @@ impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
self.frontiter = None; self.frontiter = None;
{ {
let f = &mut self.f;
let frontiter = &mut self.frontiter; let frontiter = &mut self.frontiter;
init = self.iter.try_fold(init, |acc, x| { init = self.iter.try_fold(init, |acc, x| {
let mut mid = f(x).into_iter(); let mut mid = x.into_iter();
let r = mid.try_fold(acc, &mut fold); let r = mid.try_fold(acc, &mut fold);
*frontiter = Some(mid); *frontiter = Some(mid);
r r
@ -2497,27 +2492,24 @@ impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
where Fold: FnMut(Acc, Self::Item) -> Acc, where Fold: FnMut(Acc, Self::Item) -> Acc,
{ {
self.frontiter.into_iter() self.frontiter.into_iter()
.chain(self.iter.map(self.f).map(U::into_iter)) .chain(self.iter.map(IntoIterator::into_iter))
.chain(self.backiter) .chain(self.backiter)
.fold(init, |acc, iter| iter.fold(acc, &mut fold)) .fold(init, |acc, iter| iter.fold(acc, &mut fold))
} }
} }
#[stable(feature = "rust1", since = "1.0.0")] #[unstable(feature = "iterator_flatten", issue = "0")]
impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F> where impl<I, U> DoubleEndedIterator for Flatten<I, U>
F: FnMut(I::Item) -> U, where I: DoubleEndedIterator, U: DoubleEndedIterator,
U: IntoIterator, I::Item: IntoIterator<IntoIter = U, Item = U::Item>
U::IntoIter: DoubleEndedIterator
{ {
#[inline] #[inline]
fn next_back(&mut self) -> Option<U::Item> { fn next_back(&mut self) -> Option<U::Item> {
loop { loop {
if let Some(ref mut inner) = self.backiter { if let Some(ref mut inner) = self.backiter {
if let Some(y) = inner.next_back() { if let elt@Some(_) = inner.next_back() { return elt }
return Some(y)
}
} }
match self.iter.next_back().map(&mut self.f) { match self.iter.next_back() {
None => return self.frontiter.as_mut().and_then(|it| it.next_back()), None => return self.frontiter.as_mut().and_then(|it| it.next_back()),
next => self.backiter = next.map(IntoIterator::into_iter), next => self.backiter = next.map(IntoIterator::into_iter),
} }
@ -2534,10 +2526,9 @@ impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F> wher
self.backiter = None; self.backiter = None;
{ {
let f = &mut self.f;
let backiter = &mut self.backiter; let backiter = &mut self.backiter;
init = self.iter.try_rfold(init, |acc, x| { init = self.iter.try_rfold(init, |acc, x| {
let mut mid = f(x).into_iter(); let mut mid = x.into_iter();
let r = mid.try_rfold(acc, &mut fold); let r = mid.try_rfold(acc, &mut fold);
*backiter = Some(mid); *backiter = Some(mid);
r r
@ -2558,15 +2549,15 @@ impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F> wher
where Fold: FnMut(Acc, Self::Item) -> Acc, where Fold: FnMut(Acc, Self::Item) -> Acc,
{ {
self.frontiter.into_iter() self.frontiter.into_iter()
.chain(self.iter.map(self.f).map(U::into_iter)) .chain(self.iter.map(IntoIterator::into_iter))
.chain(self.backiter) .chain(self.backiter)
.rfold(init, |acc, iter| iter.rfold(acc, &mut fold)) .rfold(init, |acc, iter| iter.rfold(acc, &mut fold))
} }
} }
#[unstable(feature = "fused", issue = "35602")] #[unstable(feature = "fused", issue = "0")]
impl<I, U, F> FusedIterator for FlatMap<I, U, F> impl<I: FusedIterator, U: Iterator> FusedIterator for Flatten<I, U>
where I: FusedIterator, U: IntoIterator, F: FnMut(I::Item) -> U {} where I::Item: IntoIterator<IntoIter = U, Item = U::Item> {}
/// An iterator that yields `None` forever after the underlying iterator /// An iterator that yields `None` forever after the underlying iterator
/// yields `None` once. /// yields `None` once.

View file

@ -93,6 +93,7 @@
#![feature(doc_spotlight)] #![feature(doc_spotlight)]
#![feature(rustc_const_unstable)] #![feature(rustc_const_unstable)]
#![feature(iterator_repeat_with)] #![feature(iterator_repeat_with)]
#![feature(iterator_flatten)]
#[prelude_import] #[prelude_import]
#[allow(unused)] #[allow(unused)]

View file

@ -836,6 +836,8 @@ fn test_iterator_scan() {
assert_eq!(i, ys.len()); assert_eq!(i, ys.len());
} }
// Note: We test flatten() by testing flat_map().
#[test] #[test]
fn test_iterator_flat_map() { fn test_iterator_flat_map() {
let xs = [0, 3, 6]; let xs = [0, 3, 6];