further optimize Chain#distinctBy

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

@@ -123,13 +123,19 @@ sealed abstract class Chain[+A] {
   /**
    * Returns true if there are no elements in this collection.
    */
-  def isEmpty: Boolean = !this.isInstanceOf[Chain.NonEmpty[_]]
+  def isEmpty: Boolean = this eq Chain.Empty
 
   /**
    * Returns false if there are no elements in this collection.
    */
   final def nonEmpty: Boolean = !isEmpty
 
+  // Temporary solution for length comparison in some special cases.
+  // TODO: Consider implementing more generic `lengthCompare` method.
+  private def isEmptyOrSingleton: Boolean =
+    // Chain.Append and Chain.Wrap always contain more than one element.
+    isEmpty || this.isInstanceOf[Chain.Singleton[_]]
+
   /**
    * Concatenates this with `c` in O(1) runtime.
    */
@@ -642,17 +648,19 @@ sealed abstract class Chain[+A] {
    * }}}
    */
   def distinctBy[B](f: A => B)(implicit O: Order[B]): Chain[A] = {
-    implicit val ord: Ordering[B] = O.toOrdering
-
-    val alreadyIn = mutable.TreeSet.empty[B]
-
-    foldLeft(Chain.empty[A]) { (elementsSoFar, a) =>
-      val b = f(a)
-      if (alreadyIn.add(b)) {
-        elementsSoFar :+ a
-      } else {
-        elementsSoFar
+    if (isEmptyOrSingleton) this
+    else {
+      implicit val ord: Ordering[B] = O.toOrdering
+
+      var result = Chain.empty[A]
+      val seen = mutable.TreeSet.empty[B]
+      val it = iterator
+      while (it.hasNext) {
+        val next = it.next()
+        if (seen.add(f(next)))
+          result ++= one(next)
       }
+      result
     }
   }