core::iter::Flatten: update FlatMap & Flatten according to discussion

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, Flatten, 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,8 +998,8 @@ 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 equivalent of
+    /// You can think of [`flat_map(f)`][flat_map] as the semantic equivalent
-    /// [`map`]ping, and then [`flatten`]ing as in `map(f).flatten()`.
+    /// 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
@@ -1025,7 +1026,7 @@ 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,
     {
-        self.map(f).flatten()
+        FlatMap { inner: flatten_compat(self.map(f)) }
     }
 
     /// Creates an iterator that flattens nested structure.
@@ -1060,11 +1061,24 @@ pub trait Iterator {
     ///                           .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");
+    /// ```
     #[inline]
     #[unstable(feature = "iterator_flatten", issue = "0")]
-    fn flatten(self) -> Flatten<Self, <Self::Item as IntoIterator>::IntoIter>
+    fn flatten(self) -> Flatten<Self>
     where Self: Sized, Self::Item: IntoIterator {
-        Flatten { iter: self, frontiter: None, backiter: None }
+        Flatten { inner: flatten_compat(self) }
     }
 
     /// Creates an iterator which ends after the first [`None`].

src/libcore/iter/mod.rs

@@ -2403,14 +2403,87 @@ 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
 /// of the produced iterators.
 ///
-/// This `type` is created by the [`flat_map`] method on [`Iterator`]. See its
+/// This `struct` is created by the [`flat_map`] method on [`Iterator`]. See its
 /// documentation for more.
 ///
 /// [`flat_map`]: trait.Iterator.html#method.flat_map
 /// [`Iterator`]: trait.Iterator.html
 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
 #[stable(feature = "rust1", since = "1.0.0")]
-type FlatMap<I, U, F> = Flatten<Map<I, F>, <U as IntoIterator>::IntoIter>;
+pub struct FlatMap<I, U: IntoIterator, F> {
+    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("inner", &self.inner).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;
+
+    #[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.
@@ -2422,16 +2495,102 @@ type FlatMap<I, U, F> = Flatten<Map<I, F>, <U as IntoIterator>::IntoIter>;
 /// [`Iterator`]: trait.Iterator.html
 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
 #[unstable(feature = "iterator_flatten", issue = "0")]
+pub struct Flatten<I: Iterator>
+where I::Item: IntoIterator {
+    inner: FlattenCompat<I, <I::Item as IntoIterator>::IntoIter>,
+}
+
+#[unstable(feature = "iterator_flatten", issue = "0")]
+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 = "0")]
+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 = "0")]
+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 = "0")]
+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)]
-pub struct Flatten<I, U> {
+struct FlattenCompat<I, U> {
     iter: I,
     frontiter: Option<U>,
     backiter: Option<U>,
 }
 
-#[unstable(feature = "iterator_flatten", issue = "0")]
+impl<I, U> Iterator for FlattenCompat<I, U>
-impl<I: Iterator, U: Iterator> Iterator for Flatten<I, U>
+    where I: Iterator, U: Iterator,
-    where I::Item: IntoIterator<IntoIter = U, Item = U::Item>
+          I::Item: IntoIterator<IntoIter = U, Item = U::Item>
 {
     type Item = U::Item;
 
@@ -2498,8 +2657,7 @@ impl<I: Iterator, U: Iterator> Iterator for Flatten<I, U>
     }
 }
 
-#[unstable(feature = "iterator_flatten", issue = "0")]
+impl<I, U> DoubleEndedIterator for FlattenCompat<I, U>
-impl<I, U> DoubleEndedIterator for Flatten<I, U>
     where I: DoubleEndedIterator, U: DoubleEndedIterator,
           I::Item: IntoIterator<IntoIter = U, Item = U::Item>
 {
@@ -2555,10 +2713,6 @@ impl<I, U> DoubleEndedIterator for Flatten<I, U>
     }
 }
 
-#[unstable(feature = "fused", issue = "35602")]
-impl<I: FusedIterator, U: Iterator> FusedIterator for Flatten<I, U>
-    where I::Item: IntoIterator<IntoIter = U, Item = U::Item> {}
-
 /// An iterator that yields `None` forever after the underlying iterator
 /// yields `None` once.
 ///

src/libcore/tests/iter.rs

@@ -836,8 +836,6 @@ fn test_iterator_scan() {
     assert_eq!(i, ys.len());
 }
 
-// Note: We test flatten() by testing flat_map().
-
 #[test]
 fn test_iterator_flat_map() {
     let xs = [0, 3, 6];
@@ -876,6 +874,44 @@ fn test_iterator_flat_map_fold() {
     assert_eq!(i, 0);
 }
 
+#[test]
+fn test_iterator_flatten() {
+    let xs = [0, 3, 6];
+    let ys = [0, 1, 2, 3, 4, 5, 6, 7, 8];
+    let it = xs.iter().map(|&x| (x..).step_by(1).take(3)).flatten();
+    let mut i = 0;
+    for x in it {
+        assert_eq!(x, ys[i]);
+        i += 1;
+    }
+    assert_eq!(i, ys.len());
+}
+
+/// Test `Flatten::fold` with items already picked off the front and back,
+/// to make sure all parts of the `Flatten` are folded correctly.
+#[test]
+fn test_iterator_flatten_fold() {
+    let xs = [0, 3, 6];
+    let ys = [1, 2, 3, 4, 5, 6, 7];
+    let mut it = xs.iter().map(|&x| x..x+3).flatten();
+    assert_eq!(it.next(), Some(0));
+    assert_eq!(it.next_back(), Some(8));
+    let i = it.fold(0, |i, x| {
+        assert_eq!(x, ys[i]);
+        i + 1
+    });
+    assert_eq!(i, ys.len());
+
+    let mut it = xs.iter().map(|&x| x..x+3).flatten();
+    assert_eq!(it.next(), Some(0));
+    assert_eq!(it.next_back(), Some(8));
+    let i = it.rfold(ys.len(), |i, x| {
+        assert_eq!(x, ys[i - 1]);
+        i - 1
+    });
+    assert_eq!(i, 0);
+}
+
 #[test]
 fn test_inspect() {
     let xs = [1, 2, 3, 4];
@@ -1289,6 +1325,23 @@ fn test_double_ended_flat_map() {
     assert_eq!(it.next_back(), None);
 }
 
+#[test]
+fn test_double_ended_flatten() {
+    let u = [0,1];
+    let v = [5,6,7,8];
+    let mut it = u.iter().map(|x| &v[*x..v.len()]).flatten();
+    assert_eq!(it.next_back().unwrap(), &8);
+    assert_eq!(it.next().unwrap(),      &5);
+    assert_eq!(it.next_back().unwrap(), &7);
+    assert_eq!(it.next_back().unwrap(), &6);
+    assert_eq!(it.next_back().unwrap(), &8);
+    assert_eq!(it.next().unwrap(),      &6);
+    assert_eq!(it.next_back().unwrap(), &7);
+    assert_eq!(it.next_back(), None);
+    assert_eq!(it.next(),      None);
+    assert_eq!(it.next_back(), None);
+}
+
 #[test]
 fn test_double_ended_range() {
     assert_eq!((11..14).rev().collect::<Vec<_>>(), [13, 12, 11]);
@@ -1980,3 +2033,54 @@ fn test_flat_map_try_folds() {
     assert_eq!(iter.try_rfold(0, i8::checked_add), None);
     assert_eq!(iter.next_back(), Some(35));
 }
+
+#[test]
+fn test_flatten_try_folds() {
+    let f = &|acc, x| i32::checked_add(acc*2/3, x);
+    let mr = &|x| (5*x)..(5*x + 5);
+    assert_eq!((0..10).map(mr).flatten().try_fold(7, f), (0..50).try_fold(7, f));
+    assert_eq!((0..10).map(mr).flatten().try_rfold(7, f), (0..50).try_rfold(7, f));
+    let mut iter = (0..10).map(mr).flatten();
+    iter.next(); iter.next_back(); // have front and back iters in progress
+    assert_eq!(iter.try_rfold(7, f), (1..49).try_rfold(7, f));
+
+    let mut iter = (0..10).map(|x| (4*x)..(4*x + 4)).flatten();
+    assert_eq!(iter.try_fold(0, i8::checked_add), None);
+    assert_eq!(iter.next(), Some(17));
+    assert_eq!(iter.try_rfold(0, i8::checked_add), None);
+    assert_eq!(iter.next_back(), Some(35));
+}
+
+#[test]
+fn test_functor_laws() {
+    // identity:
+    fn identity<T>(x: T) -> T { x }
+    assert_eq!((0..10).map(identity).sum::<usize>(), (0..10).sum());
+
+    // composition:
+    fn f(x: usize) -> usize { x + 3 }
+    fn g(x: usize) -> usize { x * 2 }
+    fn h(x: usize) -> usize { g(f(x)) }
+    assert_eq!((0..10).map(f).map(g).sum::<usize>(), (0..10).map(h).sum());
+}
+
+#[test]
+fn test_monad_laws_left_identity() {
+    fn f(x: usize) -> impl Iterator<Item = usize> {
+        (0..10).map(move |y| x * y)
+    }
+    assert_eq!(once(42).flat_map(f.clone()).sum::<usize>(), f(42).sum());
+}
+
+#[test]
+fn test_monad_laws_right_identity() {
+    assert_eq!((0..10).flat_map(|x| once(x)).sum::<usize>(), (0..10).sum());
+}
+
+#[test]
+fn test_monad_laws_associativity() {
+    fn f(x: usize) -> impl Iterator<Item = usize> { 0..x }
+    fn g(x: usize) -> impl Iterator<Item = usize> { (0..x).rev() }
+    assert_eq!((0..10).flat_map(f).flat_map(g).sum::<usize>(),
+                (0..10).flat_map(|x| f(x).flat_map(g)).sum::<usize>());
+}

src/libcore/tests/lib.rs

@@ -25,6 +25,8 @@
 #![feature(inclusive_range)]
 #![feature(inclusive_range_syntax)]
 #![feature(iterator_try_fold)]
+#![feature(iterator_flatten)]
+#![feature(conservative_impl_trait)]
 #![feature(iter_rfind)]
 #![feature(iter_rfold)]
 #![feature(iterator_repeat_with)]