From 3b327aee3ca01ee4adfed978a2c480726160d21d Mon Sep 17 00:00:00 2001
From: Chua Chee Seng <cheeseng@amaseng.com>
Date: Sat, 11 Jul 2026 13:55:54 +0800
Subject: [PATCH] Avoid intermediate Map allocations and tuple boxing in
 groupFlatMap

The Result/AbstractFunction1 pattern finalized builders via repeated
calls to immutable.Map#updated, which rebuilds trie nodes on every
call and discards all but the final version as garbage.

Replace it with a single Map builder populated via foreachEntry:

  val b = immutable.Map.newBuilder[K, CC[B]]
  m.foreachEntry((k, bldr) => b += ((k, bldr.result())))
  b.result()

This avoids both problems in one pass: results are staged once and
frozen once (no repeated .updated trie rebuilds), and foreachEntry
traverses m's internal table directly instead of going through
m.iterator, so no (K, Builder) pair is boxed per entry just to be
immediately unpacked again (as view.mapValues(_.result()).toMap
would do). Same O(n + k) complexity (n = elements, k = groups),
fewer allocations on both counts.

Also enhance tests in group-flat-map.scala to cover more cases:
duplicate keys from expansion, cross-element duplicates within a
group, mixed expansion sizes, and empty-expansion groups.
---
 library/src/scala/collection/Iterable.scala | 11 +++-------
 tests/run/group-flat-map.scala              | 23 +++++++++++++++++++++
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/library/src/scala/collection/Iterable.scala b/library/src/scala/collection/Iterable.scala
index d85ab1b8d7a7..395eaf28ea81 100644
--- a/library/src/scala/collection/Iterable.scala
+++ b/library/src/scala/collection/Iterable.scala
@@ -649,14 +649,9 @@ transparent trait IterableOps[+A, +CC[_], +C] extends Any with IterableOnce[A] w
       val bldr = m.getOrElseUpdate(k, iterableFactory.newBuilder[B])
       bldr ++= f(elem)
     }
-    class Result extends runtime.AbstractFunction1[(K, Builder[B, CC[B]]), Unit] {
-      var built = immutable.Map.empty[K, CC[B]]
-      def apply(kv: (K, Builder[B, CC[B]])) =
-        built = built.updated(kv._1, kv._2.result())
-    }
-    val result = new Result
-    m.foreach(result)
-    result.built
+    val b = immutable.Map.newBuilder[K, CC[B]]
+    m.foreachEntry((k, bldr) => b += ((k, bldr.result())))
+    b.result()
   }
 
   /** Partitions this $coll into a map according to a discriminator function `key`. All the values that
diff --git a/tests/run/group-flat-map.scala b/tests/run/group-flat-map.scala
index 0a93ea78c8e5..3324465a9ebf 100644
--- a/tests/run/group-flat-map.scala
+++ b/tests/run/group-flat-map.scala
@@ -25,3 +25,26 @@
   val xs = List("ab", "c", "de")
   assert(xs.groupFlatMap(_.length)(_.toList) ==
     xs.groupMap(_.length)(_.toList).view.mapValues(_.flatten).toMap)
+
+  // Test: duplicates from expansion are preserved in List result
+  // (flatmap semantics: each element expands to multiple elements)
+  val dupExpanded = List(1, 2).groupFlatMap(identity)(x => List(x, x))
+  assert(dupExpanded == Map(1 -> List(1, 1), 2 -> List(2, 2)))
+
+  // Test: duplicates across different elements in same group are preserved
+  // (encounter order + flatMap semantics)
+  val crossDup = List("aa", "ab").groupFlatMap(_.head)(s => List(s.length))
+  assert(crossDup == Map('a' -> List(2, 2)))
+
+  // Test: multiple groups with mixed expansion sizes
+  val mixed = List(1, 2, 3, 4).groupFlatMap(_ % 2)(x => List(x, x * 10, x * 100))
+  assert(mixed(1) == List(1, 10, 100, 3, 30, 300))
+  assert(mixed(0) == List(2, 20, 200, 4, 40, 400))
+
+  // Test: empty expansion creates empty group (not suppressed)
+  // This is the documented behavior
+  val emptyExpansions = List(1, 2, 3).groupFlatMap(identity)(x => if x == 2 then Nil else List(x))
+  // Note: key 2 gets an empty group (with a builder that was created but never added to)
+  assert(emptyExpansions(1) == List(1))
+  assert(emptyExpansions(2) == List.empty) // Empty group should exist
+  assert(emptyExpansions(3) == List(3))