scala mutable MultiMap implementation

MultiMap.scala

package com.github.gist.jdolson.collection.mutable

import scala.collection.mutable.{Set, Map, Iterable, Builder, Cloneable}
import scala.collection.generic.{CanBuildFrom, Shrinkable}
import scala.collection.{Iterator, IterableLike, breakOut}

object MultiMap {
  type Coll = MultiMap[_, _]

  implicit def canBuildFrom[A, B] = new CanBuildFrom[Coll, (A, B), MultiMap[A, B]] {
    def apply() = newBuilder[A, B]
    def apply(from: Coll) = newBuilder[A, B]
  }

  def newBuilder[A, B]: Builder[(A, B), MultiMap[A, B]] = empty

  def empty[A, B] = new MultiMap(Map.empty[A, Set[B]])

  def apply[A, B](elems: (A, B)*) = (newBuilder[A, B] ++= elems).result()

  def apply[A, B](elems: TraversableOnce[(A, B)]): MultiMap[A, B] = (newBuilder[A, B] ++= elems).result()

  def buildFrom[A, B](from: TraversableOnce[(A, TraversableOnce[B])]) = {
    val m = empty[A, B]
    for ((k, vs) <- from) m.addBindings(k, vs)
    m
  }

  private val hashSeed = "MultiMap".hashCode
}

class MultiMap[A, B] private (m: Map[A, Set[B]])
    extends Iterable[(A, B)]
       with IterableLike[(A, B), MultiMap[A, B]]
       with Builder[(A, B), MultiMap[A, B]]
       with Shrinkable[(A, B)]
       with Cloneable[MultiMap[A, B]]
{ self =>

  def empty = MultiMap.empty

  override protected[this] def newBuilder = MultiMap.newBuilder

  override def clone = MultiMap.buildFrom(m)

  def iterator: Iterator[(A, B)] = for ((k, vs) <- m.iterator; v <- vs.iterator) yield (k, v)

  def get(key: A): Set[B] = m.getOrElse(key, Set.empty[B])

  def clear(): Unit = m.clear()

  def result() = this

  def contains(key: A): Boolean = get(key).nonEmpty
  def contains(key: A, value: B): Boolean = get(key).contains(value)
  def contains(kv: (A, B)): Boolean = contains(kv._1, kv._2)

  def + [B1 >: B] (kv: (A, B1)): MultiMap[A, B1] = clone.asInstanceOf[MultiMap[A, B1]] += kv

  def += (kv: (A, B)): this.type = {
    addBinding(kv._1, kv._2)
    this
  }

  def addBinding(key: A, value: B) {
    m.getOrElseUpdate(key, Set.empty[B]) += value
  }

  def addBindings(key: A, values: TraversableOnce[B]) {
    m.getOrElseUpdate(key, Set.empty[B]) ++= values
  }

  def - (key: A) = clone -= key

  /** Remove all bindings for the given key */
  def -= (key: A): this.type = {
    m -= key
    this
  }

  def -= (kv: (A, B)): this.type = {
    removeBinding(kv._1, kv._2)
    this
  }

  def removeBinding(key: A, value: B): Unit = m.get(key) match {
    case Some(vs) => if ((vs -= value).isEmpty) m -= key
    case None =>
  }

  def removeBindings(key: A, values: Traversable[B]): Unit = m.get(key) match {
    case Some(vs) => if ((vs --= values).isEmpty) m -= key
    case None =>
  }

  override def equals(that: Any): Boolean = that match {
    case that: MultiMap[a, _] => (this eq that) || (that canEqual this) && (this.size == that.size) && {
      try {
        m.forall { case (k, vs) => that.get(k.asInstanceOf[a]) == vs }
      }
      catch { case _: ClassCastException => false }
    }
    case _ => false
  }

  override def hashCode = util.MurmurHash.symmetricHash(this, MultiMap.hashSeed)

  def asMap: Map[A, Set[B]] = m

  def inverse: MultiMap[B, A] = {
    val i = MultiMap.empty[B, A]
    foreach((k, v) => i.addBinding(v, k))
    i
  }

  def keys: collection.Iterable[A] = m.keys
  def keySet: collection.Set[A] = m.keySet
  def keysIterator: Iterator[A] = m.keysIterator

  def values: collection.Iterable[B] = new collection.Iterable[B] {
    def iterator = self.valuesIterator
    override def size = self.size
    override def foreach[U](f: B => U) = self.foreach((k, v) => f(v))
  }
  def valueSet: collection.Set[B] = m.flatMap { case (_, vs) => vs } (breakOut)
  def valuesIterator: Iterator[B] = m.valuesIterator.flatMap(_.iterator)

  // overrides for efficiency

  override def foreach[U](f: ((A, B)) => U): Unit = foreach(Function.untupled(f))

  def foreach[U](f: (A, B) => U): Unit = for ((k, vs) <- m; v <- vs) f(k, v)

  override def size = m.foldLeft(0) { case (n, (_, vs)) => n + vs.size }

  override def stringPrefix = "MultiMap"
}