Improve optimizations

typelevel · Feb 1, 2025 · 44b076d · 44b076d
1 parent 572fbfa
commit 44b076d
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 24 deletions.
diff --git a/core/src/main/scala/cats/data/NonEmptyList.scala b/core/src/main/scala/cats/data/NonEmptyList.scala
@@ -345,15 +345,22 @@ final case class NonEmptyList[+A](head: A, tail: List[A]) extends NonEmptyCollec
   override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptyList[A] = {
     implicit val ord: Ordering[B] = O.toOrdering
 
-    val bldr = List.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
+    if (tail.isEmpty) this
+    else {
+      val bldr = List.newBuilder[A]
+      val seen = mutable.TreeSet.empty[B]
+      var rest = tail
+      seen.add(f(head))
+      while (rest.nonEmpty) {
+        val next = rest.head
+        rest = rest.tail
+        if (seen.add(f(next))) {
+          bldr += next
+        }
+      }
+
+      NonEmptyList(head, bldr.result())
     }
-    NonEmptyList.fromListUnsafe(bldr.result())
   }
 
   /**

diff --git a/core/src/main/scala/cats/data/NonEmptySeq.scala b/core/src/main/scala/cats/data/NonEmptySeq.scala
@@ -241,15 +241,19 @@ final class NonEmptySeq[+A] private (val toSeq: Seq[A]) extends AnyVal with NonE
   override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptySeq[A] = {
     implicit val ord: Ordering[B] = O.toOrdering
 
-    val bldr = Seq.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
+    if (toSeq.sizeIs == 1) this
+    else {
+      val bldr = Seq.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
+      }
+
+      NonEmptySeq.fromSeqUnsafe(bldr.result())
     }
-    NonEmptySeq.fromSeqUnsafe(bldr.result())
   }
 
   /**

diff --git a/core/src/main/scala/cats/data/NonEmptyVector.scala b/core/src/main/scala/cats/data/NonEmptyVector.scala
@@ -251,15 +251,19 @@ final class NonEmptyVector[+A] private (val toVector: Vector[A])
   override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptyVector[A] = {
     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
+    if (toVector.sizeIs == 1) this
+    else {
+      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
+      }
+
+      NonEmptyVector.fromVectorUnsafe(bldr.result())
     }
-    NonEmptyVector.fromVectorUnsafe(bldr.result())
   }
 
   /**