Merge pull request scala/scala#7248 from joshlemer/setbuilder-2

optimize s.c.i.HashSetBuilder

Author: retronym

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

@@ -725,19 +725,15 @@ private final class BitmapIndexedMapNode[K, +V](
 
   override def copy(): BitmapIndexedMapNode[K, V] = {
     val contentClone = new Array[Any](content.length)
-    var i = 0
-    val dataIndices = bitCount(dataMap) * 2
-    while (i < dataIndices) {
-      contentClone(i) = content(i)
-      i += 1
-    }
+    val dataIndices = bitCount(dataMap) * TupleLength
+    Array.copy(content, 0, contentClone, 0, dataIndices)
+    var i = dataIndices
     while (i < content.length) {
       contentClone(i) = content(i).asInstanceOf[MapNode[K, V]].copy()
       i += 1
     }
     new BitmapIndexedMapNode[K, V](dataMap, nodeMap, contentClone, originalHashes.clone(), size)
   }
-
 }
 
 private final class HashCollisionMapNode[K, +V ](
@@ -1142,7 +1138,6 @@ private[immutable] final class HashMapBuilder[K, V] extends Builder[(K, V), Hash
 
   /** Copy elements to new mutable structure */
   private def copyElems(): Unit = {
-    aliased = null
     rootNode = rootNode.copy()
   }
 

library/src/scala/collection/immutable/ChampHashSet.scala

@@ -9,7 +9,9 @@ import Hashing.improve
 import java.lang.Integer.{bitCount, numberOfTrailingZeros}
 import java.lang.System.arraycopy
 
+import scala.collection.immutable.Set.Set4
 import scala.util.hashing.MurmurHash3
+import scala.runtime.Statics.releaseFence
 
 /** This class implements immutable sets using a Compressed Hash-Array Mapped Prefix-tree.
   * See paper https://michael.steindorfer.name/publications/oopsla15.pdf for more details.
@@ -26,6 +28,8 @@ final class HashSet[A] private[immutable] (val rootNode: SetNode[A], val cachedJ
     with SetOps[A, HashSet, HashSet[A]]
     with StrictOptimizedIterableOps[A, HashSet, HashSet[A]] {
 
+  releaseFence()
+
   override def iterableFactory: IterableFactory[HashSet] = HashSet
 
   override def knownSize: Int = rootNode.size
@@ -68,6 +72,13 @@ final class HashSet[A] private[immutable] (val rootNode: SetNode[A], val cachedJ
     else this
   }
 
+  override def concat(that: IterableOnce[A]): HashSet[A] = {
+    val builder = iterableFactory.newBuilder[A]
+    builder ++= this
+    builder ++= that
+    builder.result()
+  }
+
   override def tail: HashSet[A] = this - head
 
   override def init: HashSet[A] = this - last
@@ -214,9 +225,15 @@ private[immutable] sealed abstract class SetNode[A] extends Node[SetNode[A]] {
 
   def subsetOf(that: SetNode[A], shift: Int): Boolean
 
+  def copy(): SetNode[A]
 }
 
-private final class BitmapIndexedSetNode[A](val dataMap: Int, val nodeMap: Int, val content: Array[Any], val originalHashes: Array[Int], val size: Int) extends SetNode[A] {
+private final class BitmapIndexedSetNode[A](
+   var dataMap: Int,
+   var nodeMap: Int,
+   var content: Array[Any],
+   var originalHashes: Array[Int],
+   var size: Int) extends SetNode[A] {
 
   import Node._
   import SetNode._
@@ -564,8 +581,8 @@ private final class BitmapIndexedSetNode[A](val dataMap: Int, val nodeMap: Int,
         (this eq node) ||
           (this.nodeMap == node.nodeMap) &&
             (this.dataMap == node.dataMap) &&
-             java.util.Arrays.equals(this.originalHashes, node.originalHashes) &&
-              deepContentEquality(this.content, node.content, content.length)
+            java.util.Arrays.equals(this.originalHashes, node.originalHashes) &&
+            deepContentEquality(this.content, node.content, content.length)
       case _ => false
     }
 
@@ -588,13 +605,24 @@ private final class BitmapIndexedSetNode[A](val dataMap: Int, val nodeMap: Int,
   override def hashCode(): Int =
     throw new UnsupportedOperationException("Trie nodes do not support hashing.")
 
+  override def copy(): BitmapIndexedSetNode[A] = {
+    val contentClone = new Array[Any](content.length)
+    val dataIndices = bitCount(dataMap)
+    Array.copy(content, 0, contentClone, 0, dataIndices)
+    var i = dataIndices
+    while (i < content.length) {
+      contentClone(i) = content(i).asInstanceOf[SetNode[A]].copy()
+      i += 1
+    }
+    new BitmapIndexedSetNode[A](dataMap, nodeMap, contentClone, originalHashes.clone(), size)
+  }
 }
 
-private final class HashCollisionSetNode[A](val originalHash: Int, val hash: Int, val content: Vector[A]) extends SetNode[A] {
+private final class HashCollisionSetNode[A](val originalHash: Int, val hash: Int, var content: Vector[Any]) extends SetNode[A] {
 
   import Node._
 
-  require(content.size >= 2)
+  require(content.length >= 2)
 
   def contains(element: A, originalHash: Int, hash: Int, shift: Int): Boolean =
     this.hash == hash && content.contains(element)
@@ -621,7 +649,7 @@ private final class HashCollisionSetNode[A](val originalHash: Int, val hash: Int
       // assert(updatedContent.size == content.size - 1)
 
       updatedContent.size match {
-        case 1 => new BitmapIndexedSetNode[A](bitposFrom(maskFrom(hash, 0)), 0, updatedContent.toArray, Array(originalHash), 1)
+        case 1 => new BitmapIndexedSetNode[A](bitposFrom(maskFrom(hash, 0)), 0, Array(updatedContent(0)), Array(originalHash), 1)
         case _ => new HashCollisionSetNode[A](originalHash, hash, updatedContent)
       }
     }
@@ -635,23 +663,28 @@ private final class HashCollisionSetNode[A](val originalHash: Int, val hash: Int
 
   def hasPayload: Boolean = true
 
-  def payloadArity: Int = content.size
+  def payloadArity: Int = content.length
 
-  def getPayload(index: Int): A = content(index)
+  def getPayload(index: Int): A = content(index).asInstanceOf[A]
 
   override def getHash(index: Int): Int = originalHash
 
   def sizePredicate: Int = SizeMoreThanOne
 
-  def size: Int = content.size
+  def size: Int = content.length
 
-  def foreach[U](f: A => U): Unit = content.foreach(f)
+  def foreach[U](f: A => U): Unit = {
+    var i = 0
+    while (i < content.length) {
+      f(getPayload(i))
+      i += 1
+    }
+  }
 
   def subsetOf(that: SetNode[A], shift: Int): Boolean = if (this eq that) true else that match {
     case node: BitmapIndexedSetNode[A] => false
-    case node: HashCollisionSetNode[A] => {
+    case node: HashCollisionSetNode[A] =>
       this.payloadArity <= node.payloadArity && this.content.forall(node.content.contains)
-    }
   }
 
   override def equals(that: Any): Boolean =
@@ -667,6 +700,8 @@ private final class HashCollisionSetNode[A](val originalHash: Int, val hash: Int
   override def hashCode(): Int =
     throw new UnsupportedOperationException("Trie nodes do not support hashing.")
 
+  override def copy() = new HashCollisionSetNode[A](originalHash, hash, content)
+
 }
 
 private final class SetIterator[A](rootNode: SetNode[A])
@@ -740,16 +775,211 @@ object HashSet extends IterableFactory[HashSet] {
       case _ => (newBuilder[A] ++= source).result()
     }
 
-  def newBuilder[A]: Builder[A, HashSet[A]] =
-    new ImmutableBuilder[A, HashSet[A]](empty) {
-      def addOne(element: A): this.type = {
-        elems = elems + element
-        this
-      }
-    }
+  def newBuilder[A]: Builder[A, HashSet[A]] = new HashSetBuilder
 
   // scalac generates a `readReplace` method to discard the deserialized state (see https://github.com/scala/bug/issues/10412).
   // This prevents it from serializing it in the first place:
   private[this] def writeObject(out: ObjectOutputStream): Unit = ()
   private[this] def readObject(in: ObjectInputStream): Unit = ()
 }
+
+private[collection] final class HashSetBuilder[A] extends Builder[A, HashSet[A]] {
+  import Node._
+  import SetNode._
+
+  private def newEmptyRootNode = new BitmapIndexedSetNode[A](0, 0, Array(), Array(), 0)
+
+  /** The last given out HashSet as a return value of `result()`, if any, otherwise null.
+    * Indicates that on next add, the elements should be copied to an identical structure, before continuing
+    * mutations. */
+  private var aliased: HashSet[A] = _
+
+  private def isAliased: Boolean = aliased != null
+
+  /** The root node of the partially build hashmap */
+  private var rootNode: SetNode[A] = newEmptyRootNode
+
+  /** The cached hash of the partially-built hashmap */
+  private var hash: Int = 0
+
+  /** Inserts element `elem` into array `as` at index `ix`, shifting right the trailing elems */
+  private def insertElement(as: Array[Int], ix: Int, elem: Int): Array[Int] = {
+    if (ix < 0) throw new ArrayIndexOutOfBoundsException
+    if (ix > as.length) throw new ArrayIndexOutOfBoundsException
+    val result = new Array[Int](as.length + 1)
+    arraycopy(as, 0, result, 0, ix)
+    result(ix) = elem
+    arraycopy(as, ix, result, ix + 1, as.length - ix)
+    result
+  }
+
+  /** Inserts key-value into the bitmapIndexMapNode. Requires that this is a new key-value pair */
+  private def insertValue[A1 >: A](bm: BitmapIndexedSetNode[A], bitpos: Int, key: A, originalHash: Int, keyHash: Int): Unit = {
+    val dataIx = bm.dataIndex(bitpos)
+    val idx = TupleLength * dataIx
+
+    val src = bm.content
+    val dst = new Array[Any](src.length + TupleLength)
+
+    // copy 'src' and insert 2 element(s) at position 'idx'
+    arraycopy(src, 0, dst, 0, idx)
+    dst(idx) = key
+    arraycopy(src, idx, dst, idx + TupleLength, src.length - idx)
+
+    val dstHashes = insertElement(bm.originalHashes, dataIx, originalHash)
+
+    bm.dataMap = bm.dataMap | bitpos
+    bm.content = dst
+    bm.originalHashes = dstHashes
+    bm.size += 1
+  }
+
+  /** Removes element at index `ix` from array `as`, shifting the trailing elements right */
+  private def removeElement(as: Array[Int], ix: Int): Array[Int] = {
+    if (ix < 0) throw new ArrayIndexOutOfBoundsException
+    if (ix > as.length - 1) throw new ArrayIndexOutOfBoundsException
+    val result = new Array[Int](as.length - 1)
+    arraycopy(as, 0, result, 0, ix)
+    arraycopy(as, ix + 1, result, ix, as.length - ix - 1)
+    result
+  }
+
+  /** Mutates `bm` to replace inline data at bit position `bitpos` with node `node` */
+  private def migrateFromInlineToNode(bm: BitmapIndexedSetNode[A], bitpos: Int, node: SetNode[A]): Unit = {
+    val dataIx = bm.dataIndex(bitpos)
+    val idxOld = TupleLength * dataIx
+    val idxNew = bm.content.length - TupleLength - bm.nodeIndex(bitpos)
+
+    val src = bm.content
+    val dst = new Array[Any](src.length - TupleLength + 1)
+
+    // copy 'src' and remove 2 element(s) at position 'idxOld' and
+    // insert 1 element(s) at position 'idxNew'
+    // assert(idxOld <= idxNew)
+    arraycopy(src, 0, dst, 0, idxOld)
+    arraycopy(src, idxOld + TupleLength, dst, idxOld, idxNew - idxOld)
+    dst(idxNew) = node
+    arraycopy(src, idxNew + TupleLength, dst, idxNew + 1, src.length - idxNew - TupleLength)
+
+    val dstHashes = removeElement(bm.originalHashes, dataIx)
+
+    bm.dataMap ^= bitpos
+    bm.nodeMap |= bitpos
+    bm.content = dst
+    bm.originalHashes = dstHashes
+    bm.size = bm.size - 1 + node.size
+  }
+
+  /** Mutates `bm` to replace inline data at bit position `bitpos` with updated key/value */
+  private def setValue[A1 >: A](bm: BitmapIndexedSetNode[A], bitpos: Int, elem: A): Unit = {
+    val dataIx = bm.dataIndex(bitpos)
+    val idx = TupleLength * dataIx
+    bm.content(idx) = elem
+  }
+
+  def update(setNode: SetNode[A], element: A, originalHash: Int, elementHash: Int, shift: Int): Unit =
+    setNode match {
+      case bm: BitmapIndexedSetNode[A] =>
+        val mask = maskFrom(elementHash, shift)
+        val bitpos = bitposFrom(mask)
+
+        if ((bm.dataMap & bitpos) != 0) {
+          val index = indexFrom(bm.dataMap, mask, bitpos)
+          val element0 = bm.getPayload(index)
+          val element0UnimprovedHash = bm.getHash(index)
+
+          if (element0UnimprovedHash == originalHash && element0 == element) {
+            setValue(bm, bitpos, element0)
+          } else {
+            val element0Hash = improve(element0UnimprovedHash)
+            val subNodeNew = bm.mergeTwoKeyValPairs(element0, element0UnimprovedHash, element0Hash, element, originalHash, elementHash, shift + BitPartitionSize)
+            hash += elementHash
+            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
+          update(subNode, element, originalHash, elementHash, shift + BitPartitionSize)
+          bm.size += subNode.size - beforeSize
+        } else {
+          insertValue(bm, bitpos, element, originalHash, elementHash)
+          hash += elementHash
+        }
+      case hc: HashCollisionSetNode[A] =>
+        val index = hc.content.indexOf(element)
+        if (index < 0) {
+          hash += elementHash
+          hc.content = hc.content.appended(element)
+        } else {
+          hc.content = hc.content.updated(index, element)
+        }
+    }
+
+  /** If currently referencing aliased structure, copy elements to new mutable structure */
+  private def ensureUnaliased():Unit = {
+    if (isAliased) copyElems()
+    aliased = null
+  }
+
+  /** Copy elements to new mutable structure */
+  private def copyElems(): Unit = {
+    rootNode = rootNode.copy()
+  }
+
+  override def result(): HashSet[A] =
+    if (rootNode.size == 0) {
+      HashSet.empty
+    } else if (aliased != null) {
+      aliased
+    } else {
+      aliased = new HashSet(rootNode, hash)
+      releaseFence()
+      aliased
+    }
+
+  override def addOne(elem: A): this.type = {
+    ensureUnaliased()
+    val h = elem.##
+    val im = improve(h)
+    update(rootNode, elem, h, im, 0)
+    this
+  }
+
+  override def addAll(xs: IterableOnce[A]) = {
+    ensureUnaliased()
+    xs match {
+      case hm: HashSet[A] =>
+        new ChampBaseIterator(hm.rootNode) {
+          while(hasNext) {
+            val originalHash = currentValueNode.getHash(currentValueCursor)
+            update(
+              setNode = rootNode,
+              element = currentValueNode.getPayload(currentValueCursor),
+              originalHash = originalHash,
+              elementHash = improve(originalHash),
+              shift = 0
+            )
+            currentValueCursor += 1
+          }
+        }
+      case other =>
+        val it = other.iterator
+        while(it.hasNext) addOne(it.next())
+    }
+
+    this
+  }
+
+  override def clear(): Unit = {
+    aliased = null
+    if (rootNode.size > 0) {
+      // if rootNode is empty, we will not have given it away anyways, we instead give out the reused Set.empty
+      rootNode = newEmptyRootNode
+    }
+    hash = 0
+  }
+
+  private[collection] def size: Int = rootNode.size
+}
+