Add distinct method on NonEmptyList and NonEmptyVector (#1243)

typelevel · Aug 12, 2016 · cb7e2df · cb7e2df
1 parent 3fb3ed3
commit cb7e2df
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 1 deletion.
diff --git a/core/src/main/scala/cats/data/NonEmptyList.scala b/core/src/main/scala/cats/data/NonEmptyList.scala
@@ -5,6 +5,7 @@ import cats.instances.list._
 import cats.syntax.order._
 
 import scala.annotation.tailrec
+import scala.collection.immutable.TreeSet
 import scala.collection.mutable.ListBuffer
 
 /**
@@ -106,6 +107,20 @@ final case class NonEmptyList[A](head: A, tail: List[A]) {
     toList.iterator.map(A.show).mkString("NonEmptyList(", ", ", ")")
 
   override def toString: String = s"NonEmpty$toList"
+
+  /**
+   * Remove duplicates. Duplicates are checked using `Order[_]` instance.
+   */
+  def distinct(implicit O: Order[A]): NonEmptyList[A] = {
+    implicit val ord = O.toOrdering
+
+    val buf = ListBuffer.empty[A]
+    tail.foldLeft(TreeSet(head)) { (elementsSoFar, a) =>
+      if (elementsSoFar(a)) elementsSoFar else { buf += a; elementsSoFar + a }
+    }
+
+    NonEmptyList(head, buf.toList)
+  }
 }
 
 object NonEmptyList extends NonEmptyListInstances {

diff --git a/core/src/main/scala/cats/data/NonEmptyVector.scala b/core/src/main/scala/cats/data/NonEmptyVector.scala
@@ -2,7 +2,7 @@ package cats
 package data
 
 import scala.annotation.tailrec
-import scala.collection.immutable.VectorBuilder
+import scala.collection.immutable.{TreeSet, VectorBuilder}
 import cats.instances.vector._
 
 /**
@@ -130,6 +130,20 @@ final class NonEmptyVector[A] private (val toVector: Vector[A]) extends AnyVal {
   def length: Int = toVector.length
 
   override def toString: String = s"NonEmpty${toVector.toString}"
+
+  /**
+   * Remove duplicates. Duplicates are checked using `Order[_]` instance.
+   */
+  def distinct(implicit O: Order[A]): NonEmptyVector[A] = {
+    implicit val ord = O.toOrdering
+
+    val buf = Vector.newBuilder[A]
+    tail.foldLeft(TreeSet(head)) { (elementsSoFar, a) =>
+      if (elementsSoFar(a)) elementsSoFar else { buf += a; elementsSoFar + a }
+    }
+
+    NonEmptyVector(head, buf.result())
+  }
 }
 
 private[data] sealed trait NonEmptyVectorInstances {

diff --git a/tests/src/test/scala/cats/tests/NonEmptyListTests.scala b/tests/src/test/scala/cats/tests/NonEmptyListTests.scala
@@ -178,6 +178,12 @@ class NonEmptyListTests extends CatsSuite {
       (i :: nel).toList should === (i :: nel.toList)
     }
   }
+
+  test("NonEmptyList#distinct is consistent with List#distinct") {
+    forAll { nel: NonEmptyList[Int] =>
+      nel.distinct.toList should === (nel.toList.distinct)
+    }
+  }
 }
 
 class ReducibleNonEmptyListCheck extends ReducibleCheck[NonEmptyList]("NonEmptyList") {

diff --git a/tests/src/test/scala/cats/tests/NonEmptyVectorTests.scala b/tests/src/test/scala/cats/tests/NonEmptyVectorTests.scala
@@ -271,6 +271,12 @@ class NonEmptyVectorTests extends CatsSuite {
   test("Cannot create a new NonEmptyVector[Int] from apply with a an empty Vector") {
     "val bad: NonEmptyVector[Int] = NonEmptyVector(Vector.empty[Int])" shouldNot compile
   }
+
+  test("NonEmptyVector#distinct is consistent with Vector#distinct") {
+    forAll { nonEmptyVector: NonEmptyVector[Int] =>
+      nonEmptyVector.distinct.toVector should === (nonEmptyVector.toVector.distinct)
+    }
+  }
 }
 
 class ReducibleNonEmptyVectorCheck extends ReducibleCheck[NonEmptyVector]("NonEmptyVector") {