From 3e3f22b2fc11dbe84de4a7ac0504815202e4cb5c Mon Sep 17 00:00:00 2001
From: Lukas Rytz <lukas.rytz@gmail.com>
Date: Mon, 16 Sep 2019 11:14:44 +0200
Subject: [PATCH] Don't use iterator to modify the underlying Map in
 `mapValuesInPlace`

Override the method in HashMap with a more efficient implementation.
---
 build.sbt                                         |  1 +
 .../scala/collection/mutable/HashMap.scala        | 15 +++++++++++++++
 src/library/scala/collection/mutable/Map.scala    | 13 +++++++++++--
 .../scala/collection/mutable/HashMapTest.scala    |  7 +++++++
 4 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/build.sbt b/build.sbt
index a3ea81a3f880..dbb9be2c8e01 100644
--- a/build.sbt
+++ b/build.sbt
@@ -149,6 +149,7 @@ val mimaFilterSettings = Seq {
     ProblemFilters.exclude[DirectMissingMethodProblem]("scala.collection.immutable.VectorPointer.gotoPosWritable1"),
     ProblemFilters.exclude[DirectMissingMethodProblem]("scala.collection.immutable.VectorPointer.gotoPosWritable1$default$4"),
     ProblemFilters.exclude[DirectMissingMethodProblem]("scala.collection.immutable.VectorPointer.nullSlotAndCopy$default$3"),
+    ProblemFilters.exclude[DirectMissingMethodProblem]("scala.collection.mutable.HashMap.mapValuesInPlaceImpl"),
   ),
 }
 
diff --git a/src/library/scala/collection/mutable/HashMap.scala b/src/library/scala/collection/mutable/HashMap.scala
index 6432b4ccfe52..fa07b836b1c6 100644
--- a/src/library/scala/collection/mutable/HashMap.scala
+++ b/src/library/scala/collection/mutable/HashMap.scala
@@ -507,6 +507,21 @@ class HashMap[K, V](initialCapacity: Int, loadFactor: Double)
     this
   }
 
+  // TODO in 2.14: rename to `mapValuesInPlace` and override the base version (not binary compatible)
+  private[mutable] def mapValuesInPlaceImpl(f: (K, V) => V): this.type = {
+    val len = table.length
+    var i = 0
+    while (i < len) {
+      var n = table(i)
+      while (n ne null) {
+        n.value = f(n.key, n.value)
+        n = n.next
+      }
+      i += 1
+    }
+    this
+  }
+
   override def mapFactory: MapFactory[HashMap] = HashMap
 
   override protected[this] def stringPrefix = "HashMap"
diff --git a/src/library/scala/collection/mutable/Map.scala b/src/library/scala/collection/mutable/Map.scala
index f956d3e321ac..6f1e253f8bab 100644
--- a/src/library/scala/collection/mutable/Map.scala
+++ b/src/library/scala/collection/mutable/Map.scala
@@ -199,8 +199,17 @@ trait MapOps[K, V, +CC[X, Y] <: MapOps[X, Y, CC, _], +C <: MapOps[K, V, CC, C]]
     * @return   the map itself.
     */
   def mapValuesInPlace(f: (K, V) => V): this.type = {
-    iterator foreach {
-      case (key, value) => update(key, f(key, value))
+    if (nonEmpty) this match {
+      case hm: mutable.HashMap[K, V] => hm.mapValuesInPlaceImpl(f)
+      case _ =>
+        val array = this.toArray[Any]
+        val arrayLength = array.length
+        var i = 0
+        while (i < arrayLength) {
+          val (k, v) = array(i).asInstanceOf[(K, V)]
+          update(k, f(k, v))
+          i += 1
+        }
     }
     this
   }
diff --git a/test/junit/scala/collection/mutable/HashMapTest.scala b/test/junit/scala/collection/mutable/HashMapTest.scala
index 5b76da1aff6c..0b9732ffa83c 100644
--- a/test/junit/scala/collection/mutable/HashMapTest.scala
+++ b/test/junit/scala/collection/mutable/HashMapTest.scala
@@ -251,6 +251,13 @@ class HashMapTest {
     assertEquals(hashMapCount4.updateWith(countingKey2)(transform), None)
     assertEquals(2, count) // once during hashtable construction, once during updateWith
     assertEquals(hashMapCount4, mutable.HashMap(countingKey1 -> "a"))
+  }
 
+  @Test
+  def t11737(): Unit = {
+    val m = mutable.HashMap(
+      0 -> 1, 1 -> 1, 2 -> 1, 3 -> 1, 4 -> 1, 5 -> 1, 6 -> 1, 7 -> 1, 8 -> 1, 9 -> 1, 241 -> 1)
+      .mapValuesInPlace((_, _) => -1)
+    assertTrue(m.toString, m.forall(_._2 == -1))
   }
 }