From 9a7e26156259560ac546a35dd285abe44728b1f5 Mon Sep 17 00:00:00 2001
From: Erick Tryzelaar <erick.tryzelaar@gmail.com>
Date: Mon, 7 Jan 2013 08:49:41 -0800
Subject: [PATCH] core: add vec.partition and vec.partitioned

---
 src/libcore/vec.rs | 81 ++++++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 79 insertions(+), 2 deletions(-)

diff --git a/src/libcore/vec.rs b/src/libcore/vec.rs
index e84ec40f7a9..ab0b5049f8e 100644
--- a/src/libcore/vec.rs
+++ b/src/libcore/vec.rs
@@ -405,6 +405,46 @@ pub fn rsplitn<T: Copy>(v: &[T], n: uint, f: fn(t: &T) -> bool) -> ~[~[T]] {
     result
 }
 
+/**
+ * Partitions a vector into two new vectors: those that satisfies the
+ * predicate, and those that do not.
+ */
+pub fn partition<T>(v: ~[T], f: fn(&T) -> bool) -> (~[T], ~[T]) {
+    let mut lefts  = ~[];
+    let mut rights = ~[];
+
+    do v.consume |_, elt| {
+        if f(&elt) {
+            lefts.push(elt);
+        } else {
+            rights.push(elt);
+        }
+    }
+
+    (lefts, rights)
+}
+
+/**
+ * Partitions a vector into two new vectors: those that satisfies the
+ * predicate, and those that do not.
+ */
+pub pure fn partitioned<T: Copy>(v: &[T], f: fn(&T) -> bool) -> (~[T], ~[T]) {
+    let mut lefts  = ~[];
+    let mut rights = ~[];
+
+    for each(v) |elt| {
+        unsafe {
+            if f(elt) {
+                lefts.push(*elt);
+            } else {
+                rights.push(*elt);
+            }
+        }
+    }
+
+    (lefts, rights)
+}
+
 // Mutators
 
 /// Removes the first element from a vector and return it
@@ -1664,7 +1704,6 @@ impl<T> &[T]: ImmutableVector<T> {
     pure fn filter_map<U: Copy>(&self, f: fn(t: &T) -> Option<U>) -> ~[U] {
         filter_map(*self, f)
     }
-
 }
 
 pub trait ImmutableEqVector<T: Eq> {
@@ -1714,8 +1753,8 @@ impl<T: Eq> &[T]: ImmutableEqVector<T> {
 
 pub trait ImmutableCopyableVector<T> {
     pure fn filter(&self, f: fn(t: &T) -> bool) -> ~[T];
-
     pure fn rfind(&self, f: fn(t: &T) -> bool) -> Option<T>;
+    pure fn partitioned(&self, f: fn(&T) -> bool) -> (~[T], ~[T]);
 }
 
 /// Extension methods for vectors
@@ -1743,6 +1782,15 @@ impl<T: Copy> &[T]: ImmutableCopyableVector<T> {
     pure fn rfind(&self, f: fn(t: &T) -> bool) -> Option<T> {
         rfind(*self, f)
     }
+
+    /**
+     * Partitions the vector into those that satisfies the predicate, and
+     * those that do not.
+     */
+    #[inline]
+    pure fn partitioned(&self, f: fn(&T) -> bool) -> (~[T], ~[T]) {
+        partitioned(*self, f)
+    }
 }
 
 pub trait OwnedVector<T> {
@@ -1757,6 +1805,7 @@ pub trait OwnedVector<T> {
     fn truncate(&mut self, newlen: uint);
     fn retain(&mut self, f: pure fn(t: &T) -> bool);
     fn consume(self, f: fn(uint, v: T));
+    fn partition(self, f: pure fn(&T) -> bool) -> (~[T], ~[T]);
 }
 
 impl<T> ~[T]: OwnedVector<T> {
@@ -1814,6 +1863,15 @@ impl<T> ~[T]: OwnedVector<T> {
     fn consume(self, f: fn(uint, v: T)) {
         consume(self, f)
     }
+
+    /**
+     * Partitions the vector into those that satisfies the predicate, and
+     * those that do not.
+     */
+    #[inline]
+    fn partition(self, f: fn(&T) -> bool) -> (~[T], ~[T]) {
+        partition(self, f)
+    }
 }
 
 pub trait OwnedCopyableVector<T: Copy> {
@@ -3105,6 +3163,25 @@ mod tests {
                        ~[~[1, 2, 3, 4], ~[5]];
     }
 
+    #[test]
+    fn test_partition() {
+        assert (~[]).partition(|x: &int| *x < 3) == (~[], ~[]);
+        assert (~[1, 2, 3]).partition(|x: &int| *x < 4) == (~[1, 2, 3], ~[]);
+        assert (~[1, 2, 3]).partition(|x: &int| *x < 2) == (~[1], ~[2, 3]);
+        assert (~[1, 2, 3]).partition(|x: &int| *x < 0) == (~[], ~[1, 2, 3]);
+    }
+
+    #[test]
+    fn test_partitioned() {
+        assert (~[]).partitioned(|x: &int| *x < 3) == (~[], ~[]);
+        assert (~[1, 2, 3]).partitioned(|x: &int| *x < 4) ==
+               (~[1, 2, 3], ~[]);
+        assert (~[1, 2, 3]).partitioned(|x: &int| *x < 2) ==
+               (~[1], ~[2, 3]);
+        assert (~[1, 2, 3]).partitioned(|x: &int| *x < 0) ==
+               (~[], ~[1, 2, 3]);
+    }
+
     #[test]
     #[should_fail]
     #[ignore(cfg(windows))]