Merge pull request #9727 from NthPortal/topic/TrieMap-fixes-2/PR

Fix TrieMap#{filterInPlace,mapValuesInPlace}

project/MimaFilters.scala

@@ -42,7 +42,13 @@ object MimaFilters extends AutoPlugin {
     ProblemFilters.exclude[IncompatibleMethTypeProblem]("scala.reflect.io.FileZipArchive#LeakyEntry.this"),
     ProblemFilters.exclude[MissingClassProblem]("scala.reflect.io.FileZipArchive$zipFilePool$"),
 
-    )
+    // #9727
+    ProblemFilters.exclude[DirectMissingMethodProblem]("scala.collection.concurrent.TrieMap.filterInPlaceImpl"),
+    ProblemFilters.exclude[DirectMissingMethodProblem]("scala.collection.concurrent.TrieMap.mapValuesInPlaceImpl"),
+    ProblemFilters.exclude[DirectMissingMethodProblem]("scala.collection.convert.JavaCollectionWrappers#JConcurrentMapWrapper.filterInPlaceImpl"),
+    ProblemFilters.exclude[DirectMissingMethodProblem]("scala.collection.convert.JavaCollectionWrappers#JConcurrentMapWrapper.mapValuesInPlaceImpl"),
+
+  )
 
   override val buildSettings = Seq(
     mimaFailOnNoPrevious := false, // we opt everything out, knowing we only check library/reflect

src/library/scala/collection/concurrent/Map.scala

@@ -131,4 +131,22 @@ trait Map[K, V] extends scala.collection.mutable.Map[K, V] {
       case _ => this.updateWithAux(key)(remappingFunction)
     }
   }
+
+  private[collection] def filterInPlaceImpl(p: (K, V) => Boolean): this.type = {
+    val it = iterator
+    while (it.hasNext) {
+      val (k, v) = it.next()
+      if (p(k, v)) remove(k, v)
+    }
+    this
+  }
+
+  private[collection] def mapValuesInPlaceImpl(f: (K, V) => V): this.type = {
+    val it = iterator
+    while (it.hasNext) {
+      val (k, v) = it.next()
+      replace(k, v, f(k, v))
+    }
+    this
+  }
 }

src/library/scala/collection/concurrent/TrieMap.scala

@@ -153,7 +153,7 @@ private[collection] final class INode[K, V](bn: MainNode[K, V], g: Gen, equiv: E
     *              KEY_ABSENT - key wasn't there, insert only, do not overwrite
     *              KEY_PRESENT - key was there, overwrite only, do not insert
     *              other value `v` - only overwrite if the current value is this
-    *  @param hc the hashcode of `k``
+    *  @param hc the hashcode of `k`
     *
     *  @return     null if unsuccessful, Option[V] otherwise (indicating previous value bound to the key)
     */

src/library/scala/collection/mutable/Map.scala

@@ -171,17 +171,19 @@ trait MapOps[K, V, +CC[X, Y] <: MapOps[X, Y, CC, _], +C <: MapOps[K, V, CC, C]]
     * @param p  The test predicate
     */
   def filterInPlace(p: (K, V) => Boolean): this.type = {
-    if (nonEmpty) {
-      val array = this.toArray[Any] // scala/bug#7269 toArray avoids ConcurrentModificationException
-      val arrayLength = array.length
-      var i = 0
-      while (i < arrayLength) {
-        val (k, v) = array(i).asInstanceOf[(K, V)]
-        if (!p(k, v)) {
-          this -= k
+    if (!isEmpty) this match {
+      case tm: concurrent.Map[_, _] => tm.asInstanceOf[concurrent.Map[K, V]].filterInPlaceImpl(p)
+      case _ =>
+        val array = this.toArray[Any] // scala/bug#7269 toArray avoids ConcurrentModificationException
+        val arrayLength = array.length
+        var i = 0
+        while (i < arrayLength) {
+          val (k, v) = array(i).asInstanceOf[(K, V)]
+          if (!p(k, v)) {
+            this -= k
+          }
+          i += 1
         }
-        i += 1
-      }
     }
     this
   }
@@ -197,8 +199,9 @@ 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 = {
-    if (nonEmpty) this match {
+    if (!isEmpty) this match {
       case hm: mutable.HashMap[_, _] => hm.asInstanceOf[mutable.HashMap[K, V]].mapValuesInPlaceImpl(f)
+      case tm: concurrent.Map[_, _] => tm.asInstanceOf[concurrent.Map[K, V]].mapValuesInPlaceImpl(f)
       case _ =>
         val array = this.toArray[Any]
         val arrayLength = array.length

test/junit/scala/collection/concurrent/ConcurrentMapTestHelper.scala

@@ -0,0 +1,50 @@
+/*
+ * Scala (https://www.scala-lang.org)
+ *
+ * Copyright EPFL and Lightbend, Inc.
+ *
+ * Licensed under Apache License 2.0
+ * (http://www.apache.org/licenses/LICENSE-2.0).
+ *
+ * See the NOTICE file distributed with this work for
+ * additional information regarding copyright ownership.
+ */
+
+package scala.collection.concurrent
+
+import scala.concurrent.duration.SECONDS
+
+object ConcurrentMapTestHelper {
+  def genericTest_filterInPlace(newMap: => Map[String, Int]): Unit = {
+    val tester = new ConcurrentMapTester(newMap += "k" -> 0)
+
+    tester.runTasks(5, SECONDS)(
+      _.filterInPlace((_, v) => {
+        SECONDS.sleep(2)
+        v > 0
+      }),
+      map => {
+        SECONDS.sleep(1)
+        map("k") = 1
+      },
+    )
+
+    tester.assertContainsEntry("k", 1) // can get `0` if incorrectly implemented
+  }
+
+  def genericTest_mapValuesInPlace(newMap: => Map[String, Int]): Unit = {
+    val tester = new ConcurrentMapTester(newMap += "k" -> 0)
+    tester.runTasks(5, SECONDS)(
+      _.mapValuesInPlace((_, v) => {
+        SECONDS.sleep(2)
+        v + 5
+      }),
+      map => {
+        SECONDS.sleep(1)
+        map("k") = 1
+      },
+    )
+
+    tester.assertExistsEntry("k", x => x == 1 || x == 6) // can get `5` if incorrectly implemented
+  }
+}

test/junit/scala/collection/concurrent/ConcurrentMapTester.scala

@@ -0,0 +1,27 @@
+package scala.collection.concurrent
+
+import java.util.concurrent.Executors
+import scala.concurrent.duration.TimeUnit
+
+class ConcurrentMapTester[K, V](map: Map[K, V]) {
+  def runTasks(executionTimeout: Long, unit: TimeUnit)(tasks: (Map[K, V] => Unit)*): Unit = {
+    val exec = Executors.newCachedThreadPool()
+    for (task <- tasks) exec.execute(() => task(map))
+    exec.shutdown()
+    exec.awaitTermination(executionTimeout, unit)
+  }
+
+  @throws[AssertionError]
+  def assertContainsEntry(k: K, v: V): Unit = {
+    val value = map.get(k)
+    assert(value.isDefined, s"map does not contain key '$k'")
+    assert(value.contains(v), s"key '$k' is mapped to '${value.get}', not to '$v'")
+  }
+
+  @throws[AssertionError]
+  def assertExistsEntry(k: K, p: V => Boolean): Unit = {
+    val value = map.get(k)
+    assert(value.isDefined, s"map does not contain key '$k'")
+    assert(value.exists(p), s"key '$k' is mapped to '${value.get}', which does not match the predicate")
+  }
+}

test/junit/scala/collection/concurrent/TrieMapTest.scala

@@ -58,6 +58,16 @@ class TrieMapTest {
     check(List(("k", "v")))(_.view.mapValues(x => x))
   }
 
+  @Test
+  def filterInPlace(): Unit = {
+    ConcurrentMapTestHelper.genericTest_filterInPlace(TrieMap.empty)
+  }
+
+  @Test
+  def mapValuesInPlace(): Unit = {
+    ConcurrentMapTestHelper.genericTest_mapValuesInPlace(TrieMap.empty)
+  }
+
   @Test
   def customHashingAndEquiv_10481(): Unit = {
     val h = new Hashing[Int] { def hash(i: Int) = i % 4 }

test/junit/scala/collection/convert/JConcurrentMapWrapperTest.scala

@@ -0,0 +1,30 @@
+package scala.collection.convert
+
+import org.junit.Test
+
+import java.util.concurrent.{ConcurrentHashMap, ConcurrentSkipListMap}
+
+import scala.collection.concurrent.ConcurrentMapTestHelper
+import scala.jdk.CollectionConverters._
+
+class JConcurrentMapWrapperTest {
+  @Test
+  def CHM_filterInPlace(): Unit = {
+    ConcurrentMapTestHelper.genericTest_filterInPlace(new ConcurrentHashMap[String, Int].asScala)
+  }
+
+  @Test
+  def CHM_mapValuesInPlace(): Unit = {
+    ConcurrentMapTestHelper.genericTest_mapValuesInPlace(new ConcurrentHashMap[String, Int].asScala)
+  }
+
+  @Test
+  def CSLM_filterInPlace(): Unit = {
+    ConcurrentMapTestHelper.genericTest_filterInPlace(new ConcurrentSkipListMap[String, Int].asScala)
+  }
+
+  @Test
+  def CSLM_mapValuesInPlace(): Unit = {
+    ConcurrentMapTestHelper.genericTest_mapValuesInPlace(new ConcurrentSkipListMap[String, Int].asScala)
+  }
+}