From 8a5057af2e19bea374600669fee49af1ccf89d08 Mon Sep 17 00:00:00 2001 From: Aaron Turon Date: Tue, 16 Dec 2014 23:44:29 -0800 Subject: [PATCH] std::iter: Add partition and unzip methods to iterators --- src/libcore/iter.rs | 91 +++++++++++++++++++++++++++++++++++------- src/libcore/prelude.rs | 2 +- src/libstd/prelude.rs | 2 +- 3 files changed, 78 insertions(+), 17 deletions(-) diff --git a/src/libcore/iter.rs b/src/libcore/iter.rs index 9c3e53a1ace..028d2ce1cbb 100644 --- a/src/libcore/iter.rs +++ b/src/libcore/iter.rs @@ -59,6 +59,7 @@ pub use self::MinMaxResult::*; use clone::Clone; use cmp; use cmp::Ord; +use default::Default; use mem; use num::{ToPrimitive, Int}; use ops::{Add, Deref, FnMut}; @@ -68,20 +69,6 @@ use uint; #[deprecated = "renamed to Extend"] pub use self::Extend as Extendable; -/// Conversion from an `Iterator` -#[unstable = "may be replaced by a more general conversion trait"] -pub trait FromIterator { - /// Build a container with elements from an external iterator. - fn from_iter>(iterator: T) -> Self; -} - -/// A type growable from an `Iterator` implementation -#[unstable = "just renamed as part of collections reform"] -pub trait Extend { - /// Extend a container with the elements yielded by an arbitrary iterator - fn extend>(&mut self, iterator: T); -} - /// An interface for dealing with "external iterators". These types of iterators /// can be resumed at any time as all state is stored internally as opposed to /// being located on the call stack. @@ -106,6 +93,20 @@ pub trait Iterator { fn size_hint(&self) -> (uint, Option) { (0, None) } } +/// Conversion from an `Iterator` +#[unstable = "may be replaced by a more general conversion trait"] +pub trait FromIterator { + /// Build a container with elements from an external iterator. + fn from_iter>(iterator: T) -> Self; +} + +/// A type growable from an `Iterator` implementation +#[unstable = "just renamed as part of collections reform"] +pub trait Extend { + /// Extend a container with the elements yielded by an arbitrary iterator + fn extend>(&mut self, iterator: T); +} + #[unstable = "new convention for extension traits"] /// An extension trait providing numerous methods applicable to all iterators. pub trait IteratorExt: Iterator { @@ -223,7 +224,6 @@ pub trait IteratorExt: Iterator { Enumerate{iter: self, count: 0} } - /// Creates an iterator that has a `.peek()` method /// that returns an optional reference to the next element. /// @@ -471,6 +471,35 @@ pub trait IteratorExt: Iterator { FromIterator::from_iter(self) } + /// Loops through the entire iterator, collecting all of the elements into + /// one of two containers, depending on a predicate. The elements of the + /// first container satisfy the predicate, while the elements of the second + /// do not. + /// + /// ``` + /// let vec = vec![1i, 2i, 3i, 4i]; + /// let (even, odd): (Vec, Vec) = vec.into_iter().partition(|&n| n % 2 == 0); + /// assert_eq!(even, vec![2, 4]); + /// assert_eq!(odd, vec![1, 3]); + /// ``` + #[unstable = "recently added as part of collections reform"] + fn partition(mut self, mut f: F) -> (B, B) where + B: Default + Extend, F: FnMut(&A) -> bool + { + let mut left: B = Default::default(); + let mut right: B = Default::default(); + + for x in self { + if f(&x) { + left.extend(Some(x).into_iter()) + } else { + right.extend(Some(x).into_iter()) + } + } + + (left, right) + } + /// Loops through `n` iterations, returning the `n`th element of the /// iterator. /// @@ -661,6 +690,38 @@ pub trait IteratorExt: Iterator { #[unstable = "trait is unstable"] impl IteratorExt for I where I: Iterator {} +/// Extention trait for iterators of pairs. +pub trait IteratorPairExt: Iterator<(A, B)> { + /// Converts an iterator of pairs into a pair of containers. + /// + /// Loops through the entire iterator, collecting the first component of + /// each item into one new container, and the second component into another. + fn unzip(mut self) -> (FromA, FromB) where + FromA: FromIterator + Extend, FromB: FromIterator + Extend + { + struct SizeHint(uint, Option); + impl Iterator for SizeHint { + fn next(&mut self) -> Option { None } + fn size_hint(&self) -> (uint, Option) { + (self.0, self.1) + } + } + + let (lo, hi) = self.size_hint(); + let mut ts: FromA = FromIterator::from_iter(SizeHint(lo, hi)); + let mut us: FromB = FromIterator::from_iter(SizeHint(lo, hi)); + + for (t, u) in self { + ts.extend(Some(t).into_iter()); + us.extend(Some(u).into_iter()); + } + + (ts, us) + } +} + +impl IteratorPairExt for I where I: Iterator<(A, B)> {} + /// A range iterator able to yield elements from both ends /// /// A `DoubleEndedIterator` can be thought of as a deque in that `next()` and `next_back()` exhaust diff --git a/src/libcore/prelude.rs b/src/libcore/prelude.rs index fd1598db8cd..cecb4938153 100644 --- a/src/libcore/prelude.rs +++ b/src/libcore/prelude.rs @@ -52,7 +52,7 @@ pub use cmp::{Ordering, Equiv}; pub use cmp::Ordering::{Less, Equal, Greater}; pub use iter::{FromIterator, Extend, IteratorExt}; pub use iter::{Iterator, DoubleEndedIterator, DoubleEndedIteratorExt, RandomAccessIterator}; -pub use iter::{IteratorCloneExt, CloneIteratorExt}; +pub use iter::{IteratorCloneExt, CloneIteratorExt, IteratorPairExt}; pub use iter::{IteratorOrdExt, MutableDoubleEndedIterator, ExactSizeIterator}; pub use num::{ToPrimitive, FromPrimitive}; pub use option::Option; diff --git a/src/libstd/prelude.rs b/src/libstd/prelude.rs index fc59f06ae6c..11eb569e5b5 100644 --- a/src/libstd/prelude.rs +++ b/src/libstd/prelude.rs @@ -66,7 +66,7 @@ #[doc(no_inline)] pub use iter::{FromIterator, Extend, ExactSizeIterator}; #[doc(no_inline)] pub use iter::{Iterator, IteratorExt, DoubleEndedIterator}; #[doc(no_inline)] pub use iter::{DoubleEndedIteratorExt, CloneIteratorExt}; -#[doc(no_inline)] pub use iter::{RandomAccessIterator, IteratorCloneExt}; +#[doc(no_inline)] pub use iter::{RandomAccessIterator, IteratorCloneExt, IteratorPairExt}; #[doc(no_inline)] pub use iter::{IteratorOrdExt, MutableDoubleEndedIterator}; #[doc(no_inline)] pub use num::{ToPrimitive, FromPrimitive}; #[doc(no_inline)] pub use boxed::Box;