Skip to content

Commit e72e6d3

Browse files
authored
Merge pull request scala/scala#8948 from retronym/topic/group-by-2.12.x
[nomerge] Reduce allocations in groupBy
2 parents 0926d73 + 44fc8c1 commit e72e6d3

File tree

2 files changed

+73
-10
lines changed

2 files changed

+73
-10
lines changed

library/src/scala/collection/TraversableLike.scala

Lines changed: 71 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import scala.annotation.unchecked.{uncheckedVariance => uV}
2020
import parallel.ParIterable
2121
import scala.collection.immutable.{::, List, Nil}
2222
import scala.language.higherKinds
23+
import scala.runtime.AbstractFunction0
2324

2425
/** A template trait for traversable collections of type `Traversable[A]`.
2526
*
@@ -451,17 +452,77 @@ trait TraversableLike[+A, +Repr] extends Any
451452
}
452453

453454
def groupBy[K](f: A => K): immutable.Map[K, Repr] = {
454-
val m = mutable.Map.empty[K, Builder[A, Repr]]
455-
for (elem <- this) {
456-
val key = f(elem)
457-
val bldr = m.getOrElseUpdate(key, newBuilder)
458-
bldr += elem
459-
}
460-
val b = immutable.Map.newBuilder[K, Repr]
461-
for ((k, v) <- m)
462-
b += ((k, v.result))
455+
object grouper extends AbstractFunction0[Builder[A, Repr]] with Function1[A, Unit] {
456+
var k0, k1, k2, k3: K = null.asInstanceOf[K]
457+
var v0, v1, v2, v3 = (null : Builder[A, Repr])
458+
var size = 0
459+
var hashMap: mutable.HashMap[K, Builder[A, Repr]] = null
460+
override def apply(): mutable.Builder[A, Repr] = {
461+
size += 1
462+
newBuilder
463+
}
464+
def apply(elem: A): Unit = {
465+
val key = f(elem)
466+
val bldr = builderFor(key)
467+
bldr += elem
468+
}
469+
def builderFor(key: K): Builder[A, Repr] =
470+
size match {
471+
case 0 =>
472+
k0 = key
473+
v0 = apply()
474+
v0
475+
case 1 =>
476+
if (k0 == key) v0
477+
else {k1 = key; v1 = apply(); v1 }
478+
case 2 =>
479+
if (k0 == key) v0
480+
else if (k1 == key) v1
481+
else {k2 = key; v2 = apply(); v2 }
482+
case 3 =>
483+
if (k0 == key) v0
484+
else if (k1 == key) v1
485+
else if (k2 == key) v2
486+
else {k3 = key; v3 = apply(); v3 }
487+
case 4 =>
488+
if (k0 == key) v0
489+
else if (k1 == key) v1
490+
else if (k2 == key) v2
491+
else if (k3 == key) v3
492+
else {
493+
hashMap = new mutable.HashMap
494+
hashMap += ((k0, v0))
495+
hashMap += ((k1, v1))
496+
hashMap += ((k2, v2))
497+
hashMap += ((k3, v3))
498+
val bldr = apply()
499+
hashMap(key) = bldr
500+
bldr
501+
}
502+
case _ =>
503+
hashMap.getOrElseUpdate(key, apply())
504+
}
463505

464-
b.result
506+
def result(): immutable.Map[K, Repr] =
507+
size match {
508+
case 0 => immutable.Map.empty
509+
case 1 => new immutable.Map.Map1(k0, v0.result())
510+
case 2 => new immutable.Map.Map2(k0, v0.result(), k1, v1.result())
511+
case 3 => new immutable.Map.Map3(k0, v0.result(), k1, v1.result(), k2, v2.result())
512+
case 4 => new immutable.Map.Map4(k0, v0.result(), k1, v1.result(), k2, v2.result(), k3, v3.result())
513+
case _ =>
514+
val it = hashMap.entriesIterator0
515+
val m1 = immutable.HashMap.newBuilder[K, Repr]
516+
while (it.hasNext) {
517+
val entry = it.next()
518+
m1.+=((entry.key, entry.value.result()))
519+
}
520+
m1.result()
521+
}
522+
523+
}
524+
this.seq.foreach(grouper)
525+
grouper.result()
465526
}
466527

467528
def forall(p: A => Boolean): Boolean = {

library/src/scala/collection/mutable/HashMap.scala

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,8 @@ extends AbstractMap[A, B]
172172
def next() = iter.next().value
173173
}
174174

175+
private[collection] def entriesIterator0: Iterator[DefaultEntry[A, B]] = entriesIterator
176+
175177
/** Toggles whether a size map is used to track hash map statistics.
176178
*/
177179
def useSizeMap(t: Boolean) = if (t) {

0 commit comments

Comments
 (0)