Add mutable Map.updateWith (#508)

Author: emanb29

compat/src/main/scala-2.11_2.12/scala/collection/compat/PackageShared.scala

@@ -271,6 +271,10 @@ private[compat] trait PackageShared {
       self: scala.collection.Map[K, V]): MapExtensionMethods[K, V] =
     new MapExtensionMethods[K, V](self)
 
+  implicit def toMutableMapExtensionMethods[K, V](
+      self: scala.collection.mutable.Map[K, V]): MutableMapExtensionMethods[K, V] =
+    new MutableMapExtensionMethods[K, V](self)
+
   implicit def toMapViewExtensionMethods[K, V, C <: scala.collection.Map[K, V]](
       self: IterableView[(K, V), C]): MapViewExtensionMethods[K, V, C] =
     new MapViewExtensionMethods[K, V, C](self)
@@ -523,6 +527,19 @@ class MapExtensionMethods[K, V](private val self: scala.collection.Map[K, V]) ex
 
 }
 
+class MutableMapExtensionMethods[K, V](private val self: scala.collection.mutable.Map[K, V])
+    extends AnyVal {
+
+  def updateWith(key: K)(remappingFunction: (Option[V]) => Option[V]): Option[V] = {
+    val updatedEntry = remappingFunction(self.get(key))
+    updatedEntry match {
+      case Some(v) => self.update(key, v)
+      case None    => self.remove(key)
+    }
+    updatedEntry
+  }
+}
+
 class MapViewExtensionMethods[K, V, C <: scala.collection.Map[K, V]](
     private val self: IterableView[(K, V), C])
     extends AnyVal {

compat/src/test/scala/test/scala/collection/MutableMapTest.scala

@@ -0,0 +1,38 @@
+/*
+ * 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 test.scala.collection
+
+import org.junit.Test
+import org.junit.Assert._
+
+import scala.collection.compat._
+import scala.collection.{mutable => m}
+
+class MutableMapTest {
+
+  @Test
+  def updateWith: Unit = {
+    val map = m.Map("a" -> 1, "c" -> 3)
+    map.updateWith("b") {
+      case None    => Some(2)
+      case Some(x) => Some(-1)
+    }
+    map.updateWith("c") {
+      case None    => Some(-1)
+      case Some(_) => None
+    }
+    assertEquals(map.get("b"), Some(2))
+    assertFalse(map.contains("c"))
+  }
+
+}