bimap scala

BiMap.scala

import cats.Order
import scala.collection.SortedMap

/**
  * A data structure representing a bidirectional mapping between
  * two key types. Each value in the bimap is associated with exactly
  * one value of the opposite type.
  */
final case class BiMap[A, B](ab: SortedMap[A, B], ba: SortedMap[B, A]) { self =>
  // TODO: remove unsafe constructor ^^^

  /** Find the left key corresponding to a given right key. */
  def getA(b: B): Option[A] = ba.get(b)
  /** Find the right key corresponding to a given left key. */
  def getB(a: A): Option[B] = ab.get(a)

  /**
    * Same as `getA`, but throws exception when
    * the key is not in the bimap.
    */
  @SuppressWarnings(Array("org.wartremover.warts.Throw"))
  def unsafeGetA(b: B): A = getA(b).getOrElse(
    throw new IllegalArgumentException("Right key not found"))

  /**
    * Same as `getB`, but throws exception when
    * the key is not in the bimap.
    */
  @SuppressWarnings(Array("org.wartremover.warts.Throw"))
  def unsafeGetB(a: A): B = getB(a).getOrElse(
    throw new IllegalArgumentException("Left key not found"))

  /**
    * Delete a value and its twin from a bimap.
    * When the value is not a member of the bimap,
    * the original bimap is returned.
    */
  def deleteA(a: A): BiMap[A, B] =
    getB(a) match {
      case Some(b) => copy(ab = ab - a, ba = ba -b)
      case None => this
    }

  /**
    * Same as `deleteA`, for the right key.
    */
  def deleteB(b: B): BiMap[A, B] =
    getA(b) match {
      case Some(a) => copy(ab = ab - a, ba = ba -b)
      case None => this
    }

  /**
    * Insert a pair of values into the bimap, associating them.
    * If either of the values is already in the bimap, any
    * overlapping bindings are deleted.
    */
  def insert(ab: (A, B)): BiMap[A, B] =
    ab match {
      case (a, b) => deleteA(a).deleteB(b).copy(
        ab = self.ab + ab,
        ba = self.ba + ab.swap
      )
    }
}
object BiMap {
  /**
    * Build a bimap from a list of pairs. If there are any overlapping
    * pairs in the list, the later ones will override the earlier ones.
    */
  def of[A , B](values: (A, B)*)(
    implicit orderA: Order[A], orderB: Order[B]
  ): BiMap[A, B] = values.foldLeft(BiMap.empty[A, B])(_.insert(_))

  /** Create an empty bimap. */
  def empty[A, B](
    implicit orderA: Order[A], orderB: Order[B]
  ): BiMap[A, B] = BiMap(
    SortedMap.empty[A, B](orderA.toOrdering),
    SortedMap.empty[B, A](orderB.toOrdering)
  )
}