trait MultiMap[A, B] extends mutable.MultiMap[A, B] {
  def foreachBinding(f: ((A, B)) => Unit): Unit =
    this.foreach {
      case (key, values) =>
        values.foreach { value =>
          f((key, value))
        }
    }

  def partitionByBindings(
    f: ((A, B)) => Boolean
  ): (mutable.MultiMap[A, B], mutable.MultiMap[A, B]) = {
    val truePairs = MultiMap[A, B]()
    val falsePairs = MultiMap[A, B]()

    this.foreachBinding {
      case binding @ (key, value) =>
        if (f(binding))
          truePairs.addBinding(key, value)
        else
          falsePairs.addBinding(key, value)
    }

    (truePairs, falsePairs)
  }
}

object MultiMap {
  def empty[A, B]: MultiMap[A, B] =
    new mutable.HashMap[A, mutable.Set[B]] with MultiMap[A, B]

  def apply[A, B](elems: (A, B)*): MultiMap[A, B] = {
    val result = empty[A, B]
    elems.foreach {
      case (key, value) => result.addBinding(key, value)
    }
    result
  }
}