BidirectionalMultiMap.scala

BidirectionalMultiMap.scala

/*
  Immutable bidirectional multimap with unique set of values and keys.
 */
case class BiMultiMap[K, V](
                             forward: Map[K, Set[V]] = Map.empty,
                             backward: Map[V, Set[K]] = Map.empty
                           ) {
  import BiMultiMap._
  def +(kv: (K, Set[V])): BiMultiMap[K, V] = {
    val (nextForward, nextBackward) = kv match {
      case (key, values) => {
        values.foldLeft((forward, backward)) {
          case ((fw, bw), value) => (fw.addBinding(key, value), bw.addBinding(value, key))
        }
      }
    }
    this.copy(forward = nextForward, backward = nextBackward)
  }

  def -(key: K): BiMultiMap[K, V] = {
    forward.get(key) match {
      case Some(values) => {
        val (nextForward, nextBackward) = values.foldLeft((forward, backward)) {
            case ((fw, bw), value) => (fw.removeBinding(key, value), bw.removeBinding(value, key))
          }
        copy(forward = nextForward, backward = nextBackward)
      }
      case _ => this
    }
  }

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

  def seq: Map[K, Set[V]] = forward.seq

  def get(key: K): Set[V] = forward.get(key) match {
    case Some(set) => set
    case _ => Set.empty
  }

  def contains(k: K) : Boolean = forward.contains(k)

  def exists(p: ((K, Set[V])) => Boolean): Boolean = forward.exists(p)

  def isEmpty: Boolean = forward.isEmpty

  def inverse: BiMultiMap[V, K] =
    this.copy(forward = backward, backward = forward)
}

object BiMultiMap {
  def apply[K, V](): BiMultiMap[K, V] = BiMultiMap(Map.empty, Map.empty)
  def empty[K, V]: BiMultiMap[K, V] = BiMultiMap()

  implicit class SetMultiMap[A, B](val map: Map[A, Set[B]]) extends AnyVal {
    def addBinding(key: A, value: B): Map[A, Set[B]] =
      map + (key -> { map.getOrElse(key, Set.empty) + value })

    def removeBinding(key: A, value: B): Map[A, Set[B]] = map.get(key) match {
      case None => map
      case Some(set) if set.size == 1 && set.head == value => map - key
      case Some(set) => map + (key -> (set - value))
    }
  }
}