diff --git a/compat/src/main/scala-2.11_2.12/scala/collection/compat/PackageShared.scala b/compat/src/main/scala-2.11_2.12/scala/collection/compat/PackageShared.scala index f8541e10..693f8227 100644 --- a/compat/src/main/scala-2.11_2.12/scala/collection/compat/PackageShared.scala +++ b/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 { diff --git a/compat/src/test/scala/test/scala/collection/MutableMapTest.scala b/compat/src/test/scala/test/scala/collection/MutableMapTest.scala new file mode 100644 index 00000000..49e9b0ee --- /dev/null +++ b/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")) + } + +}