Implement @retronym's suggestions

Author: joshlemer

library/src/scala/collection/immutable/ChampHashMap.scala

@@ -935,7 +935,7 @@ private[immutable] final class HashMapBuilder[K, V] extends Builder[(K, V), Hash
   }
 
   /** Upserts a key/value pair into mapNode, mutably */
-  private def update(mapNode: MapNode[K, V], key: K, value: V, originalHash: Int, keyHash: Int, shift: Int): Unit = {
+  private[immutable] def update(mapNode: MapNode[K, V], key: K, value: V, originalHash: Int, keyHash: Int, shift: Int): Unit = {
     mapNode match {
       case bm: BitmapIndexedMapNode[K, V] =>
         val mask = maskFrom(keyHash, shift)
@@ -984,55 +984,6 @@ private[immutable] final class HashMapBuilder[K, V] extends Builder[(K, V), Hash
     }
   }
 
-  /** Inserts the this key/value only if the key is not present in the map already
-    * This method is safe to all, even when aliased, since it will only ensure unaliased if it is going to mutate */
-  private def updateIfNotExists(mapNode: MapNode[K, V], key: K, value: V, originalHash: Int, keyHash: Int, shift: Int): Unit = {
-    mapNode match {
-      case bm: BitmapIndexedMapNode[K, V] =>
-        val mask = maskFrom(keyHash, shift)
-        val bitpos = bitposFrom(mask)
-        if ((bm.dataMap & bitpos) != 0) {
-          val index = indexFrom(bm.dataMap, mask, bitpos)
-          val key0 = bm.getKey(index)
-
-          if (key0 == key) {
-            () // do nothing
-          } else {
-            ensureUnaliased()
-
-            val value0 = bm.getValue(index)
-            val key0UnimprovedHash = bm.getHash(index)
-            val key0Hash = improve(key0UnimprovedHash)
-
-            val subNodeNew: MapNode[K, V] =
-              bm.mergeTwoKeyValPairs(key0, value0, key0UnimprovedHash, key0Hash, key, value, originalHash, keyHash, shift + BitPartitionSize)
-
-            hash += keyHash
-            migrateFromInlineToNode(bm, bitpos, subNodeNew)
-          }
-
-        } else if ((bm.nodeMap & bitpos) != 0) {
-          val index = indexFrom(bm.nodeMap, mask, bitpos)
-          val subNode = bm.getNode(index)
-          val beforeSize = subNode.size
-          updateIfNotExists(subNode, key, value, originalHash, keyHash, shift + BitPartitionSize)
-          bm.size += subNode.size - beforeSize
-        } else {
-          ensureUnaliased()
-          insertValue(bm, bitpos, key, originalHash, keyHash, value)
-          hash += keyHash
-        }
-      case hc: HashCollisionMapNode[K, V] =>
-        val index = hc.indexOf(key)
-        if (index < 0) {
-          ensureUnaliased()
-          hash += keyHash
-          hc.content = hc.content.appended((key, value))
-        }
-    }
-  }
-
-
   /** If currently referencing aliased structure, copy elements to new mutable structure */
   private def ensureUnaliased() = {
     if (isAliased) copyElems()
@@ -1064,11 +1015,10 @@ private[immutable] final class HashMapBuilder[K, V] extends Builder[(K, V), Hash
     this
   }
 
-
-  private[collection] def addOneIfNotExists(key: K, value: V): this.type = {
-    val h = key.##
-    val im = improve(h)
-    updateIfNotExists(rootNode, key, value, h, im, 0)
+  def addOne(key: K, value: V): this.type = {
+    ensureUnaliased()
+    val originalHash = key.##
+    update(rootNode, key, value, originalHash, improve(originalHash), 0)
     this
   }
 

library/src/scala/collection/immutable/Map.scala

@@ -333,15 +333,8 @@ object Map extends MapFactory[Map] {
     }
   }
 
-  final class Map4[K, +V](
-    private[collection] val key1: K,
-    private[collection] val value1: V,
-    private[collection] val key2: K,
-    private[collection] val value2: V,
-    private[collection] val key3: K,
-    private[collection] val value3: V,
-    private[collection] val key4: K,
-    private[collection] val value4: V) extends AbstractMap[K, V] with StrictOptimizedIterableOps[(K, V), Iterable, Map[K, V]] {
+  final class Map4[K, +V](key1: K, value1: V, key2: K, value2: V, key3: K, value3: V, key4: K, value4: V)
+    extends AbstractMap[K, V] with StrictOptimizedIterableOps[(K, V), Iterable, Map[K, V]] {
 
     override def size: Int = 4
     override def knownSize: Int = 4
@@ -406,6 +399,13 @@ object Map extends MapFactory[Map] {
     override def foreach[U](f: ((K, V)) => U): Unit = {
       f((key1, value1)); f((key2, value2)); f((key3, value3)); f((key4, value4))
     }
+
+    private[immutable] def buildTo[V1 >: V](builder: HashMapBuilder[K, V1]): Unit = {
+      builder.addOne(key1, value1)
+      builder.addOne(key2, value2)
+      builder.addOne(key3, value3)
+      builder.addOne(key4, value4)
+    }
   }
 
   // scalac generates a `readReplace` method to discard the deserialized state (see https://github.com/scala/bug/issues/10412).
@@ -418,46 +418,49 @@ object Map extends MapFactory[Map] {
 @SerialVersionUID(3L)
 abstract class AbstractMap[K, +V] extends scala.collection.AbstractMap[K, V] with Map[K, V]
 
-private[collection] final class MapBuilderImpl[K, V] extends Builder[(K, V), Map[K, V]] {
+private[immutable] final class MapBuilderImpl[K, V] extends Builder[(K, V), Map[K, V]] {
   private[this] var elems: Map[K, V] = Map.empty
+  private[this] var switchedToHashMapBuilder: Boolean = false
   private[this] var hashMapBuilder: HashMapBuilder[K, V] = _
 
   override def clear(): Unit = {
     elems = Map.empty
     if (hashMapBuilder != null) {
       hashMapBuilder.clear()
     }
+    switchedToHashMapBuilder = false
   }
 
-  override def result(): Map[K, V] = {
-    if (hashMapBuilder == null || hashMapBuilder.size == 0) {
-      elems
-    } else if (elems.size == 4) {
-      val map4 = elems.asInstanceOf[Map4[K, V]]
+  override def result(): Map[K, V] =
+    if (switchedToHashMapBuilder) hashMapBuilder.result() else elems
 
-      hashMapBuilder.addOneIfNotExists(map4.key1, map4.value1)
-      hashMapBuilder.addOneIfNotExists(map4.key2, map4.value2)
-      hashMapBuilder.addOneIfNotExists(map4.key3, map4.value3)
-      hashMapBuilder.addOneIfNotExists(map4.key4, map4.value4)
-
-      hashMapBuilder.result()
-    } else {
-      // should never happen...
-      elems.foreach { case (k, v) => hashMapBuilder.addOneIfNotExists(k, v) }
-      hashMapBuilder.result()
-    }
-  }
-  override def addOne(elem: (K, V)) = {
-    if (elems.size < 4) {
+  def addOne(elem: (K, V)) = {
+    if (switchedToHashMapBuilder) {
+      hashMapBuilder.addOne(elem)
+    } else if (elems.size < 4) {
       elems = elems + elem
     } else {
-      if (hashMapBuilder == null) {
-        hashMapBuilder = new HashMapBuilder
+      // assert(elems.size == 4)
+      if (elems.contains(elem._1)) {
+        elems = elems + elem
+      } else {
+        switchedToHashMapBuilder = true
+        if (hashMapBuilder == null) {
+          hashMapBuilder = new HashMapBuilder
+        }
+        elems.asInstanceOf[Map4[K, V]].buildTo(hashMapBuilder)
+        hashMapBuilder.addOne(elem)
       }
-
-      hashMapBuilder.addOne(elem)
     }
 
     this
   }
+
+  override def addAll(xs: IterableOnce[(K, V)]): this.type =
+    if (switchedToHashMapBuilder) {
+      hashMapBuilder.addAll(xs)
+      this
+    } else {
+      super.addAll(xs)
+    }
 }