diff --git a/src/libcore/iter/iterator.rs b/src/libcore/iter/iterator.rs index 33adb3f49dd0d..e68e059385299 100644 --- a/src/libcore/iter/iterator.rs +++ b/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(self, f: F) -> FlatMap 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::>(); + /// 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::>(); + /// assert_eq!(d2, [&[1, 2], &[3, 4], &[5, 6], &[7, 8]]); + /// + /// let d1 = d3.iter().flatten().flatten().collect::>(); + /// 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 + where Self: Sized, Self::Item: IntoIterator { + Flatten { inner: flatten_compat(self) } } /// Creates an iterator which ends after the first [`None`]. diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index 1a2da83429af9..623cad754dd72 100644 --- a/src/libcore/iter/mod.rs +++ b/src/libcore/iter/mod.rs @@ -2410,12 +2410,15 @@ impl Iterator for Scan 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 { - iter: I, - f: F, - frontiter: Option, - backiter: Option, + inner: FlattenCompat, ::IntoIter> +} + +#[stable(feature = "rust1", since = "1.0.0")] +impl Clone for FlatMap + where ::IntoIter: Clone +{ + fn clone(&self) -> Self { FlatMap { inner: self.inner.clone() } } } #[stable(feature = "core_impl_debug", since = "1.9.0")] @@ -2423,11 +2426,7 @@ impl fmt::Debug for FlatMap 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 Iterator for FlatMap { type Item = U::Item; + #[inline] + fn next(&mut self) -> Option { self.inner.next() } + + #[inline] + fn size_hint(&self) -> (usize, Option) { self.inner.size_hint() } + + #[inline] + fn try_fold(&mut self, init: Acc, fold: Fold) -> R where + Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try + { + self.inner.try_fold(init, fold) + } + + #[inline] + fn 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 DoubleEndedIterator for FlatMap + where F: FnMut(I::Item) -> U, + U: IntoIterator, + U::IntoIter: DoubleEndedIterator +{ + #[inline] + fn next_back(&mut self) -> Option { self.inner.next_back() } + + #[inline] + fn try_rfold(&mut self, init: Acc, fold: Fold) -> R where + Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try + { + self.inner.try_rfold(init, fold) + } + + #[inline] + fn rfold(self, init: Acc, fold: Fold) -> Acc + where Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.inner.rfold(init, fold) + } +} + +#[unstable(feature = "fused", issue = "35602")] +impl FusedIterator for FlatMap + 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 +where I::Item: IntoIterator { + inner: FlattenCompat::IntoIter>, +} + +#[unstable(feature = "iterator_flatten", issue = "48213")] +impl fmt::Debug for Flatten + where I: Iterator + fmt::Debug, U: Iterator + fmt::Debug, + I::Item: IntoIterator, +{ + 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 Clone for Flatten + where I: Iterator + Clone, U: Iterator + Clone, + I::Item: IntoIterator, +{ + fn clone(&self) -> Self { Flatten { inner: self.inner.clone() } } +} + +#[unstable(feature = "iterator_flatten", issue = "48213")] +impl Iterator for Flatten + where I: Iterator, U: Iterator, + I::Item: IntoIterator +{ + type Item = U::Item; + + #[inline] + fn next(&mut self) -> Option { self.inner.next() } + + #[inline] + fn size_hint(&self) -> (usize, Option) { self.inner.size_hint() } + + #[inline] + fn try_fold(&mut self, init: Acc, fold: Fold) -> R where + Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try + { + self.inner.try_fold(init, fold) + } + + #[inline] + fn 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 DoubleEndedIterator for Flatten + where I: DoubleEndedIterator, U: DoubleEndedIterator, + I::Item: IntoIterator +{ + #[inline] + fn next_back(&mut self) -> Option { self.inner.next_back() } + + #[inline] + fn try_rfold(&mut self, init: Acc, fold: Fold) -> R where + Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try + { + self.inner.try_rfold(init, fold) + } + + #[inline] + fn rfold(self, init: Acc, fold: Fold) -> Acc + where Fold: FnMut(Acc, Self::Item) -> Acc, + { + self.inner.rfold(init, fold) + } +} + +#[unstable(feature = "fused", issue = "35602")] +impl FusedIterator for Flatten + where I: FusedIterator, U: Iterator, + I::Item: IntoIterator {} + +/// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`. +fn flatten_compat(iter: I) -> FlattenCompat { + FlattenCompat { iter, frontiter: None, backiter: None } +} + +/// Real logic of both `Flatten` and `FlatMap` which simply delegate to +/// this type. +#[derive(Clone, Debug)] +struct FlattenCompat { + iter: I, + frontiter: Option, + backiter: Option, +} + +impl Iterator for FlattenCompat + where I: Iterator, U: Iterator, + I::Item: IntoIterator +{ + type Item = U::Item; + #[inline] fn next(&mut self) -> Option { 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 Iterator for FlatMap 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 Iterator for FlatMap 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 DoubleEndedIterator for FlatMap where - F: FnMut(I::Item) -> U, - U: IntoIterator, - U::IntoIter: DoubleEndedIterator +impl DoubleEndedIterator for FlattenCompat + where I: DoubleEndedIterator, U: DoubleEndedIterator, + I::Item: IntoIterator { #[inline] fn next_back(&mut self) -> Option { 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 DoubleEndedIterator for FlatMap 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 DoubleEndedIterator for FlatMap 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 FusedIterator for FlatMap - where I: FusedIterator, U: IntoIterator, F: FnMut(I::Item) -> U {} - /// An iterator that yields `None` forever after the underlying iterator /// yields `None` once. /// diff --git a/src/libcore/lib.rs b/src/libcore/lib.rs index d2162d307e038..3dd30ee1c69e2 100644 --- a/src/libcore/lib.rs +++ b/src/libcore/lib.rs @@ -93,6 +93,7 @@ #![feature(doc_spotlight)] #![feature(rustc_const_unstable)] #![feature(iterator_repeat_with)] +#![feature(iterator_flatten)] #[prelude_import] #[allow(unused)] diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index f91c919d7447d..edd75f7795ed7 100644 --- a/src/libcore/tests/iter.rs +++ b/src/libcore/tests/iter.rs @@ -874,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]; @@ -1287,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::>(), [13, 12, 11]); @@ -1978,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(x: T) -> T { x } + assert_eq!((0..10).map(identity).sum::(), (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::(), (0..10).map(h).sum()); +} + +#[test] +fn test_monad_laws_left_identity() { + fn f(x: usize) -> impl Iterator { + (0..10).map(move |y| x * y) + } + assert_eq!(once(42).flat_map(f.clone()).sum::(), f(42).sum()); +} + +#[test] +fn test_monad_laws_right_identity() { + assert_eq!((0..10).flat_map(|x| once(x)).sum::(), (0..10).sum()); +} + +#[test] +fn test_monad_laws_associativity() { + fn f(x: usize) -> impl Iterator { 0..x } + fn g(x: usize) -> impl Iterator { (0..x).rev() } + assert_eq!((0..10).flat_map(f).flat_map(g).sum::(), + (0..10).flat_map(|x| f(x).flat_map(g)).sum::()); +} diff --git a/src/libcore/tests/lib.rs b/src/libcore/tests/lib.rs index 3e901a9d442ce..7954d52f6b1e3 100644 --- a/src/libcore/tests/lib.rs +++ b/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)]