add groupBy to NonEmptyList and groupByNel to List syntax (#1580)

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

@@ -7,6 +7,7 @@ import cats.syntax.order._
 import scala.annotation.tailrec
 import scala.collection.immutable.TreeSet
 import scala.collection.mutable.ListBuffer
+import scala.collection.{immutable, mutable}
 
 /**
  * A data type which represents a non empty list of A, with
@@ -217,6 +218,29 @@ final case class NonEmptyList[+A](head: A, tail: List[A]) {
     NonEmptyList((head, 0), bldr.result)
   }
 
+  /**
+    * Groups elements inside of this `NonEmptyList` using a mapping function
+    *
+    * {{{
+    * scala> import cats.data.NonEmptyList
+    * scala> val nel = NonEmptyList.of(12, -2, 3, -5)
+    * scala> nel.groupBy(_ >= 0)
+    * res0: Map[Boolean, cats.data.NonEmptyList[Int]] = Map(false -> NonEmptyList(-2, -5), true -> NonEmptyList(12, 3))
+    * }}}
+    */
+  def groupBy[B](f: A => B): Map[B, NonEmptyList[A]] = {
+    val m = mutable.Map.empty[B, mutable.Builder[A, List[A]]]
+    for { elem <- toList } {
+      m.getOrElseUpdate(f(elem), List.newBuilder[A]) += elem
+    }
+    val b = immutable.Map.newBuilder[B, NonEmptyList[A]]
+    for { (k, v) <- m } {
+      val head :: tail = v.result // we only create non empty list inside of the map `m`
+      b += ((k, NonEmptyList(head, tail)))
+    }
+    b.result
+  }
+
 }
 
 object NonEmptyList extends NonEmptyListInstances {

core/src/main/scala/cats/syntax/list.scala

@@ -9,4 +9,6 @@ trait ListSyntax {
 
 final class ListOps[A](val la: List[A]) extends AnyVal {
   def toNel: Option[NonEmptyList[A]] = NonEmptyList.fromList(la)
+  def groupByNel[B](f: A => B): Map[B, NonEmptyList[A]] =
+    toNel.fold(Map.empty[B, NonEmptyList[A]])(_.groupBy(f))
 }

tests/src/test/scala/cats/tests/ListTests.scala

@@ -29,6 +29,12 @@ class ListTests extends CatsSuite {
     List.empty[Int].toNel should === (None)
   }
 
+  test("groupByNel should be consistent with groupBy")(
+    forAll { (fa: List[Int], f: Int => Int) =>
+      fa.groupByNel(f).mapValues(_.toList) should === (fa.groupBy(f))
+    }
+  )
+
   test("show"){
     List(1, 2, 3).show should === ("List(1, 2, 3)")
     (Nil: List[Int]).show should === ("List()")

tests/src/test/scala/cats/tests/NonEmptyListTests.scala

@@ -231,6 +231,12 @@ class NonEmptyListTests extends CatsSuite {
       nel.zipWithIndex.toList should === (nel.toList.zipWithIndex)
     }
   }
+
+  test("NonEmptyList#groupBy is consistent with List#groupBy") {
+    forAll { (nel: NonEmptyList[Int], f: Int => Int) =>
+      nel.groupBy(f).mapValues(_.toList) should === (nel.toList.groupBy(f))
+    }
+  }
 }
 
 class ReducibleNonEmptyListCheck extends ReducibleCheck[NonEmptyList]("NonEmptyList") {