How to implement a Set with a user-defined equality

Question

This question is related to this one: is it possible to create a Set-like class (meaning that it\'s extending the Set trait) in Scala where the equality used to

Answer 1

We created this solution as a group in a Scala training with Miles Sabin (@milessabin).

import scala.collection.mutable.ListBuffer
import scala.collection.generic.CanBuildFrom
import scala.collection.SetLike
import scala.collection.mutable.Builder

/**
 * we extend Set[T] to provide the Set-like interface
 * we extends SetLike[T, EqualitySet[T]] to specify that Set methods will return
 *   instances of type EqualitySet (and not simply Set)
 */
trait EqualitySet[T] extends Set[T] with SetLike[T, EqualitySet[T]] { outer =>
  /** we need to provide an Equals[T] instance to create an EqualitySet[T] */
  implicit def equality: Equals[T]  

  /** our internal implementation as a list of elements */
  protected val set = ListBuffer[T]()

  /** we need to implements those 4 methods */
  def contains(t: T) = set.exists(equality.isEqual(_, t))
  def +(t: T) = { if (!contains(t)) set += t; this }
  def -(t: T) = { set -= t; this }
  def iterator = set.iterator

  /** we must be able to provide an empty set with the proper equality definition */
  override def empty = new EqualitySet[T] { 
    override def equality = outer.equality
  }
}  

/** 
 * Companion object for the EqualitySet class
 */
object EqualitySet {

  /** 
   * this implicit is absolutely necessary to be able to preserve the resulting
   * collection type when calling `filter`
   */
  implicit def canBuildFrom[T] = new CanBuildFrom[EqualitySet[T], T, EqualitySet[T]] {
    def apply(from: EqualitySet[T]): Builder[T, EqualitySet[T]] = 
      new Builder[T, EqualitySet[T]] {
        // use a ListBuffer internally to accumulate elements
        private val elems = ListBuffer[T]()
        def +=(t: T) = { 
          if (!elems.exists(from.equality.isEqual(_, t))) elems += t
          this
        } 
        def clear() = elems.clear

        // when we finish building the collection
        // we can return an EqualitySet with the original equality relation
        def result() = new EqualitySet[T] {
          override val set = elems
          override def equality = from.equality
        }
     }
     def apply(): Builder[T, EqualitySet[T]] = 
       sys.error("this can't be implemented, because no equality instance is provided")
  }

  /** @return an EqualitySet for a type T having an Equals instance */
  def apply[T : Equals](ts: T*) = {
    var set  = new EqualitySet[T] {
      def equality = implicitly[Equals[T]]
    }.empty
    ts.foreach { t => set += t }
    set
  }
}

Then, when we use the code above we get:

scala> val set = EqualitySet[Coordinate](Coordinate(-1, 2), 
                                         Coordinate(-1, 3), 
                                         Coordinate(1, 4))

set: java.lang.Object with test.EqualitySet[Coordinate] = 
     Set(Coordinate(-1,2)
         Coordinate(1,4))

scala> val set2 = set.filter(_.i > 0)

// still an EqualitySet[Coordinate] \o/ */
set2: test.EqualitySet[Coordinate] = Set(Coordinate(1,4))