From 96c7b83fb8f2be52a7f73d89b46895344adfac9b Mon Sep 17 00:00:00 2001
From: Yann Bolliger <yann.bolliger@sqooba.io>
Date: Wed, 26 Feb 2020 11:12:10 +0100
Subject: [PATCH] Add partitionMap to 2.11 and 2.12

---
 .../collection/compat/PackageShared.scala     | 19 +++++++++++++--
 .../scala/collection/CollectionTest.scala     | 23 +++++++++++++++----
 2 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/compat/src/main/scala-2.11_2.12/scala/collection/compat/PackageShared.scala b/compat/src/main/scala-2.11_2.12/scala/collection/compat/PackageShared.scala
index 62f6d728..dc60c030 100644
--- a/compat/src/main/scala-2.11_2.12/scala/collection/compat/PackageShared.scala
+++ b/compat/src/main/scala-2.11_2.12/scala/collection/compat/PackageShared.scala
@@ -233,7 +233,7 @@ class IteratorExtensionMethods[A](private val self: c.Iterator[A]) extends AnyVa
     self.sameElements(that.iterator)
   }
   def concat[B >: A](that: c.TraversableOnce[B]): c.TraversableOnce[B] = self ++ that
-  def tapEach[U](f: A => U): c.Iterator[A] = self.map(a => { f(a); a })
+  def tapEach[U](f: A => U): c.Iterator[A]                             = self.map(a => { f(a); a })
 }
 
 class TraversableOnceExtensionMethods[A](private val self: c.TraversableOnce[A]) extends AnyVal {
@@ -319,6 +319,21 @@ class TraversableLikeExtensionMethods[A, Repr](private val self: c.GenTraversabl
   def tapEach[U](f: A => U)(implicit bf: CanBuildFrom[Repr, A, Repr]): Repr =
     self.map(a => { f(a); a })
 
+  def partitionMap[A1, A2, That, Repr1, Repr2](f: A => Either[A1, A2])(
+      implicit bf1: CanBuildFrom[Repr, A1, Repr1],
+      bf2: CanBuildFrom[Repr, A2, Repr2]
+  ): (Repr1, Repr2) = {
+    val l = bf1()
+    val r = bf2()
+    self.foreach { x =>
+      f(x) match {
+        case Left(x1)  => l += x1
+        case Right(x2) => r += x2
+      }
+    }
+    (l.result(), r.result())
+  }
+
   def groupMap[K, B, That](key: A => K)(f: A => B)(
       implicit bf: CanBuildFrom[Repr, B, That]): Map[K, That] = {
     val map = m.Map.empty[K, m.Builder[B, That]]
@@ -338,7 +353,7 @@ class TraversableLikeExtensionMethods[A, Repr](private val self: c.GenTraversabl
       val k = key(elem)
       val v = map.get(k) match {
         case Some(b) => reduce(b, f(elem))
-        case None => f(elem)
+        case None    => f(elem)
       }
       map.put(k, v)
     }
diff --git a/compat/src/test/scala/test/scala/collection/CollectionTest.scala b/compat/src/test/scala/test/scala/collection/CollectionTest.scala
index 96c73da9..b8ea12ae 100644
--- a/compat/src/test/scala/test/scala/collection/CollectionTest.scala
+++ b/compat/src/test/scala/test/scala/collection/CollectionTest.scala
@@ -36,8 +36,8 @@ class CollectionTest {
     val bT: BitSet = b
     assertEquals(BitSet(1, 2, 3), b)
 
-    val c          = xs.to(PriorityQueue)
-    val cT: PriorityQueue[Int]  = c
+    val c                      = xs.to(PriorityQueue)
+    val cT: PriorityQueue[Int] = c
     assert(PriorityQueue(1, 2, 3) sameElements c)
 
     val ys = List(1 -> "a", 2 -> "b")
@@ -48,8 +48,8 @@ class CollectionTest {
     assertTrue(m.isInstanceOf[Map[_, _]])
 
     // Stream.to(Seq) doesn't evaluate the stream
-    val strm = 1 #:: {throw new Exception("not lazy")} #:: Stream.empty[Int]
-    val strmsq: Seq[Int] = strm.to(Seq)
+    val strm                   = 1 #:: { throw new Exception("not lazy") } #:: Stream.empty[Int]
+    val strmsq: Seq[Int]       = strm.to(Seq)
     var strmln: LinearSeq[Int] = strm.to(LinearSeq)
   }
 
@@ -101,10 +101,23 @@ class CollectionTest {
     assertEquals(Map(3 -> 3, 4 -> 1), res)
   }
 
+  @Test
+  def partitionMapTest(): Unit = {
+    val empty = Seq.empty[Int].partitionMap(Right(_))
+    assertEquals((Seq(), Seq()), empty)
+
+    val res = Seq("foo", "test", "bar", "baz")
+      .partitionMap {
+        case s if s.contains("a") => Left(s)
+        case s                    => Right(s.length)
+      }
+    assertEquals((Seq("bar", "baz"), Seq("foo".length, "test".length)), res)
+  }
+
   @Test
   def tapEach(): Unit = {
     var count = 0
-    val it = Iterator(1, 2, 3).tapEach(count += _)
+    val it    = Iterator(1, 2, 3).tapEach(count += _)
     assertEquals(0, count)
     it.foreach(_ => ())
     assertEquals(6, count)