add Chain#distinctBy

typelevel · Apr 13, 2022 · cf16b6a · cf16b6a
1 parent f494931
commit cf16b6a
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 10 deletions.
diff --git a/core/src/main/scala/cats/data/Chain.scala b/core/src/main/scala/cats/data/Chain.scala
@@ -27,8 +27,8 @@ import cats.kernel.instances.StaticMethods
 
 import scala.annotation.tailrec
 import scala.collection.immutable.SortedMap
-import scala.collection.immutable.TreeSet
 import scala.collection.immutable.{IndexedSeq => ImIndexedSeq}
+import scala.collection.mutable
 import scala.collection.mutable.ListBuffer
 
 import Chain.{
@@ -151,6 +151,10 @@ sealed abstract class Chain[+A] extends ChainCompat[A] {
    */
   final def nonEmpty: Boolean = !isEmpty
 
+  // Quick check whether the chain is either empty or contains one element only.
+  @inline private def isEmptyOrSingleton: Boolean =
+    isEmpty || this.isInstanceOf[Chain.Singleton[_]]
+
   /**
    * Concatenates this with `c` in O(1) runtime.
    */
@@ -760,19 +764,59 @@ sealed abstract class Chain[+A] extends ChainCompat[A] {
 
   /**
    * Remove duplicates. Duplicates are checked using `Order[_]` instance.
+   * 
+   * Example:
+   * {{{
+   * scala> import cats.data.Chain
+   * scala> val chain = Chain(1, 2, 2, 3)
+   * scala> chain.distinct
+   * res0: cats.data.Chain[Int] = Chain(1, 2, 3)
+   * }}}
    */
   def distinct[AA >: A](implicit O: Order[AA]): Chain[AA] = {
-    implicit val ord: Ordering[AA] = O.toOrdering
-
-    var alreadyIn = TreeSet.empty[AA]
+    if (isEmptyOrSingleton) this
+    else {
+      implicit val ord: Ordering[AA] = O.toOrdering
+
+      val bldr = Vector.newBuilder[AA]
+      val seen = mutable.TreeSet.empty[AA]
+      val it = iterator
+      while (it.hasNext) {
+        val next = it.next()
+        if (seen.add(next))
+          bldr += next
+      }
+      // Result can contain a single element only.
+      Chain.fromSeq(bldr.result())
+    }
+  }
 
-    foldLeft(Chain.empty[AA]) { (elementsSoFar, b) =>
-      if (alreadyIn.contains(b)) {
-        elementsSoFar
-      } else {
-        alreadyIn += b
-        elementsSoFar :+ b
+  /**
+   * Remove duplicates by a predicate. Duplicates are checked using `Order[_]` instance.
+   * 
+   * Example:
+   * {{{
+   * scala> import cats.data.Chain
+   * scala> val chain = Chain(1, 2, 3, 4)
+   * scala> chain.distinctBy(_ / 2)
+   * res0: cats.data.Chain[Int] = Chain(1, 2, 4)
+   * }}}
+   */
+  def distinctBy[B](f: A => B)(implicit O: Order[B]): Chain[A] = {
+    if (isEmptyOrSingleton) this
+    else {
+      implicit val ord: Ordering[B] = O.toOrdering
+
+      val bldr = Vector.newBuilder[A]
+      val seen = mutable.TreeSet.empty[B]
+      val it = iterator
+      while (it.hasNext) {
+        val next = it.next()
+        if (seen.add(f(next)))
+          bldr += next
       }
+      // Result can contain a single element only.
+      Chain.fromSeq(bldr.result())
     }
   }
 

diff --git a/tests/shared/src/test/scala/cats/tests/ChainSuite.scala b/tests/shared/src/test/scala/cats/tests/ChainSuite.scala
@@ -380,6 +380,12 @@ class ChainSuite extends CatsSuite {
     }
   }
 
+  test("Chain#distinctBy is consistent with List#distinctBy") {
+    forAll { (a: Chain[Int], f: Int => String) =>
+      assertEquals(a.distinctBy(f).toList, a.toList.distinctBy(f))
+    }
+  }
+
   test("=== is consistent with == (issue #2540)") {
     assertEquals(Chain.one(1) |+| Chain.one(2) |+| Chain.one(3), Chain.fromSeq(List(1, 2, 3)))