Add groupBy for NonEmptySet (#2285)

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

@@ -346,6 +346,21 @@ sealed class NonEmptySetOps[A](val value: NonEmptySet[A]) {
     */
   def zipWithIndex: NonEmptySet[(A, Int)] =
     NonEmptySetImpl.create(toSortedSet.zipWithIndex)
+
+  /**
+    * Groups elements inside this `NonEmptySet` according to the `Order`
+    * of the keys produced by the given mapping function.
+    */
+  def groupBy[B](f: A => B)(implicit B: Order[B]): NonEmptyMap[B, NonEmptySet[A]] = {
+     reduceLeftTo(a => NonEmptyMap.one(f(a), NonEmptySet.one(a))) { (acc, a) =>
+       val key = f(a)
+       val result = acc.lookup(key) match {
+         case Some(nes) => nes.add(a)
+         case _ => NonEmptySet.one(a)
+       }
+       acc.add((key, result))
+     }
+  }
 }
 
 private[data] sealed abstract class NonEmptySetInstances {

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

@@ -64,4 +64,5 @@ trait AllSyntaxBinCompat1
     with NestedSyntax
     with BinestedSyntax
     with ParallelFlatSyntax
+    with SetSyntax
     with ValidatedExtensionSyntax

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

@@ -52,4 +52,5 @@ package object syntax {
   object validated extends ValidatedSyntax with ValidatedExtensionSyntax
   object vector extends VectorSyntax
   object writer extends WriterSyntax
+  object set extends SetSyntax
 }

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

@@ -0,0 +1,52 @@
+package cats.syntax
+
+import scala.collection.immutable.{SortedSet, SortedMap}
+import cats.data.NonEmptySet
+import cats.Order
+
+trait SetSyntax {
+  implicit final def catsSyntaxSet[A](se: SortedSet[A]): SetOps[A] = new SetOps(se)
+}
+
+final class SetOps[A](val se: SortedSet[A]) extends AnyVal {
+
+  /**
+    * Returns an Option of NonEmptySet from a SortedSet
+    *
+    * Example:
+    * {{{
+    * scala> import scala.collection.immutable.SortedSet
+    * scala> import cats.data.NonEmptySet
+    * scala> import cats.implicits._
+    *
+    * scala> val result1: SortedSet[Int] = SortedSet(1, 2)
+    * scala> result1.toNes
+    * res0: Option[NonEmptySet[Int]] = Some(TreeSet(1, 2))
+    *
+    * scala> val result2: SortedSet[Int] = SortedSet.empty[Int]
+    * scala> result2.toNes
+    * res1: Option[NonEmptySet[Int]] = None
+    * }}}
+    */
+  def toNes: Option[NonEmptySet[A]] = NonEmptySet.fromSet(se)
+
+  /**
+    * Groups elements inside this `SortedSet` according to the `Order` of the keys
+    * produced by the given mapping function.
+    *
+    * {{{
+    * scala> import cats.data.NonEmptySet
+    * scala> import scala.collection.immutable.{SortedMap, SortedSet}
+    * scala> import cats.implicits._
+    *
+    * scala> val sortedSet = SortedSet(12, -2, 3, -5)
+    *
+    * scala> sortedSet.groupByNes(_ >= 0)
+    * res0: SortedMap[Boolean, NonEmptySet[Int]] = Map(false -> TreeSet(-5, -2), true -> TreeSet(3, 12))
+    * }}}
+    */
+  def groupByNes[B](f: A => B)(implicit B: Order[B]): SortedMap[B, NonEmptySet[A]] = {
+    implicit val ordering = B.toOrdering
+    toNes.fold(SortedMap.empty[B, NonEmptySet[A]])(_.groupBy(f).toSortedMap)
+  }
+}

tests/src/test/scala/cats/tests/NonEmptySetSuite.scala

@@ -233,4 +233,9 @@ class NonEmptySetSuite extends CatsSuite {
     }
   }
 
+  test("NonEmptySet#groupBy is consistent with Set#groupBy") {
+    forAll { (nes: NonEmptySet[Int], f: Int => Int) =>
+      nes.groupBy(f).map(_.toSortedSet).toSortedMap should === (nes.toSortedSet.groupBy(f))
+    }
+  }
 }

tests/src/test/scala/cats/tests/SyntaxSuite.scala

@@ -1,8 +1,10 @@
 package cats
 package tests
 
+import scala.collection.immutable.SortedSet
+import scala.collection.immutable.SortedMap
 import cats.arrow.Compose
-import cats.data.{Binested, Nested}
+import cats.data.{Binested, Nested, NonEmptyList, NonEmptySet}
 import cats.instances.AllInstances
 import cats.syntax.{AllSyntax, AllSyntaxBinCompat}
 
@@ -358,5 +360,22 @@ object SyntaxSuite extends AllSyntaxBinCompat with AllInstances with AllSyntax {
 
     val binested: Binested[F, G, H, A, B] = fgahb.binested
   }
+
+  def testNonEmptySet[A, B: Order] : Unit = {
+    val f = mock[A => B]
+    val set = mock[SortedSet[A]]
+
+    val nes: Option[NonEmptySet[A]] = set.toNes
+    val grouped: SortedMap[B, NonEmptySet[A]] = set.groupByNes(f)
+  }
+
+  def testNonEmptyList[A, B: Order] : Unit = {
+    val f = mock[A => B]
+    val list = mock[List[A]]
+
+    val nel: Option[NonEmptyList[A]] = list.toNel
+    val grouped: SortedMap[B, NonEmptyList[A]] = list.groupByNel(f)
+  }
+
 }