From 2ac680e710d8b11ee5075dcf0ef307209a8a40c5 Mon Sep 17 00:00:00 2001
From: obarros <tchubarros@gmail.com>
Date: Sat, 9 Jun 2018 12:33:04 +0100
Subject: [PATCH] Add groupBy for NonEmptySet

---
 .../main/scala/cats/data/NonEmptySet.scala    | 15 ++++++
 core/src/main/scala/cats/syntax/all.scala     |  1 +
 core/src/main/scala/cats/syntax/package.scala |  1 +
 core/src/main/scala/cats/syntax/set.scala     | 52 +++++++++++++++++++
 .../scala/cats/tests/NonEmptySetSuite.scala   |  5 ++
 5 files changed, 74 insertions(+)
 create mode 100644 core/src/main/scala/cats/syntax/set.scala

diff --git a/core/src/main/scala/cats/data/NonEmptySet.scala b/core/src/main/scala/cats/data/NonEmptySet.scala
index ade9b6184ff..8ad0a4ed9a8 100644
--- a/core/src/main/scala/cats/data/NonEmptySet.scala
+++ b/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 {
diff --git a/core/src/main/scala/cats/syntax/all.scala b/core/src/main/scala/cats/syntax/all.scala
index a8ce68e849e..971a53b55a1 100644
--- a/core/src/main/scala/cats/syntax/all.scala
+++ b/core/src/main/scala/cats/syntax/all.scala
@@ -64,3 +64,4 @@ trait AllSyntaxBinCompat1
     with NestedSyntax
     with BinestedSyntax
     with ParallelFlatSyntax
+    with SetSyntax
diff --git a/core/src/main/scala/cats/syntax/package.scala b/core/src/main/scala/cats/syntax/package.scala
index bf95218ba10..2e925afae6c 100644
--- a/core/src/main/scala/cats/syntax/package.scala
+++ b/core/src/main/scala/cats/syntax/package.scala
@@ -52,4 +52,5 @@ package object syntax {
   object validated extends ValidatedSyntax
   object vector extends VectorSyntax
   object writer extends WriterSyntax
+  object set extends SetSyntax
 }
diff --git a/core/src/main/scala/cats/syntax/set.scala b/core/src/main/scala/cats/syntax/set.scala
new file mode 100644
index 00000000000..0fa5323537f
--- /dev/null
+++ b/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)
+  }
+}
diff --git a/tests/src/test/scala/cats/tests/NonEmptySetSuite.scala b/tests/src/test/scala/cats/tests/NonEmptySetSuite.scala
index 36250049e4c..30b5dee8957 100644
--- a/tests/src/test/scala/cats/tests/NonEmptySetSuite.scala
+++ b/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))
+    }
+  }
 }