Merge pull request scala#6760 from valydia/scala/collection-strawman#235

`partitionWith` support in `IterableOps`

Author: szeiger

src/library/scala/collection/Iterable.scala

@@ -648,6 +648,31 @@ trait IterableOps[+A, +CC[_], +C] extends Any with IterableOnce[A] with Iterable
   def collect[B](pf: PartialFunction[A, B]): CC[B] =
     fromIterable(new View.Collect(this, pf))
 
+  /** A pair of ${coll}s, first, consisting of the values that are produced by `f` applied to this $coll elements contained in [[scala.util.Left]] and, second,
+    *  all the values contained in [[scala.util.Right]].
+    *
+    *  Example:
+    *  {{{
+    *    val xs = $Coll(1, "one", 2, "two", 3, "three") partitionWith {
+    *     case i: Int => Left(i)
+    *     case s: String => Right(s)
+    *    }
+    *    // xs == ($Coll(1, 2, 3),
+    *    //        $Coll(one, two, three))
+    *  }}}
+    *
+    *  @tparam A1         element type of the first resulting collection
+    *  @tparam A2         element type of the second resulting collection
+    *  @param f          split function that map the element of the $coll into an [[scala.util.Either]][A1, A2]
+    *
+    *  @return           a pair of ${coll}s, first, consisting of the values that are produced by `f` applied to this $coll
+    *               elements contained in [[scala.util.Left]] and, second, all the values contained in [[scala.util.Right]].
+    */
+  def partitionWith[A1, A2](f: A => Either[A1, A2]): (CC[A1], CC[A2]) = {
+    val mp = new View.PartitionWith(this, f)
+    (fromIterable(mp.left), fromIterable(mp.right))
+  }
+
   /** Returns a new $coll containing the elements from the left hand operand followed by the elements from the
     *  right hand operand. The element type of the $coll is the most specific superclass encompassing
     *  the element types of the two operands.

src/library/scala/collection/StrictOptimizedIterableOps.scala

@@ -224,4 +224,17 @@ trait StrictOptimizedIterableOps[+A, +CC[_], +C]
     b.result()
   }
 
+  // Optimized, push-based version of `partitionWith`
+  override def partitionWith[A1, A2](f: A => Either[A1, A2]): (CC[A1], CC[A2]) = {
+    val l = iterableFactory.newBuilder[A1]
+    val r = iterableFactory.newBuilder[A2]
+    foreach { x =>
+      f(x) match {
+        case Left(x1) => l += x1
+        case Right(x2) => r += x2
+      }
+    }
+    (l.result(), r.result())
+  }
+
 }

src/library/scala/collection/View.scala

@@ -160,7 +160,7 @@ object View extends IterableFactory[View] {
     override def isEmpty: Boolean = underlying.isEmpty
   }
 
-  /** A view that partitions an underlying collection into two views */
+  /** A class that partitions an underlying collection into two views */
   @SerialVersionUID(3L)
   class Partition[A](val underlying: SomeIterableOps[A], val p: A => Boolean) extends Serializable {
 
@@ -183,6 +183,71 @@ object View extends IterableFactory[View] {
     override def isEmpty: Boolean = iterator.isEmpty
   }
 
+  /** A class that splits an underlying collection into two views */
+  @SerialVersionUID(3L)
+  class PartitionWith[A, A1, A2](val underlying: SomeIterableOps[A], val f: A => Either[A1, A2]) extends Serializable {
+
+    /** The view consisting of all elements of the underlying collection
+      *  that map to `Left`.
+      */
+    val left: View[A1] = new LeftPartitionedWith(this, f)
+
+
+    /** The view consisting of all elements of the underlying collection
+      *  that map to `Right`.
+      */
+    val right: View[A2] = new RightPartitionedWith(this, f)
+
+  }
+
+  @SerialVersionUID(3L)
+  class LeftPartitionedWith[A, A1, A2](partitionWith: PartitionWith[A, A1, A2], f: A => Either[A1, A2]) extends AbstractView[A1] {
+    def iterator = new AbstractIterator[A1] {
+      private val self = partitionWith.underlying.iterator
+      private var hd: A1 = _
+      private var hdDefined: Boolean = false
+      def hasNext = hdDefined || {
+        def findNext(): Boolean =
+          if (self.hasNext) {
+            f(self.next()) match {
+              case Left(a1) => hd = a1; hdDefined = true; true
+              case Right(_) => findNext()
+            }
+          } else false
+        findNext()
+      }
+      def next() =
+        if (hasNext) {
+          hdDefined = false
+          hd
+        } else Iterator.empty.next()
+    }
+  }
+
+  @SerialVersionUID(3L)
+  class RightPartitionedWith[A, A1, A2](partitionWith: PartitionWith[A, A1, A2], f: A => Either[A1, A2]) extends AbstractView[A2] {
+      def iterator = new AbstractIterator[A2] {
+        private val self = partitionWith.underlying.iterator
+        private var hd: A2 = _
+        private var hdDefined: Boolean = false
+        def hasNext = hdDefined || {
+          def findNext(): Boolean =
+            if (self.hasNext) {
+              f(self.next()) match {
+                case Left(_) => findNext()
+                case Right(a2) => hd = a2; hdDefined = true; true
+              }
+            } else false
+          findNext()
+        }
+        def next() =
+          if (hasNext) {
+            hdDefined = false
+            hd
+          } else Iterator.empty.next()
+      }
+  }
+
   /** A view that drops leading elements of the underlying collection. */
   @SerialVersionUID(3L)
   class Drop[A](underlying: SomeIterableOps[A], n: Int) extends AbstractView[A] {

test/junit/scala/collection/BuildFromTest.scala

@@ -105,7 +105,7 @@ class BuildFromTest {
     builder.result()
   }
 
-  def mapSplit[A, B, C, ToL, ToR](coll: Iterable[A])(f: A => Either[B, C])
+  def partitionWith[A, B, C, ToL, ToR](coll: Iterable[A])(f: A => Either[B, C])
               (implicit bfLeft:  BuildFrom[coll.type, B, ToL], bfRight: BuildFrom[coll.type, C, ToR]): (ToL, ToR) = {
     val left = bfLeft.newBuilder(coll)
     val right = bfRight.newBuilder(coll)
@@ -135,14 +135,14 @@ class BuildFromTest {
   }
 
   @Test
-  def mapSplitTest: Unit = {
+  def partitionWithTest: Unit = {
     val xs1 = immutable.List(1, 2, 3)
-    val (xs2, xs3) = mapSplit(xs1)(x => if (x % 2 == 0) Left(x) else Right(x.toString))
+    val (xs2, xs3) = partitionWith(xs1)(x => if (x % 2 == 0) Left(x) else Right(x.toString))
     val xs4: immutable.List[Int] = xs2
     val xs5: immutable.List[String] = xs3
 
     val xs6 = immutable.TreeMap((1, "1"), (2, "2"))
-    val (xs7, xs8) = mapSplit(xs6) { case (k, v) => Left[(String, Int), (Int, Boolean)]((v, k)) }
+    val (xs7, xs8) = partitionWith(xs6) { case (k, v) => Left[(String, Int), (Int, Boolean)]((v, k)) }
     val xs9: immutable.TreeMap[String, Int] = xs7
     val xs10: immutable.TreeMap[Int, Boolean] = xs8
   }

test/junit/scala/collection/IterableTest.scala

@@ -212,6 +212,7 @@ class IterableTest {
     Assert.assertEquals(baselist.reverse, checklist)
   }
 
+  @Test
   def unzip(): Unit = {
     val zipped = Seq((1, 'a'), (2, 'b'), (3, 'c'))
     val (s1, s2) = zipped.unzip
@@ -261,4 +262,14 @@ class IterableTest {
     Assert.assertEquals("Fu()", foo.toString)
   }
 
+  @Test
+  def partitionWith: Unit = {
+    val (left, right) = Seq(1, "1", 2, "2", 3, "3", 4, "4", 5, "5").partitionWith {
+      case i: Int => Left(i)
+      case s: String => Right(s)
+    }
+    Assert.assertEquals(left, Seq(1, 2, 3, 4 ,5))
+    Assert.assertEquals(right, Seq("1", "2", "3", "4" ,"5"))
+  }
+
 }