def groupBy[K](f: A => K): immutable.Map[K, C] = {
    val m = mutable.Map.empty[K, Builder[A, C]]
    val it = iterator
    while (it.hasNext) {
      val elem = it.next()
      val key = f(elem)
      val bldr = m.getOrElseUpdate(key, newSpecificBuilder)
      bldr += elem
    }
    var result = immutable.HashMap.empty[K, C]
    val mapIt = m.iterator
    while (mapIt.hasNext) {
      val (k, v) = mapIt.next()
      result = result.updated(k, v.result())
    }
    result
  }

  def groupByMutableTransform[K](f: A => K): immutable.Map[K, C] = {
    val m = mutable.Map.empty[K, Builder[A, C]]
    val it = iterator
    while (it.hasNext) {
      val elem = it.next()
      val key = f(elem)
      val bldr = m.getOrElseUpdate(key, newSpecificBuilder)
      bldr += elem
    }
    m.asInstanceOf[mutable.Map[K, Any]].mapValuesInPlace((_, v) => v.asInstanceOf[Builder[A, C]].result())
    val m1  = new immutable.HashMapBuilder[K, C]
    m1.addAll(m.asInstanceOf[Map[K, C]])
    m1.result()
  }

  def groupByBuilderGetOrElse[K](f: A => K): immutable.Map[K, C] = {
    val m  = new immutable.HashMapBuilder[K, Builder[A, C]]
    val it = iterator
    while (it.hasNext) {
      val elem = it.next()
      val key  = f(elem)
      val bldr = m.getOrElse(key, null) match {
        case null =>
          val builder = newSpecificBuilder
          m.addOne(key, builder)
          builder
        case v    => v
      }
      bldr += elem
    }
    m.result(_.result())
  }

package scala.collection

import java.util.concurrent.TimeUnit
import org.openjdk.jmh.annotations.{Benchmark, BenchmarkMode, Fork, Level, Measurement, Mode, OutputTimeUnit, Param, Scope, Setup, State, Threads, Warmup}
import org.openjdk.jmh.infra.Blackhole
import scala.collection.mutable.ArrayBuffer


@BenchmarkMode(Array(Mode.AverageTime))
@Fork(2)
@Threads(1)
@Warmup(iterations = 10)
@Measurement(iterations = 10)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
@State(Scope.Benchmark)
class GroupByBenchmark {
  @Param(Array("10", "100", "1000"))
  var size: Int = _
  @Param(Array("original", "BuilderGetOrElse", "MutableTransform"))
  var impl      = "original"
  @Param(Array("0"))
  var hashCodeCost: Int = 0

  case class Key(a: Int) {
    override def hashCode(): Int = {
      Blackhole.consumeCPU(hashCodeCost)
      Integer.hashCode(a)
    }
  }
  case class Groupable(a: Int) {
    val groupKey1 = new Key(a) // N groups
    val groupKey2 = new Key(a % 8) // 8 groups in total
  }
  var groupables: ArrayBuffer[Groupable] = _

  class GroupByWrapper[A](as: collection.Iterable[A]) extends collection.Iterable[A] {
    override def iterator: Iterator[A] = as.iterator
    override protected def newSpecificBuilder = new mutable.Builder[A, Iterable[A]] {
      override def clear(): Unit = ()
      override def result(): Iterable[A] = Nil
      override def addOne(elem: A): this.type = this
    }
  }

  @Setup(Level.Trial) def initKeys(): Unit = {
    groupables = ArrayBuffer.from((0 to size).iterator.map(Groupable(_)))
  }

  @Benchmark def groupByKey1(): AnyRef = {
    groupBy(groupables)(_.groupKey1)
  }
  @Benchmark def groupByKey2(): AnyRef = {
    groupBy(groupables)(_.groupKey2)
  }
  @Benchmark def groupByKey1CheapBuilder(): AnyRef = {
    groupBy(new GroupByWrapper[Groupable](groupables))(_.groupKey1)
  }
  @Benchmark def groupByKey2CheapBuilder(): AnyRef = {
    groupBy(new GroupByWrapper[Groupable](groupables))(_.groupKey2)
  }
  private def groupBy[A, B](as: collection.Iterable[A])(f: A => B) = impl match {
    case "original" => as.groupBy(f)
    case "BuilderGetOrElse" => as.groupByBuilderGetOrElse(f)
    case "MutableTransform" => as.groupByMutableTransform(f)
  }
}