Add distinctBy to NonEmptyCollection and all impls

Author: cybersaurus

core/src/main/scala-2.13+/cats/data/NonEmptyLazyList.scala

@@ -345,14 +345,17 @@ class NonEmptyLazyListOps[A](private val value: NonEmptyLazyList[A])
   /**
    * Remove duplicates. Duplicates are checked using `Order[_]` instance.
    */
-  def distinct[AA >: A](implicit O: Order[AA]): NonEmptyLazyList[AA] = {
-    implicit val ord: Ordering[AA] = O.toOrdering
+  override def distinct[AA >: A](implicit O: Order[AA]): NonEmptyLazyList[AA] = distinctBy(identity[AA])
+
+  override def distinctBy[AA >: A, B](f: A => B)(implicit O: Order[B]): NonEmptyLazyList[AA] = {
+    implicit val ord: Ordering[B] = O.toOrdering
 
     val buf = LazyList.newBuilder[AA]
-    toLazyList.foldLeft(TreeSet.empty[AA]) { (elementsSoFar, a) =>
-      if (elementsSoFar(a)) elementsSoFar
+    toLazyList.foldLeft(TreeSet.empty[B]) { (elementsSoFar, a) =>
+      val b = f(a)
+      if (elementsSoFar(b)) elementsSoFar
       else {
-        buf += a; elementsSoFar + a
+        buf += a; elementsSoFar + b
       }
     }
 

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

@@ -603,7 +603,10 @@ class NonEmptyChainOps[A](private val value: NonEmptyChain[A])
    * Remove duplicates. Duplicates are checked using `Order[_]` instance.
    */
   final def distinct[AA >: A](implicit O: Order[AA]): NonEmptyChain[AA] =
-    create(toChain.distinct[AA])
+    distinctBy(identity[AA])
+
+  final def distinctBy[AA >: A, B](f: A => B)(implicit O: Order[B]): NonEmptyChain[AA] =
+    create(toChain.distinctBy[B](f))
 
   final def sortBy[B](f: A => B)(implicit B: Order[B]): NonEmptyChain[A] = create(toChain.sortBy(f))
   final def sorted[AA >: A](implicit AA: Order[AA]): NonEmptyChain[AA] = create(toChain.sorted[AA])

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

@@ -52,6 +52,7 @@ private[cats] trait NonEmptyCollection[+A, U[+_], NE[+_]] extends Any {
   def zipWithIndex: NE[(A, Int)]
 
   def distinct[AA >: A](implicit O: Order[AA]): NE[AA]
+  def distinctBy[AA >: A, B](f: A => B)(implicit O: Order[B]): NE[AA]
   def sortBy[B](f: A => B)(implicit B: Order[B]): NE[A]
   def sorted[AA >: A](implicit AA: Order[AA]): NE[AA]
   def groupByNem[B](f: A => B)(implicit B: Order[B]): NonEmptyMap[B, NE[A]]

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

@@ -340,14 +340,17 @@ final case class NonEmptyList[+A](head: A, tail: List[A]) extends NonEmptyCollec
   /**
    * Remove duplicates. Duplicates are checked using `Order[_]` instance.
    */
-  def distinct[AA >: A](implicit O: Order[AA]): NonEmptyList[AA] = {
-    implicit val ord: Ordering[AA] = O.toOrdering
+  override def distinct[AA >: A](implicit O: Order[AA]): NonEmptyList[AA] = distinctBy(identity[AA])
+
+  override def distinctBy[AA >: A, B](f: A => B)(implicit O: Order[B]): NonEmptyList[AA] = {
+    implicit val ord: Ordering[B] = O.toOrdering
 
     val buf = ListBuffer.empty[AA]
-    tail.foldLeft(TreeSet(head: AA)) { (elementsSoFar, b) =>
+    tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) =>
+      val b = f(a)
       if (elementsSoFar(b)) elementsSoFar
       else {
-        buf += b; elementsSoFar + b
+        buf += a; elementsSoFar + b
       }
     }
 

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

@@ -236,14 +236,17 @@ final class NonEmptySeq[+A] private (val toSeq: Seq[A]) extends AnyVal with NonE
   /**
    * Remove duplicates. Duplicates are checked using `Order[_]` instance.
    */
-  def distinct[AA >: A](implicit O: Order[AA]): NonEmptySeq[AA] = {
-    implicit val ord: Ordering[AA] = O.toOrdering
+  override def distinct[AA >: A](implicit O: Order[AA]): NonEmptySeq[AA] = distinctBy(identity[AA])
+
+  override def distinctBy[AA >: A, B](f: A => B)(implicit O: Order[B]): NonEmptySeq[AA] = {
+    implicit val ord: Ordering[B] = O.toOrdering
 
     val buf = Seq.newBuilder[AA]
-    tail.foldLeft(TreeSet(head: AA)) { (elementsSoFar, a) =>
-      if (elementsSoFar(a)) elementsSoFar
+    tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) =>
+      val b = f(a)
+      if (elementsSoFar(b)) elementsSoFar
       else {
-        buf += a; elementsSoFar + a
+        buf += a; elementsSoFar + b
       }
     }
 

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

@@ -246,14 +246,17 @@ final class NonEmptyVector[+A] private (val toVector: Vector[A])
   /**
    * Remove duplicates. Duplicates are checked using `Order[_]` instance.
    */
-  def distinct[AA >: A](implicit O: Order[AA]): NonEmptyVector[AA] = {
-    implicit val ord: Ordering[AA] = O.toOrdering
+  override def distinct[AA >: A](implicit O: Order[AA]): NonEmptyVector[AA] = distinctBy(identity[AA])
+
+  override def distinctBy[AA >: A, B](f: A => B)(implicit O: Order[B]): NonEmptyVector[AA] = {
+    implicit val ord: Ordering[B] = O.toOrdering
 
     val buf = Vector.newBuilder[AA]
-    tail.foldLeft(TreeSet(head: AA)) { (elementsSoFar, a) =>
-      if (elementsSoFar(a)) elementsSoFar
+    tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) =>
+      val b = f(a)
+      if (elementsSoFar(b)) elementsSoFar
       else {
-        buf += a; elementsSoFar + a
+        buf += a; elementsSoFar + b
       }
     }
 

tests/shared/src/test/scala-2.13+/cats/tests/NonEmptyLazyListSuite.scala

@@ -49,6 +49,7 @@ class NonEmptyLazyListSuite extends NonEmptyCollectionSuite[LazyList, NonEmptyLa
   checkAll(s"NonEmptyLazyList[Int]", HashTests[NonEmptyLazyList[Int]].hash)
   checkAll(s"Hash[NonEmptyLazyList[Int]]", SerializableTests.serializable(Hash[NonEmptyLazyList[Int]]))
 
+  checkAll("NonEmptyLazyList[Int]", NonEmptyAlternativeTests[NonEmptyLazyList].nonEmptyAlternative[Int, Int, Int])
   checkAll("NonEmptyLazyList[Int]", NonEmptyAlternativeTests[NonEmptyLazyList].nonEmptyAlternative[Int, Int, Int])
   checkAll("NonEmptyAlternative[NonEmptyLazyList]",
            SerializableTests.serializable(NonEmptyAlternative[NonEmptyLazyList])
@@ -175,6 +176,12 @@ class NonEmptyLazyListSuite extends NonEmptyCollectionSuite[LazyList, NonEmptyLa
     }
   }
 
+  test("NonEmptyLazyList#distinctBy is consistent with List#distinctBy") {
+    forAll { (ci: NonEmptyLazyList[Int], f: Int => String) =>
+      assert(ci.distinctBy(f).toList === (ci.toList.distinctBy(f)))
+    }
+  }
+
   test("NonEmptyLazyList#distinct is consistent with List#distinct") {
     forAll { (ci: NonEmptyLazyList[Int]) =>
       assert(ci.distinct.toList === (ci.toList.distinct))

tests/shared/src/test/scala/cats/tests/NonEmptyChainSuite.scala

@@ -203,6 +203,12 @@ class NonEmptyChainSuite extends NonEmptyCollectionSuite[Chain, NonEmptyChain, N
     }
   }
 
+  test("NonEmptyChain#distinctBy is consistent with List#distinctBy") {
+    forAll { (ci: NonEmptyChain[Int], f: Int => String) =>
+      assert(ci.distinctBy(f).toList === (ci.toList.distinctBy(f)))
+    }
+  }
+
   test("NonEmptyChain#distinct is consistent with List#distinct") {
     forAll { (ci: NonEmptyChain[Int]) =>
       assert(ci.distinct.toList === (ci.toList.distinct))

tests/shared/src/test/scala/cats/tests/NonEmptyListSuite.scala

@@ -296,6 +296,12 @@ class NonEmptyListSuite extends NonEmptyCollectionSuite[List, NonEmptyList, NonE
     }
   }
 
+  test("NonEmptyList#distinctBy is consistent with List#distinctBy") {
+    forAll { (nel: NonEmptyList[Int], f: Int => String) =>
+      assert(nel.distinctBy(f).toList === (nel.toList.distinctBy(f)))
+    }
+  }
+
   test("NonEmptyList#reverse is consistent with List#reverse") {
     forAll { (nel: NonEmptyList[Int]) =>
       assert(nel.reverse.toList === (nel.toList.reverse))

tests/shared/src/test/scala/cats/tests/NonEmptySeqSuite.scala

@@ -90,4 +90,16 @@ class NonEmptySeqSuite extends NonEmptyCollectionSuite[Seq, NonEmptySeq, NonEmpt
       assert(a.zip(b).toSeq === (a.toSeq.zip(b.toSeq)))
     }
   }
+
+  test("NonEmptySeq#distinct is consistent with List#distinct") {
+    forAll { (nes: NonEmptySeq[Int]) =>
+      assert(nes.distinct.toList === (nes.toList.distinct))
+    }
+  }
+
+  test("NonEmptySeq#distinctBy is consistent with List#distinctBy") {
+    forAll { (nes: NonEmptySeq[Int], f: Int => String) =>
+      assert(nes.distinctBy(f).toList === (nes.toList.distinctBy(f)))
+    }
+  }
 }

tests/shared/src/test/scala/cats/tests/NonEmptyVectorSuite.scala

@@ -359,6 +359,12 @@ class NonEmptyVectorSuite extends NonEmptyCollectionSuite[Vector, NonEmptyVector
     assert(compileErrors("val bad: NonEmptyVector[Int] = NonEmptyVector(Vector.empty[Int])").nonEmpty)
   }
 
+  test("NonEmptyVector#distinctBy is consistent with Vector#distinctBy") {
+    forAll { (nonEmptyVector: NonEmptyVector[Int], f: Int => String) =>
+      assert(nonEmptyVector.distinctBy(f).toVector === (nonEmptyVector.toVector.distinctBy(f)))
+    }
+  }
+
   test("NonEmptyVector#distinct is consistent with Vector#distinct") {
     forAll { (nonEmptyVector: NonEmptyVector[Int]) =>
       assert(nonEmptyVector.distinct.toVector === (nonEmptyVector.toVector.distinct))