Immutable, sorted multi map in Scala

SortedMultiMap.scala

import scala.collection._
import scala.collection.generic.CanBuildFrom
import scala.Iterator

/**
 * Factory singleton for creating a [[ .SortedMultiMap SortedMultiMap]].
 */
object SortedMultiMap {
  /**
   * Construct an empty [[ .SortedMultiMap SortedMultiMap]].
   * @param ord the [[scala.math.Ordering ordering]] type class
   * @tparam K the type of keys
   * @tparam V the type of values
   * @return an empty `SortedMultiMap`
   */
  def empty[K, V](implicit ord: Ordering[K]): SortedMultiMap[K, V] =
    new SortedMultiMap[K,V](SortedMap.empty[K, Set[V]])

  /**
   * Constructs a new [[ .SortedMultiMap SortedMultiMap]] containing
   * passed elements.
   * @param elements one or more `key -> value` pairs
   * @param ord the [[scala.math.Ordering ordering]] type class
   * @tparam K the type of keys
   * @tparam V the type of values
   * @return a new `SortedMultiMap` containing the elements
   */
  def apply[K,V](elements: (K, V)*)(implicit ord: Ordering[K]): SortedMultiMap[K, V] = {
    val builder = new SortedMultiMapBuilder[K, V]()
    builder ++= elements
    builder.result()
  }

  /**
   * Constructs a new [[ .SortedMultiMap SortedMultiMap]] containing
   * passed elements.
   *
   * @param elements one or more `key -> Set(value)` pairs
   * @param ord the [[scala.math.Ordering ordering]] type class
   * @param d a dummy implicit to distinguish this `apply` method from another `apply` method
   * @tparam K the type of keys
   * @tparam V the type of values
   * @return a new `SortedMultiMap` containing the elements
   */
  def apply[K,V](elements: (K, Set[V])*)(implicit ord: Ordering[K], d: DummyImplicit): SortedMultiMap[K, V] =
    new SortedMultiMap[K,V](SortedMap[K, Set[V]](elements: _*))

  /**
   * Constructs a new [[ .SortedMultiMap SortedMultiMap]] containing
   * the passed delegate sorted map.
   * @param delegate the [[scala.collection.SortedMap SortedMap]] to delegate most of the operations to
   * @param ord the [[scala.math.Ordering ordering]] type class
   * @tparam K the type of keys
   * @tparam V the type of values
   * @return a new `SortedMultiMap` constructed from the delegate map
   */
  def apply[K,V](delegate: SortedMap[K, Set[V]])(implicit ord: Ordering[K]): SortedMultiMap[K, V] =
    new SortedMultiMap[K, V](delegate)


  /**
   * Implicit method that provides the means of converting any collection type containing
   * `key -> Set(value)` elements to a [[ .SortedMultiMap SortedMultiMap]].
   * In particular this implicit method returns a [[scala.collection.generic.CanBuildFrom builder factory]]
   * that is able to provide the appropriate [[scala.collection.mutable.Builder Builder]] to convert
   * an arbitrary collection of `key -> Set(value)` elements to `SortedMultiMap`.
   * @param ord the [[scala.math.Ordering ordering]] type class
   * @tparam K the type of keys
   * @tparam V the type of values
   * @return the mechanism to convert any collection containing `key -> Set(value)` elements
   */
  implicit def canBuildFrom[K, V](implicit ord: Ordering[K]): CanBuildFrom[GenTraversableOnce[_], (K, Set[V]), SortedMultiMap[K,V]] =
    new CanBuildFrom[GenTraversableOnce[_], (K, Set[V]), SortedMultiMap[K, V]] {
      override def apply(): mutable.Builder[(K, Set[V]), SortedMultiMap[K, V]] = newBuilder[K,V]
      override def apply(from: GenTraversableOnce[_]): mutable.Builder[(K, Set[V]), SortedMultiMap[K, V]] = newBuilder[K,V]
    }

  /**
   * Generates a new [[scala.collection.mutable.Builder Builder]] to construct a
   * [[ .SortedMultiMap SortedMultiMap]] by
   * adding elements of type `K -> Set[V]`
   * @param ord the [[scala.math.Ordering ordering]] type class
   * @tparam K the type of keys
   * @tparam V the type of values
   * @return a builder which can construct a `SortedMultiMap` given
   */
  def newBuilder[K,V](implicit ord: Ordering[K]): mutable.Builder[(K, Set[V]), SortedMultiMap[K, V]] =
    new mutable.MapBuilder[K, Set[V], SortedMap[K, Set[V]]](SortedMap.empty[K, Set[V]](ord)) mapResult (new SortedMultiMap[K,V](_)(ord))


  /**
   * Implicit method that provides the means of converting any sequence type containing
   * `key -> value` elements to a [[ .SortedMultiMap SortedMultiMap]].
   * In particular this implicit method returns a [[scala.collection.generic.CanBuildFrom builder factory]]
   * that is able to provide the appropriate [[scala.collection.mutable.Builder Builder]] to convert
   * a sequence of `key -> value` elements to `SortedMultiMap`.
   * @param ord the [[scala.math.Ordering ordering]] type class
   * @tparam K the type of keys
   * @tparam V the type of values
   * @return the mechanism to convert any sequence containing `key -> value` elements
   */
  implicit def canBuildFromSeq[K, V](implicit ord: Ordering[K]): CanBuildFrom[GenSeq[_], (K, V), SortedMultiMap[K,V]] =
    new CanBuildFrom[GenSeq[_], (K, V), SortedMultiMap[K,V]] {
      override def apply(): mutable.Builder[(K, V), SortedMultiMap[K, V]] =
        new SortedMultiMapBuilder[K, V]()
      override def apply(from: GenSeq[_]): mutable.Builder[(K, V), SortedMultiMap[K, V]] =
        new SortedMultiMapBuilder[K, V]()
    }

  /**
   * Class representing a special [[scala.collection.mutable.Builder Builder]] which
   * can construct a new [[ .SortedMultiMap SortedMultiMap]]
   * by processing a sequence of key value pairs.
   * @param ord the [[scala.math.Ordering ordering]] type class
   * @tparam K the type of keys
   * @tparam V the type of values
   */
  class SortedMultiMapBuilder[K, V](implicit ord: Ordering[K]) extends mutable.Builder[(K,V), SortedMultiMap[K,V]] {
    private val elems = new mutable.HashMap[K, Set[V]]

    override def +=(elem: (K, V)): SortedMultiMapBuilder.this.type = {
      elem match {
        case (key, value) =>
          elems.get(key) match {
            case None => elems(key) = Set(value)
            case Some(set) => elems(key) = set + value
          }
      }
      this
    }
    override def clear(): Unit = elems.clear()
    override def result(): SortedMultiMap[K, V] = {
      new SortedMultiMap[K,V](SortedMap.empty[K, Set[V]] ++ elems)
    }
  }

}

/**
 * A special immutable collection class providing sorted access to keys as well as allowing
 * multiple values to be bound to a specific key. Note that you have to use the newly provided
 * methods to add multiple values under one key.
 *
 * More specifically this is a wrapper around the default [[scala.collection.SortedMap SortedMap]]
 * and also provides additional methods to add / remove multiple values bound to a specific key.
 * Internally it is represented as a SortedMap with entries of the form `key -> Set(value)`.
 * Proper, internal utility methods are provided to ensure modifying operations return this collection type again.
 * Some methods are delegated to the underlying 'SortedMap' for performance reasons (incomplete list of methods !!!)
 * @param delegate the underlying `SortedMap` that handles most functionality
 * @param ord the [[scala.math.Ordering ordering]] type class
 * @tparam K the type of keys
 * @tparam V the type of values
 */
class SortedMultiMap[K,V] private (private val delegate: SortedMap[K, Set[V]])
                                  (implicit ord: Ordering[K])
 extends SortedMap[K, Set[V]]
 with SortedMapLike[K, Set[V], SortedMultiMap[K,V]] {

  /**
   * Creates a new set containing the passed values.
   * Can be overridden in subclasses to return a different set implementation.
   * @param vs the elements to add to the newly constructed set
   * @return a new set containing the passed values
   */
  protected def makeNewSet(vs: V*): Set[V] = Set(vs :_*)

  /**
   * Adds a new value bound to a specific key.
   *
   * This special method is provided by the `MultiMap` functionality.
   * The returned map will have the same elements as this one,
   * if the binding was already present.
   *
   * @param key the key to add the value to
   * @param value the value to add
   * @return a new map containing the newly added binding
   */
  def addBinding(key: K, value: V): SortedMultiMap[K, V] = {
    get(key) match {
      case None =>
        val newSet = makeNewSet(value)
        new SortedMultiMap(delegate.updated(key, newSet))(ord)
      case Some(oldSet) =>
        val newSet = oldSet + value
        new SortedMultiMap(delegate.updated(key, newSet))(ord)
    }
  }

  /**
   * Removes an existing value bound to a specific key.
   *
   * This special method is is provided by the `MultiMap` functionality.
   * The returned map will be this map, if the key does not exists.
   * The returned map will have the same elements as this map,
   * if the binding was not present.
   *
   * @param key the key to remove the value from
   * @param value the value to remove
   * @return a new map not containing the newly removed binding
   */
  def removeBinding(key: K, value: V): SortedMultiMap[K, V] = {
    get(key) match {
      case None => this
      case Some(oldSet) =>
        val newSet = oldSet - value
        if (newSet.nonEmpty)
          new SortedMultiMap(delegate.updated(key, newSet))(ord)
        else
          new SortedMultiMap(delegate - key)(ord)
    }
  }

  /**
   * Checks if there exists a binding to `key` such that
   * it satisfies the predicate `p`.
   *
   * @param key the key of the binding
   * @param p the predicate used to check if a value bound to the `key` satisfies a condition
   * @return `true` if one or more values satisfied the predicate,
   *        `false` if the key did not exist or neither value satisfied the predicate
   */
  def bindingExists(key: K, p: V => Boolean): Boolean = get(key) match {
    case None => false
    case Some(set) => set exists p
  }

  override def get(key: K): Option[Set[V]] = delegate.get(key)

  override def +[B1 >: Set[V]](kv: (K, B1)): SortedMap[K, B1] = delegate + kv

  /**
   * Add a key/value pair to this map.
   * @param kv the key/value pair to add
   * @return A new map with the new binding added to this map
   */
  def +(kv: (K,Set[V])): SortedMultiMap[K, V] = new SortedMultiMap(delegate + kv)(ord)

  /**
   * Add all key/value pairs to this map.
   * @param xs a traversable collection which contains the pairs
   * @return A new map with the new bindings added from the other collection to this map
   */
  def ++(xs: GenTraversableOnce[(K, Set[V])]): SortedMultiMap[K, V] =
    new SortedMultiMap(delegate ++ xs)(ord)

  override def -(key: K): SortedMultiMap[K, V] = new SortedMultiMap(delegate - key)(ord)

  override def rangeImpl(from: Option[K], until: Option[K]): SortedMultiMap[K, V] =
    new SortedMultiMap(delegate.rangeImpl(from, until))(ord)

  override implicit def ordering: Ordering[K] = ord

  override def iterator: Iterator[(K, Set[V])] = delegate.iterator

  override def empty: SortedMultiMap[K, V] = SortedMultiMap.empty[K,V](ord)

  override protected[this] def newBuilder: mutable.Builder[(K, Set[V]), SortedMultiMap[K, V]] =
    SortedMultiMap.newBuilder[K,V](ord)


  // overridden for performance
  override def addString(b: StringBuilder, start: String, sep: String, end: String): StringBuilder = delegate.addString(b, start, sep, end)
  override def default(key: K): Set[V] = delegate.default(key)
  override def valuesIterator: Iterator[Set[V]] = delegate.valuesIterator
  override def values: Iterable[Set[V]] = delegate.values
  override def keys: Iterable[K] = delegate.keys
  override def keysIterator: Iterator[K] = delegate.keysIterator
  override def isDefinedAt(key: K): Boolean = delegate.isDefinedAt(key)
  override def contains(key: K): Boolean = delegate.contains(key)
  override def apply(key: K): Set[V] = delegate.apply(key)
  override def getOrElse[B1 >: Set[V]](key: K, default: => B1): B1 = delegate.getOrElse(key, default)
  override def isEmpty: Boolean = delegate.isEmpty
  override def nonEmpty: Boolean = delegate.nonEmpty
  override def size: Int = delegate.size
  override def keySet: SortedSet[K] = delegate.keySet
  override def lastKey: K = delegate.lastKey
  override def firstKey: K = delegate.firstKey
  override def hasDefiniteSize: Boolean = delegate.hasDefiniteSize
}