SetBy.scala

import shapeless._
sealed trait SetBy[A, B, FT <: A => B] {
  def +(a: A): SetBy[A, B, FT]

  def values: Set[A]

  override def toString: String = s"SetBy(${values.mkString(", ")})"
}

object SetBy {
  def apply[A, B, FT <: A => B](as: A*)(implicit ev: Witness.Aux[FT]): SetBy[A, B, FT] = {
    val by = ev.value
    val mp = as.map(a => by(a) -> a).toMap
    new SetByImpl(index = mp, values = mp.values.toSet)
  }

  private class SetByImpl[A, B, FT <: A => B](val index: Map[B, A], val values: Set[A])(implicit val ev: Witness.Aux[FT]) extends SetBy[A, B, FT] {
    private def by = ev.value
    def +(a: A): SetBy[A, B, FT] = {
      val by_a = by(a)
      index.get(by_a) match {
        case None => new SetByImpl(index = index + (by_a -> a), values = values + a)
        case Some(old_a) => new SetByImpl(index = index + (by_a -> a), values = values - old_a + a)
      }
    }
  }
}