Merge pull request scala/scala#8073 from lihaoyi/fast-update-with

override mutable.HashMap#updateWith and mutable.LinkedHashMap#updateWith for performance

Author: szeiger

library/src/scala/collection/mutable/HashMap.scala

@@ -106,6 +106,57 @@ class HashMap[K, V](initialCapacity: Int, loadFactor: Double)
  }
  }
 
+ // Override updateWith for performance, so we can do the update while hashing
+ // the input key only once and performing one lookup into the hash table
+ override def updateWith(key: K)(remappingFunction: Option[V] => Option[V]): Option[V] = {
+ if (getClass != classOf[HashMap[_, _]]) {
+ // subclasses of HashMap might customise `get` ...
+ super.updateWith(key)(remappingFunction)
+ } else {
+ val hash = computeHash(key)
+ val indexedHash = index(hash)
+
+ var foundNode: Node[K, V] = null
+ var previousNode: Node[K, V] = null
+ table(indexedHash) match {
+ case null =>
+ case nd =>
+ @tailrec
+ def findNode(prev: Node[K, V], nd: Node[K, V], k: K, h: Int): Unit = {
+ if (h == nd.hash && k == nd.key) {
+ previousNode = prev
+ foundNode = nd
+ }
+ else if ((nd.next eq null) || (nd.hash > h)) ()
+ else findNode(nd, nd.next, k, h)
+ }
+
+ findNode(null, nd, key, hash)
+ }
+
+ val previousValue = foundNode match {
+ case null => None
+ case nd => Some(nd.value)
+ }
+
+ val nextValue = remappingFunction(previousValue)
+
+ (previousValue, nextValue) match {
+ case (None, None) => // do nothing
+
+ case (Some(_), None) =>
+ if (previousNode != null) previousNode.next = foundNode.next
+ else table(indexedHash) = foundNode.next
+ contentSize -= 1
+
+ case (None, Some(value)) => put0(key, value, false, hash, indexedHash)
+
+ case (Some(_), Some(newValue)) => foundNode.value = newValue
+ }
+ nextValue
+ }
+ }
+
  override def subtractAll(xs: IterableOnce[K]): this.type = {
  if (size == 0) {
  return this

library/src/scala/collection/mutable/HashTable.scala

@@ -176,7 +176,11 @@ private[collection] /*abstract class*/ trait HashTable[A, B, Entry >: Null <: Ha
  /** Remove entry from table if present.
  */
  final def removeEntry(key: A) : Entry = {
- val h = index(elemHashCode(key))
+ removeEntry0(key, index(elemHashCode(key)))
+ }
+ /** Remove entry from table if present.
+ */
+ private[collection] final def removeEntry0(key: A, h: Int) : Entry = {
  var e = table(h).asInstanceOf[Entry]
  if (e != null) {
  if (elemEquals(e.key, key)) {

library/src/scala/collection/mutable/LinkedHashMap.scala

@@ -139,15 +139,17 @@ class LinkedHashMap[K, V]
  override def remove(key: K): Option[V] = {
  val e = table.removeEntry(key)
  if (e eq null) None
- else {
- if (e.earlier eq null) firstEntry = e.later
- else e.earlier.later = e.later
- if (e.later eq null) lastEntry = e.earlier
- else e.later.earlier = e.earlier
- e.earlier = null // Null references to prevent nepotism
- e.later = null
- Some(e.value)
- }
+ else Some(remove0(e))
+ }
+
+ private[this] def remove0(e: Entry): V = {
+ if (e.earlier eq null) firstEntry = e.later
+ else e.earlier.later = e.later
+ if (e.later eq null) lastEntry = e.earlier
+ else e.later.earlier = e.earlier
+ e.earlier = null // Null references to prevent nepotism
+ e.later = null
+ e.value
  }
 
  def addOne(kv: (K, V)): this.type = { put(kv._1, kv._2); this }
@@ -176,6 +178,34 @@ class LinkedHashMap[K, V]
  else Iterator.empty.next()
  }
 
+ // Override updateWith for performance, so we can do the update while hashing
+ // the input key only once and performing one lookup into the hash table
+ override def updateWith(key: K)(remappingFunction: Option[V] => Option[V]): Option[V] = {
+ val keyIndex = table.index(table.elemHashCode(key))
+ val entry = table.findEntry0(key, keyIndex)
+
+ val previousValue =
+ if (entry == null) None
+ else Some(entry.value)
+
+ val nextValue = remappingFunction(previousValue)
+
+ (previousValue, nextValue) match {
+ case (None, None) => // do nothing
+ case (Some(_), None) =>
+ remove0(entry)
+ table.removeEntry0(key, keyIndex)
+
+ case (None, Some(value)) =>
+ table.addEntry0(table.createNewEntry(key, value), keyIndex)
+
+ case (Some(_), Some(value)) =>
+ entry.value = value
+ }
+
+ nextValue
+ }
+
  override def valuesIterator: Iterator[V] = new AbstractIterator[V] {
  private[this] var cur = firstEntry
  def hasNext = cur ne null