Bulk operations for red-black trees

szeiger · szeiger · commit 8e6aaf3afa25 · 2018-10-02T21:26:14.000+02:00
diff --git a/library/src/scala/collection/immutable/RedBlackTree.scala b/library/src/scala/collection/immutable/RedBlackTree.scala
@@ -571,4 +571,132 @@ private[collection] object RedBlackTree {
   private[this] class ValuesIterator[A: Ordering, B](tree: Tree[A, B], focus: Option[A]) extends TreeIterator[A, B, B](tree, focus) {
     override def nextResult(tree: Tree[A, B]) = tree.value
   }
+
+  // Bulk operations based on "Just Join for Parallel Ordered Sets" (https://www.cs.cmu.edu/~guyb/papers/BFS16.pdf):
+
+  def union[A, B](t1: Tree[A, B], t2: Tree[A, B])(implicit ordering: Ordering[A]): Tree[A, B] = blacken(_union(t1, t2))
+
+  def intersect[A, B](t1: Tree[A, B], t2: Tree[A, B])(implicit ordering: Ordering[A]): Tree[A, B] = blacken(_intersect(t1, t2))
+
+  def difference[A, B](t1: Tree[A, B], t2: Tree[A, _])(implicit ordering: Ordering[A]): Tree[A, B] =
+    blacken(_difference(t1, t2.asInstanceOf[Tree[A, B]]))
+
+  private[this] def r[A, B](t: Tree[A, B]): Int = {
+    @tailrec def h(t: Tree[A, B], i: Int): Int =
+      if(t eq null) i+1 else if(isBlackTree(t)) h(t.left, i+1) else h(t.left, i)
+    if((t eq null) || isBlackTree(t)) 2*(h(t, 0)-1)
+    else 2*h(t, 0)-1
+  }
+
+  private[this] def rotateLeft[A, B](t: Tree[A, B]): Tree[A, B] =
+    mkTree(isBlackTree(t.right), t.right.key, t.right.value,
+      mkTree(isBlackTree(t), t.key, t.value, t.left, t.right.left),
+      t.right.right)
+
+  private[this] def rotateRight[A, B](t: Tree[A, B]): Tree[A, B] =
+    mkTree(isBlackTree(t.left), t.left.key, t.left.value,
+      t.left.left,
+      mkTree(isBlackTree(t), t.key, t.value, t.left.right, t.right))
+
+  private[this] def joinRightRB[A, B](tl: Tree[A, B], k: A, v: B, tr: Tree[A, B]): Tree[A, B] = {
+    if(r(tl) == (r(tr)/2)*2)
+      RedTree(k, v, tl, tr)
+    else {
+      val cc = isBlackTree(tl)
+      val rr = tl.right
+      val ttr = joinRightRB(rr, k, v, tr)
+      if(cc && isRedTree(ttr) && isRedTree(ttr.right)) {
+        val ttr2 = mkTree(isBlackTree(ttr), ttr.key, ttr.value, ttr.left, blacken(ttr.right))
+        val tt = mkTree(cc, tl.key, tl.value, tl.left, ttr2)
+        rotateLeft(tt)
+      } else mkTree(cc, tl.key, tl.value, tl.left, ttr)
+    }
+  }
+
+  private[this] def joinLeftRB[A, B](tl: Tree[A, B], k: A, v: B, tr: Tree[A, B]): Tree[A, B] = {
+    if(r(tr) == (r(tl)/2)*2)
+      RedTree(k, v, tl, tr)
+    else {
+      val cc = isBlackTree(tr)
+      val ll = tr.left
+      val ttl = joinLeftRB(tl, k, v, ll)
+      if(cc && isRedTree(ttl) && isRedTree(ttl.left)) {
+        val ttl2 = mkTree(isBlackTree(ttl), ttl.key, ttl.value, blacken(ttl.left), ttl.right)
+        val tt = mkTree(cc, tr.key, tr.value, ttl2, tr.right)
+        rotateRight(tt)
+      } else mkTree(cc, tr.key, tr.value, ttl, tr.right)
+    }
+  }
+
+  private[this] def join[A, B](tl: Tree[A, B], k: A, v: B, tr: Tree[A, B]): Tree[A, B] = {
+    val rtl = r(tl)
+    val rtr = r(tr)
+    if(rtl/2 > rtr/2) {
+      val tt = joinRightRB(tl, k, v, tr)
+      if(isRedTree(tt) && isRedTree(tt.right)) blacken(tt)
+      else tt
+    } else if(rtr/2 > rtl/2) {
+      val tt = joinLeftRB(tl, k, v, tr)
+      if(isRedTree(tt) && isRedTree(tt.left)) blacken(tt)
+      else tt
+    } else mkTree(!(isBlackTree(tl) && isBlackTree(tr)), k, v, tl, tr)
+  }
+
+  private[this] def split[A, B](t: Tree[A, B], k: A, v: B)(implicit ordering: Ordering[A]): (Tree[A, B], Boolean, Tree[A, B]) = {
+    if(t eq null) (null, false, null)
+    else {
+      val cmp = ordering.compare(k, t.key)
+      if(cmp == 0) (t.left, true, t.right)
+      else if(cmp < 0) {
+        val (ll, b, lr) = split(t.left, k, v)
+        (ll, b, join(lr, t.key, t.value, t.right))
+      } else {
+        val (rl, b, rr) = split(t.right, k, v)
+        (join(t.left, t.key, t.value, rl), b, rr)
+      }
+    }
+  }
+
+  private[this] def splitLast[A, B](t: Tree[A, B]): (Tree[A, B], A, B) =
+    if(t.right eq null) (t.left, t.key, t.value)
+    else {
+      val (tt, kk, vv) = splitLast(t.right)
+      (join(t.left, t.key, t.value, tt), kk, vv)
+    }
+
+  private[this] def join2[A, B](tl: Tree[A, B], tr: Tree[A, B]): Tree[A, B] =
+    if(tl eq null) tr
+    else {
+      val (ttl, k, v) = splitLast(tl)
+      join(ttl, k, v, tr)
+    }
+
+  private[this] def _union[A, B](t1: Tree[A, B], t2: Tree[A, B])(implicit ordering: Ordering[A]): Tree[A, B] =
+    if(t1 eq null) t2
+    else if(t2 eq null) t1
+    else {
+      val (l1, b, r1) = split(t1, t2.key, t2.value)
+      val tl = _union(l1, t2.left)
+      val tr = _union(r1, t2.right)
+      join(tl, t2.key, t2.value, tr)
+    }
+
+  private[this] def _intersect[A, B](t1: Tree[A, B], t2: Tree[A, B])(implicit ordering: Ordering[A]): Tree[A, B] =
+    if((t1 eq null) || (t2 eq null)) null
+    else {
+      val (l1, b, r1) = split(t1, t2.key, t2.value)
+      val tl = _intersect(l1, t2.left)
+      val tr = _intersect(r1, t2.right)
+      if(b) join(tl, t2.key, t2.value, tr)
+      else join2(tl, tr)
+    }
+
+  private[this] def _difference[A, B](t1: Tree[A, B], t2: Tree[A, B])(implicit ordering: Ordering[A]): Tree[A, B] =
+    if((t1 eq null) || (t2 eq null)) t1
+    else {
+      val (l1, b, r1) = split(t1, t2.key, t2.value)
+      val tl = _difference(l1, t2.left)
+      val tr = _difference(r1, t2.right)
+      join2(tl, tr)
+    }
 }
diff --git a/library/src/scala/collection/immutable/TreeMap.scala b/library/src/scala/collection/immutable/TreeMap.scala
@@ -3,7 +3,7 @@ package collection
 package immutable
 
 import scala.collection.immutable.{RedBlackTree => RB}
-import scala.collection.mutable.{Builder, ImmutableBuilder}
+import scala.collection.mutable.{Builder, ReusableBuilder}
 
 
 /** This class implements immutable maps using a tree.
@@ -25,7 +25,7 @@ import scala.collection.mutable.{Builder, ImmutableBuilder}
   *  @define mayNotTerminateInf
   *  @define willNotTerminateInf
   */
-final class TreeMap[K, +V] private (tree: RB.Tree[K, V])(implicit val ordering: Ordering[K])
+final class TreeMap[K, +V] private (private val tree: RB.Tree[K, V])(implicit val ordering: Ordering[K])
   extends AbstractMap[K, V]
     with SortedMap[K, V]
     with SortedMapOps[K, V, TreeMap, TreeMap[K, V]]
@@ -58,15 +58,29 @@ final class TreeMap[K, +V] private (tree: RB.Tree[K, V])(implicit val ordering:
   }
 
   override def concat[V1 >: V](that: collection.IterableOnce[(K, V1)]): TreeMap[K, V1] = {
-    val it = that.iterator
-    var t: RB.Tree[K, V1] = tree
-    while (it.hasNext) {
-      val (k, v) = it.next()
-      t = RB.update(t, k, v, overwrite = true)
+    val t = that match {
+      case tm: TreeMap[K, V] if ordering eq tm.ordering =>
+        RB.union(tree, tm.tree)
+      case _ =>
+        val it = that.iterator
+        var t: RB.Tree[K, V1] = tree
+        while (it.hasNext) {
+          val (k, v) = it.next()
+          t = RB.update(t, k, v, overwrite = true)
+        }
+        if(t eq tree) this else new TreeMap(t)
+        t
     }
     if(t eq tree) this else new TreeMap(t)
   }
 
+  override def removeAll(keys: IterableOnce[K]): TreeMap[K, V] = keys match {
+    case ts: TreeSet[K] if ordering eq ts.ordering =>
+      val t = RB.difference(tree, ts.tree)
+      if(t eq tree) this else new TreeMap(t)
+    case _ => super.removeAll(keys)
+  }
+
   /** A new TreeMap with the entry added is returned,
    *  assuming that key is <em>not</em> in the TreeMap.
    *
@@ -181,10 +195,23 @@ object TreeMap extends SortedMapFactory[TreeMap] {
         new TreeMap[K, V](t)
     }
 
-  def newBuilder[K : Ordering, V]: Builder[(K, V), TreeMap[K, V]] =
-    new ImmutableBuilder[(K, V), TreeMap[K, V]](empty) {
-      def addOne(elem: (K, V)): this.type = { elems = elems + elem; this }
-      override def addAll(xs: IterableOnce[(K, V)]): this.type = { elems = elems.concat(xs); this }
+  def newBuilder[K, V](implicit ordering: Ordering[K]): Builder[(K, V), TreeMap[K, V]] = new ReusableBuilder[(K, V), TreeMap[K, V]] {
+    private[this] var tree: RB.Tree[K, V] = null
+    def addOne(elem: (K, V)): this.type = { tree = RB.update(tree, elem._1, elem._2, overwrite = true); this }
+    override def addAll(xs: IterableOnce[(K, V)]): this.type = {
+      xs match {
+        case tm: TreeMap[K, V] if ordering eq tm.ordering =>
+          tree = RB.union(tree, tm.tree)
+        case _ =>
+          val it = xs.iterator
+          while (it.hasNext) {
+            val (k, v) = it.next()
+            tree = RB.update(tree, k, v, overwrite = true)
+          }
+      }
+      this
     }
-
+    def result(): TreeMap[K, V] = if(tree eq null) TreeMap.empty else new TreeMap[K, V](tree)
+    def clear(): Unit = { tree = null }
+  }
 }
diff --git a/library/src/scala/collection/immutable/TreeSet.scala b/library/src/scala/collection/immutable/TreeSet.scala
@@ -2,7 +2,7 @@ package scala
 package collection
 package immutable
 
-import mutable.{Builder, ImmutableBuilder}
+import mutable.{Builder, ReusableBuilder}
 import immutable.{RedBlackTree => RB}
 
 
@@ -23,7 +23,7 @@ import immutable.{RedBlackTree => RB}
   *  @define mayNotTerminateInf
   *  @define willNotTerminateInf
   */
-final class TreeSet[A] private (tree: RB.Tree[A, Null])(implicit val ordering: Ordering[A])
+final class TreeSet[A] private (private[immutable] val tree: RB.Tree[A, Null])(implicit val ordering: Ordering[A])
   extends AbstractSet[A]
     with SortedSet[A]
     with SortedSetOps[A, TreeSet, TreeSet[A]]
@@ -133,12 +133,34 @@ final class TreeSet[A] private (tree: RB.Tree[A, Null])(implicit val ordering: O
   }
 
   override def concat(that: collection.IterableOnce[A]): TreeSet[A] = {
-    val it = that.iterator
-    var t = tree
-    while (it.hasNext) t = RB.update(t, it.next(), null, overwrite = false)
+    val t = that match {
+      case ts: TreeSet[A] if ordering eq ts.ordering =>
+        RB.union(tree, ts.tree)
+      case _ =>
+        val it = that.iterator
+        var t = tree
+        while (it.hasNext) t = RB.update(t, it.next(), null, overwrite = false)
+        t
+    }
     if(t eq tree) this else newSet(t)
   }
 
+  override def intersect(that: collection.Set[A]): TreeSet[A] = that match {
+    case ts: TreeSet[A] if ordering eq ts.ordering =>
+      val t = RB.intersect(tree, ts.tree)
+      if(t eq tree) this else newSet(t)
+    case _ =>
+      super.intersect(that)
+  }
+
+  override def diff(that: collection.Set[A]): TreeSet[A] = that match {
+    case ts: TreeSet[A] if ordering eq ts.ordering =>
+      val t = RB.difference(tree, ts.tree)
+      if(t eq tree) this else newSet(t)
+    case _ =>
+      super.diff(that)
+  }
+
   override protected[this] def className = "TreeSet"
 }
 
@@ -163,10 +185,20 @@ object TreeSet extends SortedIterableFactory[TreeSet] {
         new TreeSet[E](t)
     }
 
-  def newBuilder[A : Ordering]: Builder[A, TreeSet[A]] =
-    new ImmutableBuilder[A, TreeSet[A]](empty) {
-      def addOne(elem: A): this.type = { elems = elems.incl(elem); this }
-      override def addAll(xs: IterableOnce[A]): this.type = { elems = elems.concat(xs); this }
+  def newBuilder[A](implicit ordering: Ordering[A]): Builder[A, TreeSet[A]] = new ReusableBuilder[A, TreeSet[A]] {
+    private[this] var tree: RB.Tree[A, Null] = null
+    def addOne(elem: A): this.type = { tree = RB.update(tree, elem, null, overwrite = false); this }
+    override def addAll(xs: IterableOnce[A]): this.type = {
+      xs match {
+        case ts: TreeSet[A] if ordering eq ts.ordering =>
+          tree = RB.union(tree, ts.tree)
+        case _ =>
+          val it = xs.iterator
+          while (it.hasNext) tree = RB.update(tree, it.next(), null, overwrite = false)
+      }
+      this
     }
-
+    def result(): TreeSet[A] = if(tree eq null) TreeSet.empty else new TreeSet[A](tree)
+    def clear(): Unit = { tree = null }
+  }
 }
