Rollup merge of rust-lang#48115 - Centril:feature/iterator_flatten, r=alexcrichton

Add Iterator::flatten

This adds the trait method `.flatten()` on `Iterator` which flattens one level of nesting from an iterator or (into)iterators. The method `.flat_fmap(f)` is then redefined as `.map(f).flatten()`. The implementation of `Flatten` is essentially that of what it was for `FlatMap` but removing the call to `f` at various places.

Hopefully the type alias approach should be OK as was indicated / alluded to by @bluss and @eddyb in rust-lang/rfcs#2306 (comment).

cc @scottmcm

Author: kennytm

src/libcore/iter/iterator.rs

@@ -12,7 +12,8 @@ use cmp::Ordering;
 use ops::Try;
 
 use super::{AlwaysOk, LoopState};
-use super::{Chain, Cycle, Cloned, Enumerate, Filter, FilterMap, FlatMap, Fuse};
+use super::{Chain, Cycle, Cloned, Enumerate, Filter, FilterMap, Fuse};
+use super::{Flatten, FlatMap, flatten_compat};
 use super::{Inspect, Map, Peekable, Scan, Skip, SkipWhile, StepBy, Take, TakeWhile, Rev};
 use super::{Zip, Sum, Product};
 use super::{ChainState, FromIterator, ZipImpl};
@@ -997,11 +998,15 @@ pub trait Iterator {
     /// an extra layer of indirection. `flat_map()` will remove this extra layer
     /// on its own.
     ///
+    /// You can think of [`flat_map(f)`][flat_map] as the semantic equivalent
+    /// of [`map`]ping, and then [`flatten`]ing as in `map(f).flatten()`.
+    ///
     /// Another way of thinking about `flat_map()`: [`map`]'s closure returns
     /// one item for each element, and `flat_map()`'s closure returns an
     /// iterator for each element.
     ///
     /// [`map`]: #method.map
+    /// [`flatten`]: #method.flatten
     ///
     /// # Examples
     ///
@@ -1021,7 +1026,79 @@ pub trait Iterator {
     fn flat_map<U, F>(self, f: F) -> FlatMap<Self, U, F>
         where Self: Sized, U: IntoIterator, F: FnMut(Self::Item) -> U,
     {
-        FlatMap{iter: self, f: f, frontiter: None, backiter: None }
+        FlatMap { inner: flatten_compat(self.map(f)) }
+    }
+
+    /// 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");
+    /// ```
+    ///
+    /// You can also rewrite this in terms of [`flat_map()`] which is preferable
+    /// in this case since that conveys intent clearer:
+    ///
+    /// ```
+    /// let words = ["alpha", "beta", "gamma"];
+    ///
+    /// // chars() returns an iterator
+    /// let merged: String = words.iter()
+    ///                           .flat_map(|s| s.chars())
+    ///                           .collect();
+    /// assert_eq!(merged, "alphabetagamma");
+    /// ```
+    ///
+    /// Flattening once only removes one level of nesting:
+    ///
+    /// ```
+    /// #![feature(iterator_flatten)]
+    ///
+    /// let d3 = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]];
+    ///
+    /// let d2 = d3.iter().flatten().collect::<Vec<_>>();
+    /// assert_eq!(d2, [&[1, 2], &[3, 4], &[5, 6], &[7, 8]]);
+    ///
+    /// let d1 = d3.iter().flatten().flatten().collect::<Vec<_>>();
+    /// assert_eq!(d1, [&1, &2, &3, &4, &5, &6, &7, &8]);
+    /// ```
+    ///
+    /// Here we see that `flatten()` does not perform a "deep" flatten.
+    /// Instead, only one level of nesting is removed. That is, if you
+    /// `flatten()` a three-dimensional array the result will be
+    /// two-dimensional and not one-dimensional. To get a one-dimensional
+    /// structure, you have to `flatten()` again.
+    #[inline]
+    #[unstable(feature = "iterator_flatten", issue = "48213")]
+    fn flatten(self) -> Flatten<Self>
+    where Self: Sized, Self::Item: IntoIterator {
+        Flatten { inner: flatten_compat(self) }
     }
 
     /// Creates an iterator which ends after the first [`None`].

src/libcore/iter/mod.rs

@@ -2410,24 +2410,23 @@ impl<B, I, St, F> Iterator for Scan<I, St, F> where
 /// [`Iterator`]: trait.Iterator.html
 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
 #[stable(feature = "rust1", since = "1.0.0")]
-#[derive(Clone)]
 pub struct FlatMap<I, U: IntoIterator, F> {
-    iter: I,
-    f: F,
-    frontiter: Option<U::IntoIter>,
-    backiter: Option<U::IntoIter>,
+    inner: FlattenCompat<Map<I, F>, <U as IntoIterator>::IntoIter>
+}
+
+#[stable(feature = "rust1", since = "1.0.0")]
+impl<I: Clone, U: Clone + IntoIterator, F: Clone> Clone for FlatMap<I, U, F>
+    where <U as IntoIterator>::IntoIter: Clone
+{
+    fn clone(&self) -> Self { FlatMap { inner: self.inner.clone() } }
 }
 
 #[stable(feature = "core_impl_debug", since = "1.9.0")]
 impl<I: fmt::Debug, U: IntoIterator, F> fmt::Debug for FlatMap<I, U, F>
     where U::IntoIter: fmt::Debug
 {
     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()
+        f.debug_struct("FlatMap").field("inner", &self.inner).finish()
     }
 }
 
@@ -2437,17 +2436,173 @@ impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
 {
     type Item = U::Item;
 
+    #[inline]
+    fn next(&mut self) -> Option<U::Item> { self.inner.next() }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
+
+    #[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>
+    {
+        self.inner.try_fold(init, fold)
+    }
+
+    #[inline]
+    fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
+        where Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        self.inner.fold(init, fold)
+    }
+}
+
+#[stable(feature = "rust1", since = "1.0.0")]
+impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F>
+    where F: FnMut(I::Item) -> U,
+          U: IntoIterator,
+          U::IntoIter: DoubleEndedIterator
+{
+    #[inline]
+    fn next_back(&mut self) -> Option<U::Item> { self.inner.next_back() }
+
+    #[inline]
+    fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R where
+        Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try<Ok=Acc>
+    {
+        self.inner.try_rfold(init, fold)
+    }
+
+    #[inline]
+    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
+        where Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        self.inner.rfold(init, fold)
+    }
+}
+
+#[unstable(feature = "fused", issue = "35602")]
+impl<I, U, F> FusedIterator for FlatMap<I, U, F>
+    where I: FusedIterator, U: IntoIterator, F: FnMut(I::Item) -> U {}
+
+/// 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 = "48213")]
+pub struct Flatten<I: Iterator>
+where I::Item: IntoIterator {
+    inner: FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>,
+}
+
+#[unstable(feature = "iterator_flatten", issue = "48213")]
+impl<I, U> fmt::Debug for Flatten<I>
+    where I: Iterator + fmt::Debug, U: Iterator + fmt::Debug,
+          I::Item: IntoIterator<IntoIter = U, Item = U::Item>,
+{
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_struct("Flatten").field("inner", &self.inner).finish()
+    }
+}
+
+#[unstable(feature = "iterator_flatten", issue = "48213")]
+impl<I, U> Clone for Flatten<I>
+    where I: Iterator + Clone, U: Iterator + Clone,
+          I::Item: IntoIterator<IntoIter = U, Item = U::Item>,
+{
+    fn clone(&self) -> Self { Flatten { inner: self.inner.clone() } }
+}
+
+#[unstable(feature = "iterator_flatten", issue = "48213")]
+impl<I, U> Iterator for Flatten<I>
+    where I: Iterator, U: Iterator,
+          I::Item: IntoIterator<IntoIter = U, Item = U::Item>
+{
+    type Item = U::Item;
+
+    #[inline]
+    fn next(&mut self) -> Option<U::Item> { self.inner.next() }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
+
+    #[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>
+    {
+        self.inner.try_fold(init, fold)
+    }
+
+    #[inline]
+    fn fold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
+        where Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        self.inner.fold(init, fold)
+    }
+}
+
+#[unstable(feature = "iterator_flatten", issue = "48213")]
+impl<I, U> DoubleEndedIterator for Flatten<I>
+    where I: DoubleEndedIterator, U: DoubleEndedIterator,
+          I::Item: IntoIterator<IntoIter = U, Item = U::Item>
+{
+    #[inline]
+    fn next_back(&mut self) -> Option<U::Item> { self.inner.next_back() }
+
+    #[inline]
+    fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R where
+        Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try<Ok=Acc>
+    {
+        self.inner.try_rfold(init, fold)
+    }
+
+    #[inline]
+    fn rfold<Acc, Fold>(self, init: Acc, fold: Fold) -> Acc
+        where Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        self.inner.rfold(init, fold)
+    }
+}
+
+#[unstable(feature = "fused", issue = "35602")]
+impl<I, U> FusedIterator for Flatten<I>
+    where I: FusedIterator, U: Iterator,
+          I::Item: IntoIterator<IntoIter = U, Item = U::Item> {}
+
+/// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`.
+fn flatten_compat<I, U>(iter: I) -> FlattenCompat<I, U> {
+    FlattenCompat { iter, frontiter: None, backiter: None }
+}
+
+/// Real logic of both `Flatten` and `FlatMap` which simply delegate to
+/// this type.
+#[derive(Clone, Debug)]
+struct FlattenCompat<I, U> {
+    iter: I,
+    frontiter: Option<U>,
+    backiter: Option<U>,
+}
+
+impl<I, U> Iterator for FlattenCompat<I, U>
+    where I: Iterator, U: Iterator,
+          I::Item: IntoIterator<IntoIter = U, Item = U::Item>
+{
+    type Item = U::Item;
+
     #[inline]
     fn next(&mut self) -> Option<U::Item> {
         loop {
             if let Some(ref mut inner) = self.frontiter {
-                if let Some(x) = inner.by_ref().next() {
-                    return Some(x)
-                }
+                if let elt@Some(_) = inner.next() { return elt }
             }
-            match self.iter.next().map(&mut self.f) {
+            match self.iter.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 +2628,9 @@ impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
         self.frontiter = None;
 
         {
-            let f = &mut self.f;
             let frontiter = &mut self.frontiter;
             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);
                 *frontiter = Some(mid);
                 r
@@ -2497,27 +2651,23 @@ impl<I: Iterator, U: IntoIterator, F> Iterator for FlatMap<I, U, F>
         where Fold: FnMut(Acc, Self::Item) -> Acc,
     {
         self.frontiter.into_iter()
-            .chain(self.iter.map(self.f).map(U::into_iter))
+            .chain(self.iter.map(IntoIterator::into_iter))
             .chain(self.backiter)
             .fold(init, |acc, iter| iter.fold(acc, &mut fold))
     }
 }
 
-#[stable(feature = "rust1", since = "1.0.0")]
-impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F> where
-    F: FnMut(I::Item) -> U,
-    U: IntoIterator,
-    U::IntoIter: DoubleEndedIterator
+impl<I, U> DoubleEndedIterator for FlattenCompat<I, U>
+    where I: DoubleEndedIterator, U: DoubleEndedIterator,
+          I::Item: IntoIterator<IntoIter = U, Item = U::Item>
 {
     #[inline]
     fn next_back(&mut self) -> Option<U::Item> {
         loop {
             if let Some(ref mut inner) = self.backiter {
-                if let Some(y) = inner.next_back() {
-                    return Some(y)
-                }
+                if let elt@Some(_) = inner.next_back() { return elt }
             }
-            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()),
                 next => self.backiter = next.map(IntoIterator::into_iter),
             }
@@ -2534,10 +2684,9 @@ impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F> wher
         self.backiter = None;
 
         {
-            let f = &mut self.f;
             let backiter = &mut self.backiter;
             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);
                 *backiter = Some(mid);
                 r
@@ -2558,16 +2707,12 @@ impl<I: DoubleEndedIterator, U, F> DoubleEndedIterator for FlatMap<I, U, F> wher
         where Fold: FnMut(Acc, Self::Item) -> Acc,
     {
         self.frontiter.into_iter()
-            .chain(self.iter.map(self.f).map(U::into_iter))
+            .chain(self.iter.map(IntoIterator::into_iter))
             .chain(self.backiter)
             .rfold(init, |acc, iter| iter.rfold(acc, &mut fold))
     }
 }
 
-#[unstable(feature = "fused", issue = "35602")]
-impl<I, U, F> FusedIterator for FlatMap<I, U, F>
-    where I: FusedIterator, U: IntoIterator, F: FnMut(I::Item) -> U {}
-
 /// An iterator that yields `None` forever after the underlying iterator
 /// yields `None` once.
 ///