Splitting the reduce and result operations

scala · Apr 6, 2021 · 83df9ee · 83df9ee
1 parent 4f05b5d
commit 83df9ee
Show file tree

Hide file tree

Showing 2 changed files with 120 additions and 93 deletions.
diff --git a/src/main/scala/scala/collection/next/NextIterableOnceOpsExtensions.scala b/src/main/scala/scala/collection/next/NextIterableOnceOpsExtensions.scala
@@ -13,12 +13,28 @@
 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
+  import NextIterableOnceOpsExtensions.{GroupMapGen, GroupMapGenGen}
+
+  def groupBy[K](key: A => K)(implicit valuesFactory: Factory[A, C]): immutable.Map[K, C] =
+    groupByGen(key).result
+
+  def groupByGen[K](key: A => K)(implicit valuesFactory: Factory[A, C]): GroupMapGen[A, K, A, C] =
+    groupByGenGen(key).collectValuesAs(valuesFactory)
+
+  def groupByGenGen[K](key: A => K): GroupMapGenGen[A, K, A] =
+    groupMapGenGen(key)(identity)
+
+  def groupMap[K, V](key: A => K)(f: A => V)(implicit valuesFactory: Factory[V, CC[V]]): immutable.Map[K, CC[V]] =
+    groupMapGen(key)(f).result
+
+  def groupMapGen[K, V](key: A => K)(f: A => V)(implicit valuesFactory: Factory[V, CC[V]]): GroupMapGen[A, K, V, CC[V]] =
+    groupMapGenGen(key)(f).collectValuesAs(valuesFactory)
+
+  def groupMapGenGen[K, V](key: A => K)(f: A => V): GroupMapGenGen[A, K, V] =
+    new GroupMapGenGen(col, key, f)
 
   /**
    * Partitions this IterableOnce into a map according to a discriminator function `key`. All the values that
@@ -32,70 +48,54 @@ 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] =
-    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)
+  def groupMapReduce[K, V](key: A => K)(f: A => V)(reduce: (V, V) => V): immutable.Map[K, V] =
+    groupMapGenGen(key)(f).reduceValues(reduce)
 }
 
-object NextIterableOnceOpsExtensions {
-  final case class GroupMap[A, K, V, CC[_], MC[_, _]](
+private[next] object NextIterableOnceOpsExtensions {
+  final class GroupMapGenGen[A, K, V] private[NextIterableOnceOpsExtensions](
     col: IterableOnceOps[A, AnyConstr, _],
     key: A => K,
-    f: A => V,
-    colFactory: Factory[V, CC[V]],
-    mapFactory: CustomMapFactory[MC, K]
+    f: A => V
   ) {
-    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)
+    def reduceValues(reduce: (V, V) => V): immutable.Map[K, V] =
+      reduceValuesAs(immutable.Map)(reduce)
 
-    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] = {
+    def reduceValuesAs[MC](resultFactory: Factory[(K, V), MC])(reduce: (V, V) => V): MC = {
       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)
+      resultFactory.fromSpecific(m)
     }
-  }
 
-  sealed trait CustomMapFactory[MC[_, _], K] {
-    def from[V](col: IterableOnce[(K, V)]): MC[K, V]
+    def collectValuesAs[C](valuesFactory: Factory[V, C]): GroupMapGen[A, K, V, C] =
+      new GroupMapGen(col, key, f, valuesFactory)
   }
 
-  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)
-      }
+  final class GroupMapGen[A, K, V, C] private[NextIterableOnceOpsExtensions](
+    col: IterableOnceOps[A, AnyConstr, _],
+    key: A => K,
+    f: A => V,
+    valuesFactory: Factory[V, C]
+  ) {
+    def result: immutable.Map[K, C] =
+      resultAs(immutable.Map)
 
-    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)
+    def resultAs[MC](resultFactory: Factory[(K, C), MC]): MC = {
+      val m = mutable.Map.empty[K, mutable.Builder[V, C]]
+      col.foreach { elem =>
+        val k = key(elem)
+        val v = f(elem)
+        m.get(k) match {
+          case Some(builder) => builder.addOne(v)
+          case None          => m.update(key = k, value = valuesFactory.newBuilder.addOne(v))
+        }
       }
+      resultFactory.fromSpecific(m.view.mapValues(_.result()))
+    }
   }
 }
diff --git a/src/test/scala/scala/collection/next/TestIterableOnceExtensions.scala b/src/test/scala/scala/collection/next/TestIterableOnceExtensions.scala
@@ -16,62 +16,65 @@ import org.junit.Assert._
 import org.junit.Test
 import scala.collection.IterableOnceOps
 import scala.collection.generic.IsIterableOnce
-import scala.collection.immutable.{SortedMap, SortedSet}
+import scala.collection.immutable.{ArraySeq, BitSet, SortedMap, SortedSet}
 
 final class TestIterableOnceExtensions {
   import TestIterableOnceExtensions._
 
   // groupMapReduce --------------------------------------------
   @Test
   def iteratorGroupMapReduce(): Unit = {
-    def occurrences[A](coll: IterableOnce[A]): Map[A, Int] =
-      coll.iterator.groupMapReduce(identity)(_ => 1)(_ + _)
+    def occurrences[A](data: IterableOnce[A]): Map[A, Int] =
+      data.iterator.groupMapReduce(identity)(_ => 1)(_ + _)
 
-    val xs = Seq('a', 'b', 'b', 'c', 'a', 'a', 'a', 'b')
+    val data = Seq('a', 'b', 'b', 'c', 'a', 'a', 'a', 'b')
     val expected = Map('a' -> 4, 'b' -> 3, 'c' -> 1)
-    assertEquals(expected, occurrences(xs))
+
+    assertEquals(expected, occurrences(data))
   }
 
   @Test
   def iterableOnceOpsGroupMapReduce(): Unit = {
-    def occurrences[A, CC[_], C](coll: IterableOnceOps[A, CC, C]): Map[A, Int] =
-      coll.groupMapReduce(identity)(_ => 1)(_ + _)
+    def occurrences[A, CC[_], C](data: IterableOnceOps[A, CC, C]): Map[A, Int] =
+      data.groupMapReduce(identity)(_ => 1)(_ + _)
 
-    val xs = Seq('a', 'b', 'b', 'c', 'a', 'a', 'a', 'b')
+    val data = Seq('a', 'b', 'b', 'c', 'a', 'a', 'a', 'b')
     val expected = Map('a' -> 4, 'b' -> 3, 'c' -> 1)
-    assertEquals(expected, occurrences(xs))
+
+    assertEquals(expected, occurrences(data))
   }
 
   @Test
   def anyLikeIterableOnceGroupMapReduce(): Unit = {
-    def occurrences[Repr](coll: Repr)(implicit it: IsIterableOnce[Repr]): Map[it.A, Int] =
-      it(coll).iterator.groupMapReduce(identity)(_ => 1)(_ + _)
+    def occurrences[Repr](data: Repr)(implicit it: IsIterableOnce[Repr]): Map[it.A, Int] =
+      it(data).iterator.groupMapReduce(identity)(_ => 1)(_ + _)
 
-    val xs = "abbcaaab"
+    val data = "abbcaaab"
     val expected = Map('a' -> 4, 'b' -> 3, 'c' -> 1)
-    assertEquals(expected, occurrences(xs))
+
+    assertEquals(expected, occurrences(data))
   }
 
   @Test
   def customIterableOnceOpsGroupMapReduce(): Unit = {
-    def occurrences(coll: LowerCaseString): Map[Char, Int] =
-      coll.groupMapReduce(identity)(_ => 1)(_ + _)
+    def occurrences(data: LowerCaseString): Map[Char, Int] =
+      data.groupMapReduce(identity)(_ => 1)(_ + _)
 
-    val xs = LowerCaseString("abBcAaAb")
+    val data = LowerCaseString("abBcAaAb")
     val expected = Map('a' -> 4, 'b' -> 3, 'c' -> 1)
-    assertEquals(expected, occurrences(xs))
+
+    assertEquals(expected, occurrences(data))
   }
   // -----------------------------------------------------------
 
-  // groupMapTo ------------------------------------------------
+  // GroupMapGenGen --------------------------------------------
   @Test
-  def anyCollectionGroupMapToFull(): Unit = {
+  def anyCollectionGroupMapGenResultAs(): Unit = {
     def getUniqueUsersByCountrySorted(data: List[Record]): List[(String, List[String])] =
       data
-        .groupMapTo(_.country)(_.user)
+        .groupMapGenGen(_.country)(_.user)
         .collectValuesAs(SortedSet)
-        .collectResultsAs(SortedMap)
-        .result
+        .resultAs(SortedMap)
         .view
         .mapValues(_.toList)
         .toList
@@ -95,15 +98,11 @@ final class TestIterableOnceExtensions {
   }
 
   @Test
-  def anyCollectionGroupByToFull(): Unit = {
-    def getUniqueWordsByFirstLetterSorted(data: List[String]): List[(Char, List[String])] =
+  def anyCollectionGroupMapGenGenReduce(): Unit = {
+    def getAllWordsByFirstLetterSorted(data: List[String]): List[(Char, String)] =
       data
-        .groupByTo(_.head)
-        .collectValuesAs(SortedSet)
-        .collectResultsAs(SortedMap)
-        .result
-        .view
-        .mapValues(_.toList)
+        .groupByGenGen(_.head)
+        .reduceValuesAs(SortedMap)(_ ++ " " ++ _)
         .toList
 
     val data = List(
@@ -116,23 +115,51 @@ final class TestIterableOnceExtensions {
       "Winter",
       "Banana"
     )
-
     val expected = List(
-      'A' -> List("Apple", "April", "Autumn"),
-      'B' -> List("Banana"),
-      'W' -> List("Wilson", "Winter")
+      'A' -> "Autumn April Apple Apple",
+      'B' -> "Banana Banana",
+      'W' -> "Wilson Winter"
     )
 
-    assertEquals(expected, getUniqueWordsByFirstLetterSorted(data))
+    assertEquals(expected, getAllWordsByFirstLetterSorted(data))
   }
 
   @Test
-  def anyCollectionGroupByToReduceFull(): Unit = {
-    def getAllWordsByFirstLetterSorted(data: List[String]): List[(Char, String)] =
+  def iterableOnceOpsGroupByGenSpecificFactory(): Unit = {
+    def bitsByEven(data: BitSet): Map[Boolean, BitSet] =
+      data.groupByGen(x => (x % 2) == 0).result
+
+    val data = BitSet(1, 2, 3, 4, 5)
+    val expected = Map(
+      true -> BitSet(2, 4),
+      false -> BitSet(1, 3, 5)
+    )
+
+    assertEquals(expected, bitsByEven(data))
+  }
+
+  @Test
+  def iterableOnceOpsGroupMapGenIterableFactory(): Unit = {
+    def bitsByEvenAsChars(data: BitSet): Map[Boolean, Set[Char]] =
+      data.groupMapGen(x => (x % 2) == 0)(_.toChar).result
+
+    val data = BitSet(100, 101, 102, 103, 104, 105)
+    val expected = Map(
+      true -> Set('d', 'f', 'h'),
+      false -> Set('e', 'g', 'i')
+    )
+
+    assertEquals(expected, bitsByEvenAsChars(data))
+  }
+
+  @Test
+  def iteratorGroupBy(): Unit = {
+    def getUniqueWordsByFirstLetter(data: IterableOnce[String]): List[(Char, Set[String])] =
       data
-        .groupByTo(_.head)
-        .collectResultsAs(SortedMap)
-        .reduce(_ ++ " " ++ _)
+        .iterator
+        .groupBy(_.head)
+        .view
+        .mapValues(_.toSet)
         .toList
 
     val data = List(
@@ -147,12 +174,12 @@ final class TestIterableOnceExtensions {
     )
 
     val expected = List(
-      'A' -> "Autumn April Apple Apple",
-      'B' -> "Banana Banana",
-      'W' -> "Wilson Winter"
+      'A' -> Set("Apple", "April", "Autumn"),
+      'B' -> Set("Banana"),
+      'W' -> Set("Wilson", "Winter")
     )
 
-    assertEquals(expected, getAllWordsByFirstLetterSorted(data))
+    assertEquals(expected, getUniqueWordsByFirstLetter(data))
   }
   // -----------------------------------------------------------
 }