POC: GroupMap fluent api

Author: BalmungSan

src/main/scala/scala/collection/next/NextIterableOnceOpsExtensions.scala

@@ -13,9 +13,13 @@
 package scala.collection
 package next
 
+import scala.language.implicitConversions
+
 private[next] final class NextIterableOnceOpsExtensions[A, CC[_], C](
   private val col: IterableOnceOps[A, CC, C]
 ) extends AnyVal {
+  import NextIterableOnceOpsExtensions.GroupMap
+
   /**
    * Partitions this IterableOnce into a map according to a discriminator function `key`. All the values that
    * have the same discriminator are then transformed by the `value` function and then reduced into a
@@ -28,14 +32,70 @@ private[next] final class NextIterableOnceOpsExtensions[A, CC[_], C](
    *
    * @note This will force the evaluation of the Iterator.
    */
-  def groupMapReduce[K, B](key: A => K)(f: A => B)(reduce: (B, B) => B): immutable.Map[K, B] = {
-    val m = mutable.Map.empty[K, B]
-    col.foreach { elem =>
-      m.updateWith(key = key(elem)) {
-        case Some(b) => Some(reduce(b, f(elem)))
-        case None    => Some(f(elem))
+  def groupMapReduce[K, B](key: A => K)(f: A => B)(reduce: (B, B) => B): immutable.Map[K, B] =
+    groupMapTo(key)(f).reduce(reduce)
+
+  def groupByTo[K](key: A => K): GroupMap[A, K, A, immutable.Iterable, immutable.Map] =
+    groupMapTo(key)(identity)
+
+  def groupMapTo[K, V](key: A => K)(f: A => V): GroupMap[A, K, V, immutable.Iterable, immutable.Map] =
+    new GroupMap(col, key, f, immutable.Iterable, immutable.Map)
+}
+
+object NextIterableOnceOpsExtensions {
+  final case class GroupMap[A, K, V, CC[_], MC[_, _]](
+    col: IterableOnceOps[A, Any, Any],
+    key: A => K,
+    f: A => V,
+    colFactory: Factory[V, CC[V]],
+    mapFactory: CustomMapFactory[MC, K]
+  ) {
+    def collectValuesAs[CC1[_]](factory: Factory[V, CC1[V]]): GroupMap[A, K, V, CC1, MC] =
+      this.copy(colFactory = factory)
+
+    def collectResultsAs[MC1[_, _]](factory: CustomMapFactory[MC1, K]): GroupMap[A, K, V, CC, MC1] =
+      this.copy(mapFactory = factory)
+
+    final def result: MC[K, CC[V]] = {
+      val m = mutable.Map.empty[K, mutable.Builder[V, CC[V]]]
+      col.foreach { elem =>
+        val k = key(elem)
+        val v = f(elem)
+        m.get(k) match {
+          case Some(builder) => builder.addOne(v)
+          case None          => m += (k -> colFactory.newBuilder.addOne(v))
+        }
+      }
+      mapFactory.from(m.view.mapValues(_.result()))
+    }
+
+    final def reduce(reduce: (V, V) => V): MC[K, V] = {
+      val m = mutable.Map.empty[K, V]
+      col.foreach { elem =>
+        m.updateWith(key = key(elem)) {
+          case Some(b) => Some(reduce(b, f(elem)))
+          case None    => Some(f(elem))
+        }
       }
+      mapFactory.from(m)
     }
-    m.to(immutable.Map)
+  }
+
+  sealed trait CustomMapFactory[MC[_, _], K] {
+    def from[V](col: IterableOnce[(K, V)]): MC[K, V]
+  }
+
+  object CustomMapFactory {
+    implicit def fromMapFactory[MC[_, _], K](mf: MapFactory[MC]): CustomMapFactory[MC, K] =
+      new CustomMapFactory[MC, K] {
+        override def from[V](col: IterableOnce[(K, V)]): MC[K, V] =
+          mf.from(col)
+      }
+
+    implicit def fromSortedMapFactory[MC[_, _], K : Ordering](smf: SortedMapFactory[MC]): CustomMapFactory[MC, K] =
+      new CustomMapFactory[MC, K] {
+        override def from[V](col: IterableOnce[(K, V)]): MC[K, V] =
+          smf.from(col)
+      }
   }
 }

src/test/scala/scala/collection/next/TestIterableOnceExtensions.scala

@@ -16,10 +16,12 @@ import org.junit.Assert._
 import org.junit.Test
 import scala.collection.IterableOnceOps
 import scala.collection.generic.IsIterableOnce
+import scala.collection.immutable.{SortedMap, SortedSet}
 
 final class TestIterableOnceExtensions {
-  import TestIterableOnceExtensions.LowerCaseString
+  import TestIterableOnceExtensions._
 
+  // groupMapReduce --------------------------------------------
   @Test
   def iteratorGroupMapReduce(): Unit = {
     def occurrences[A](coll: IterableOnce[A]): Map[A, Int] =
@@ -59,6 +61,100 @@ final class TestIterableOnceExtensions {
     val expected = Map('a' -> 4, 'b' -> 3, 'c' -> 1)
     assertEquals(expected, occurrences(xs))
   }
+  // -----------------------------------------------------------
+
+  // groupMapTo ------------------------------------------------
+  @Test
+  def anyCollectionGroupMapToFull(): Unit = {
+    def getUniqueUsersByCountrySorted(data: List[Record]): List[(String, List[String])] =
+      data
+        .groupMapTo(_.country)(_.user)
+        .collectValuesAs(SortedSet)
+        .collectResultsAs(SortedMap)
+        .result
+        .view
+        .mapValues(_.toList)
+        .toList
+
+    val data = List(
+      Record(user = "Luis", country = "Colombia"),
+      Record(user = "Seth", country = "USA"),
+      Record(user = "April", country =  "USA"),
+      Record(user = "Julien", country = "Suisse"),
+      Record(user = "Rob", country =  "USA"),
+      Record(user = "Seth", country = "USA")
+    )
+
+    val expected = List(
+      "Colombia" -> List("Luis"),
+      "Suisse" -> List("Julien"),
+      "USA" -> List("April", "Rob", "Seth")
+    )
+
+    assertEquals(expected, getUniqueUsersByCountrySorted(data))
+  }
+
+  @Test
+  def anyCollectionGroupByToFull(): Unit = {
+    def getUniqueWordsByFirstLetterSorted(data: List[String]): List[(Char, List[String])] =
+      data
+        .groupByTo(_.head)
+        .collectValuesAs(SortedSet)
+        .collectResultsAs(SortedMap)
+        .result
+        .view
+        .mapValues(_.toList)
+        .toList
+
+    val data = List(
+      "Autumn",
+      "Banana",
+      "April",
+      "Wilson",
+      "Apple",
+      "Apple",
+      "Winter",
+      "Banana"
+    )
+
+    val expected = List(
+      'A' -> List("Apple", "April", "Autumn"),
+      'B' -> List("Banana"),
+      'W' -> List("Wilson", "Winter")
+    )
+
+    assertEquals(expected, getUniqueWordsByFirstLetterSorted(data))
+  }
+
+  @Test
+  def anyCollectionGroupByToReduceFull(): Unit = {
+    def getAllWordsByFirstLetterSorted(data: List[String]): List[(Char, String)] =
+      data
+        .groupByTo(_.head)
+        .collectResultsAs(SortedMap)
+        .reduce(_ ++ " " ++ _)
+        .toList
+
+    val data = List(
+      "Autumn",
+      "Banana",
+      "April",
+      "Wilson",
+      "Apple",
+      "Apple",
+      "Winter",
+      "Banana"
+    )
+
+    val expected = List(
+      'A' -> "Autumn April Apple Apple",
+      'B' -> "Banana Banana",
+      'W' -> "Wilson Winter"
+    )
+
+    assertEquals(expected, getAllWordsByFirstLetterSorted(data))
+  }
+  // -----------------------------------------------------------
 }
 
 object TestIterableOnceExtensions {
@@ -81,4 +177,6 @@ object TestIterableOnceExtensions {
     override def span(p: Char => Boolean): (String, String) = ???
     override def tapEach[U](f: Char => U): String = ???
   }
+
+  final case class Record(user: String, country: String)
 }