Add more group-like methods to NonEmptyList & Chain

core/src/main/scala/cats/data/Chain.scala

@@ -356,22 +356,111 @@ sealed abstract class Chain[+A] {
 
   /**
    * Groups elements inside this `Chain` according to the `Order`
-   * of the keys produced by the given mapping function.
+   * of the keys produced by the given key function.
+   *
+   * {{{
+   * scala> import scala.collection.immutable.SortedMap
+   * scala> import cats.data.{Chain, NonEmptyChain}
+   * scala> import cats.implicits._
+   * scala> val chain = Chain(12, -2, 3, -5)
+   * scala> val expectedResult = SortedMap(false -> NonEmptyChain(-2, -5), true -> NonEmptyChain(12, 3))
+   * scala> val result = chain.groupBy(_ >= 0)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
    */
-  final def groupBy[B](f: A => B)(implicit B: Order[B]): SortedMap[B, NonEmptyChain[A]] = {
-    implicit val ordering: Ordering[B] = B.toOrdering
-    var m = SortedMap.empty[B, NonEmptyChain[A]]
-    val iter = iterator
+  final def groupBy[K](key: A => K)(implicit K: Order[K]): SortedMap[K, NonEmptyChain[A]] =
+    groupMap(key)(identity)
+
+  /**
+   * Groups elements inside this `Chain` according to the `Order`
+   * of the keys produced by the given key function.
+   * And each element in a group is transformed into a value of type B
+   * using the mapping function.
+   *
+   * {{{
+   * scala> import scala.collection.immutable.SortedMap
+   * scala> import cats.data.{Chain, NonEmptyChain}
+   * scala> import cats.implicits._
+   * scala> val chain = Chain(12, -2, 3, -5)
+   * scala> val expectedResult = SortedMap(false -> NonEmptyChain("-2", "-5"), true -> NonEmptyChain("12", "3"))
+   * scala> val result = chain.groupMap(_ >= 0)(_.toString)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
+   */
+  final def groupMap[K, B](key: A => K)(f: A => B)(implicit K: Order[K]): SortedMap[K, NonEmptyChain[B]] = {
+    implicit val ordering: Ordering[K] = K.toOrdering
+    var m = SortedMap.empty[K, NonEmptyChain[B]]
+
+    for (elem <- iterator) {
+      val k = key(elem)
+
+      m.get(k) match {
+        case Some(cat) => m = m.updated(key = k, value = cat :+ f(elem))
+        case None      => m += (k -> NonEmptyChain.one(f(elem)))
+      }
+    }
+
+    m
+  }
 
-    while (iter.hasNext) {
-      val elem = iter.next()
-      val k = f(elem)
+  /**
+   * Groups elements inside this `Chain` according to the `Order`
+   * of the keys produced by the given key function.
+   * Then each element in a group is transformed into a value of type B
+   * using the mapping function.
+   * And finally they are all reduced into a single value
+   * using their `Semigroup`
+   *
+   * {{{
+   * scala> import scala.collection.immutable.SortedMap
+   * scala> import cats.data.Chain
+   * scala> import cats.implicits._
+   * scala> val chain = Chain("Hello", "World", "Goodbye", "World")
+   * scala> val expectedResult = SortedMap("goodbye" -> 1, "hello" -> 1, "world" -> 2)
+   * scala> val result = chain.groupMapReduce(_.strip.toLowerCase)(_ => 1)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
+   */
+  final def groupMapReduce[K, B](key: A => K)(f: A => B)(implicit K: Order[K], S: Semigroup[B]): SortedMap[K, B] =
+    groupMapReduceWith(key)(f)(S.combine)
+
+  /**
+   * Groups elements inside this `Chain` according to the `Order`
+   * of the keys produced by the given key function.
+   * Then each element in a group is transformed into a value of type B
+   * using the mapping function.
+   * And finally they are all reduced into a single value
+   * using the provided combine function.
+   *
+   * {{{
+   * scala> import scala.collection.immutable.SortedMap
+   * scala> import cats.data.Chain
+   * scala> import cats.implicits._
+   * scala> val chain = Chain("Hello", "World", "Goodbye", "World")
+   * scala> val expectedResult = SortedMap("goodbye" -> 1, "hello" -> 1, "world" -> 2)
+   * scala> val result = chain.groupMapReduceWith(_.strip.toLowerCase)(_ => 1)(_ + _)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
+   */
+  final def groupMapReduceWith[K, B](key: A => K)(f: A => B)(combine: (B, B) => B)(implicit
+    K: Order[K]
+  ): SortedMap[K, B] = {
+    implicit val ordering: Ordering[K] = K.toOrdering
+    var m = SortedMap.empty[K, B]
+
+    for (elem <- iterator) {
+      val k = key(elem)
 
       m.get(k) match {
-        case None      => m += ((k, NonEmptyChain.one(elem))); ()
-        case Some(cat) => m = m.updated(k, cat :+ elem)
+        case Some(b) => m = m.updated(key = k, value = combine(b, f(elem)))
+        case None    => m += (k -> f(elem))
       }
     }
+
     m
   }
 

core/src/main/scala/cats/data/NonEmptyChain.scala

@@ -379,12 +379,163 @@ class NonEmptyChainOps[A](private val value: NonEmptyChain[A])
 
   /**
    * Groups elements inside this `NonEmptyChain` according to the `Order`
-   * of the keys produced by the given mapping function.
+   * of the keys produced by the given key function.
+   *
+   * {{{
+   * scala> import cats.data.{NonEmptyChain, NonEmptyMap}
+   * scala> import cats.implicits._
+   * scala> val nec = NonEmptyChain(12, -2, 3, -5)
+   * scala> val expectedResult = NonEmptyMap.of(false -> NonEmptyChain(-2, -5), true -> NonEmptyChain(12, 3))
+   * scala> val result = nec.groupBy(_ >= 0)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
+   */
+  final def groupBy[K](key: A => K)(implicit K: Order[K]): NonEmptyMap[K, NonEmptyChain[A]] =
+    toChain.groupBy(key).asInstanceOf[NonEmptyMap[K, NonEmptyChain[A]]]
+
+  /**
+   * Groups elements inside this `NonEmptyChain` according to the `Order`
+   * of the keys produced by the given key function.
+   *
+   * {{{
+   * scala> import cats.data.{NonEmptyChain, NonEmptyMap}
+   * scala> import cats.implicits._
+   * scala> val nec = NonEmptyChain(12, -2, 3, -5)
+   * scala> val expectedResult = NonEmptyMap.of(false -> NonEmptyChain(-2, -5), true -> NonEmptyChain(12, 3))
+   * scala> val result = nec.groupByNem(_ >= 0)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
+   */
+  final def groupByNem[K](key: A => K)(implicit K: Order[K]): NonEmptyMap[K, NonEmptyChain[A]] =
+    groupBy(key)
+
+  /**
+   * Groups elements inside this `NonEmptyChain` according to the `Order`
+   * of the keys produced by the given key function.
+   * And each element in a group is transformed into a value of type B
+   * using the mapping function.
+   *
+   * {{{
+   * scala> import cats.data.{NonEmptyChain, NonEmptyMap}
+   * scala> import cats.implicits._
+   * scala> val nec = NonEmptyChain(12, -2, 3, -5)
+   * scala> val expectedResult = NonEmptyMap.of(false -> NonEmptyChain("-2", "-5"), true -> NonEmptyChain("12", "3"))
+   * scala> val result = nec.groupMap(_ >= 0)(_.toString)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
    */
-  final def groupBy[B](f: A => B)(implicit B: Order[B]): NonEmptyMap[B, NonEmptyChain[A]] =
-    toChain.groupBy(f).asInstanceOf[NonEmptyMap[B, NonEmptyChain[A]]]
+  final def groupMap[K, B](key: A => K)(f: A => B)(implicit K: Order[K]): NonEmptyMap[K, NonEmptyChain[B]] =
+    toChain.groupMap(key)(f).asInstanceOf[NonEmptyMap[K, NonEmptyChain[B]]]
 
-  final def groupByNem[B](f: A => B)(implicit B: Order[B]): NonEmptyMap[B, NonEmptyChain[A]] = groupBy(f)
+  /**
+   * Groups elements inside this `NonEmptyChain` according to the `Order`
+   * of the keys produced by the given key function.
+   * And each element in a group is transformed into a value of type B
+   * using the mapping function.
+   *
+   * {{{
+   * scala> import cats.data.{NonEmptyChain, NonEmptyMap}
+   * scala> import cats.implicits._
+   * scala> val nec = NonEmptyChain(12, -2, 3, -5)
+   * scala> val expectedResult = NonEmptyMap.of(false -> NonEmptyChain("-2", "-5"), true -> NonEmptyChain("12", "3"))
+   * scala> val result = nec.groupMapNem(_ >= 0)(_.toString)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
+   */
+  final def groupMapNem[K, B](key: A => K)(f: A => B)(implicit K: Order[K]): NonEmptyMap[K, NonEmptyChain[B]] =
+    groupMap(key)(f)
+
+  /**
+   * Groups elements inside this `NonEmptyChain` according to the `Order`
+   * of the keys produced by the given key function.
+   * Then each element in a group is transformed into a value of type B
+   * using the mapping function.
+   * And finally they are all reduced into a single value
+   * using their `Semigroup`
+   *
+   * {{{
+   * scala> import cats.data.{NonEmptyChain, NonEmptyMap}
+   * scala> import cats.implicits._
+   * scala> val nec = NonEmptyChain("Hello", "World", "Goodbye", "World")
+   * scala> val expectedResult = NonEmptyMap.of("goodbye" -> 1, "hello" -> 1, "world" -> 2)
+   * scala> val result = nec.groupMapReduce(_.strip.toLowerCase)(_ => 1)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
+   */
+  final def groupMapReduce[K, B](key: A => K)(f: A => B)(implicit K: Order[K], B: Semigroup[B]): NonEmptyMap[K, B] =
+    toChain.groupMapReduce(key)(f).asInstanceOf[NonEmptyMap[K, B]]
+
+  /**
+   * Groups elements inside this `NonEmptyChain` according to the `Order`
+   * of the keys produced by the given key function.
+   * Then each element in a group is transformed into a value of type B
+   * using the mapping function.
+   * And finally they are all reduced into a single value
+   * using their `Semigroup`
+   *
+   * {{{
+   * scala> import cats.data.{NonEmptyChain, NonEmptyMap}
+   * scala> import cats.implicits._
+   * scala> val nec = NonEmptyChain("Hello", "World", "Goodbye", "World")
+   * scala> val expectedResult = NonEmptyMap.of("goodbye" -> 1, "hello" -> 1, "world" -> 2)
+   * scala> val result = nec.groupMapReduceNem(_.strip.toLowerCase)(_ => 1)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
+   */
+  final def groupMapReduceNem[K, B](key: A => K)(f: A => B)(implicit K: Order[K], B: Semigroup[B]): NonEmptyMap[K, B] =
+    groupMapReduce(key)(f)
+
+  /**
+   * Groups elements inside this `NonEmptyChain` according to the `Order`
+   * of the keys produced by the given key function.
+   * Then each element in a group is transformed into a value of type B
+   * using the mapping function.
+   * And finally they are all reduced into a single value
+   * using the provided combine function.
+   *
+   * {{{
+   * scala> import cats.data.{NonEmptyChain, NonEmptyMap}
+   * scala> import cats.implicits._
+   * scala> val nec = NonEmptyChain("Hello", "World", "Goodbye", "World")
+   * scala> val expectedResult = NonEmptyMap.of("goodbye" -> 1, "hello" -> 1, "world" -> 2)
+   * scala> val result = nec.groupMapReduceWith(_.strip.toLowerCase)(_ => 1)(_ + _)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
+   */
+  final def groupMapReduceWith[K, B](key: A => K)(f: A => B)(combine: (B, B) => B)(implicit
+    K: Order[K]
+  ): NonEmptyMap[K, B] =
+    toChain.groupMapReduceWith(key)(f)(combine).asInstanceOf[NonEmptyMap[K, B]]
+
+  /**
+   * Groups elements inside this `NonEmptyChain` according to the `Order`
+   * of the keys produced by the given key function.
+   * Then each element in a group is transformed into a value of type B
+   * using the mapping function.
+   * And finally they are all reduced into a single value
+   * using the provided combine function.
+   *
+   * {{{
+   * scala> import cats.data.{NonEmptyChain, NonEmptyMap}
+   * scala> import cats.implicits._
+   * scala> val nec = NonEmptyChain("Hello", "World", "Goodbye", "World")
+   * scala> val expectedResult = NonEmptyMap.of("goodbye" -> 1, "hello" -> 1, "world" -> 2)
+   * scala> val result = nec.groupMapReduceWithNem(_.strip.toLowerCase)(_ => 1)(_ + _)
+   * scala> result === expectedResult
+   * res0: Boolean = true
+   * }}}
+   */
+  final def groupMapReduceWithNem[K, B](key: A => K)(f: A => B)(combine: (B, B) => B)(implicit
+    K: Order[K]
+  ): NonEmptyMap[K, B] =
+    groupMapReduceWith(key)(f)(combine)
 
   final def iterator: Iterator[A] = toChain.iterator