diff --git a/core/src/main/scala/cats/data/NonEmptyList.scala b/core/src/main/scala/cats/data/NonEmptyList.scala index aae4265fdf..207fe365a9 100644 --- a/core/src/main/scala/cats/data/NonEmptyList.scala +++ b/core/src/main/scala/cats/data/NonEmptyList.scala @@ -343,18 +343,24 @@ final case class NonEmptyList[+A](head: A, tail: List[A]) extends NonEmptyCollec override def distinct[AA >: A](implicit O: Order[AA]): NonEmptyList[AA] = distinctBy(identity[AA]) override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptyList[A] = { - implicit val ord: Ordering[B] = O.toOrdering - - val buf = ListBuffer.empty[A] - tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) => - val b = f(a) - if (elementsSoFar(b)) elementsSoFar - else { - buf += a; elementsSoFar + b + if (tail.isEmpty) this + else { + implicit val ord: Ordering[B] = O.toOrdering + + 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, buf.toList) + NonEmptyList(head, bldr.result()) + } } /** diff --git a/core/src/main/scala/cats/data/NonEmptySeq.scala b/core/src/main/scala/cats/data/NonEmptySeq.scala index 457a14b26e..4597a70f84 100644 --- a/core/src/main/scala/cats/data/NonEmptySeq.scala +++ b/core/src/main/scala/cats/data/NonEmptySeq.scala @@ -239,18 +239,21 @@ final class NonEmptySeq[+A] private (val toSeq: Seq[A]) extends AnyVal with NonE override def distinct[AA >: A](implicit O: Order[AA]): NonEmptySeq[AA] = distinctBy(identity[AA]) override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptySeq[A] = { - implicit val ord: Ordering[B] = O.toOrdering - - val buf = Seq.newBuilder[A] - tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) => - val b = f(a) - if (elementsSoFar(b)) elementsSoFar - else { - buf += a; elementsSoFar + b + if (toSeq.lengthCompare(1) == 0) this + else { + 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 } - } - NonEmptySeq(head, buf.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 index 7dffe88ce3..b80b6b27ef 100644 --- a/core/src/main/scala/cats/data/NonEmptyVector.scala +++ b/core/src/main/scala/cats/data/NonEmptyVector.scala @@ -249,18 +249,21 @@ final class NonEmptyVector[+A] private (val toVector: Vector[A]) override def distinct[AA >: A](implicit O: Order[AA]): NonEmptyVector[AA] = distinctBy(identity[AA]) override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptyVector[A] = { - implicit val ord: Ordering[B] = O.toOrdering - - val buf = Vector.newBuilder[A] - tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) => - val b = f(a) - if (elementsSoFar(b)) elementsSoFar - else { - buf += a; elementsSoFar + b + if (toVector.lengthCompare(1) == 0) 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 } - } - NonEmptyVector(head, buf.result()) + NonEmptyVector.fromVectorUnsafe(bldr.result()) + } } /**