Deletion from a TreeSet/TreeMap of a key that doesn't exist should not allocate

pass the tree into balance/balLeft/balRight methods so that is can be reused is unchanged

prefer use of isBlack/isRed to isRedTree when we know the tree is not null

Author: mkeskells

library/src/scala/collection/immutable/RedBlackTree.scala

@@ -211,19 +211,25 @@ private[collection] object NewRedBlackTree {
 
   def tail[A, B](tree: Tree[A, B]): Tree[A, B] = {
     def _tail(tree: Tree[A, B]): Tree[A, B] =
-      if(tree eq null) throw new NoSuchElementException("empty tree")
-      else if(tree.left eq null) tree.right
-      else if(isBlackTree(tree.left)) balLeft(tree.key, tree.value, _tail(tree.left), tree.right)
-      else RedTree(tree.key, tree.value, _tail(tree.left), tree.right)
+      if (tree eq null) throw new NoSuchElementException("empty tree")
+      else {
+        val tl = tree.left
+        if (tl eq null) tree.right
+        else if (tl.isBlack) balLeft(tree, _tail(tl), tree.right)
+        else tree.redWithLeft(_tail(tree.left))
+      }
     blacken(_tail(tree))
   }
 
   def init[A, B](tree: Tree[A, B]): Tree[A, B] = {
     def _init(tree: Tree[A, B]): Tree[A, B] =
-      if(tree eq null) throw new NoSuchElementException("empty tree")
-      else if(tree.right eq null) tree.left
-      else if(isBlackTree(tree.right)) balRight(tree.key, tree.value, tree.left, _init(tree.right))
-      else RedTree(tree.key, tree.value, tree.left, _init(tree.right))
+      if (tree eq null) throw new NoSuchElementException("empty tree")
+      else {
+        val tr = tree.right
+        if (tr eq null) tree.left
+        else if (tr.isBlack) balRight(tree, tree.left, _init(tr))
+        else tree.redWithRight(_init(tr))
+      }
     blacken(_init(tree))
   }
 
@@ -306,7 +312,7 @@ private[collection] object NewRedBlackTree {
     else tree
   }
 
-  def isBlack(tree: Tree[_, _]) = (tree eq null) || isBlackTree(tree)
+  def isBlack(tree: Tree[_, _]) = (tree eq null) || tree.isBlack
 
   @`inline` private[this] def isRedTree(tree: Tree[_, _]) = (tree ne null) && tree.isRed
   @`inline` private[this] def isBlackTree(tree: Tree[_, _]) = (tree ne null) && tree.isBlack
@@ -318,8 +324,10 @@ private[collection] object NewRedBlackTree {
   private[this] def maybeBlacken[A, B](t: Tree[A, B]): Tree[A, B] =
     if(isBlack(t)) t else if(isRedTree(t.left) || isRedTree(t.right)) t.black else t
 
-  private[this] def mkTree[A, B](isBlack: Boolean, k: A, v: B, l: Tree[A, B], r: Tree[A, B]) =
-    if (isBlack) BlackTree(k, v, l, r) else RedTree(k, v, l, r)
+  private[this] def mkTree[A, B](isBlack: Boolean, key: A, value: B, left: Tree[A, B], right: Tree[A, B]) = {
+    val sizeAndColour = sizeOf(left) + sizeOf(right) + 1 | (if(isBlack) initialBlackCount else initialRedCount)
+    new Tree(key, value.asInstanceOf[AnyRef], left, right, sizeAndColour)
+  }
 
   /** Create a new balanced tree where `newLeft` replaces `tree.left`. */
   private[this] def balanceLeft[A, B1](tree: Tree[A, B1], newLeft: Tree[A, B1]): Tree[A, B1] = {
@@ -722,6 +730,15 @@ private[collection] object NewRedBlackTree {
         new Tree(key, value.asInstanceOf[AnyRef], newLeft, _right, initialBlackCount | size)
       }
     }
+    private[NewRedBlackTree] def redWithLeft[B1 >: B](newLeft: Tree[A, B1]): Tree[A, B1] = {
+      //assertNotMutable(this)
+      //assertNotMutable(newLeft)
+      if ((newLeft eq _left) && isRed) this
+      else {
+        val size = sizeOf(newLeft) + sizeOf(_right) + 1
+        new Tree(key, value.asInstanceOf[AnyRef], newLeft, _right, initialRedCount | size)
+      }
+    }
     private[NewRedBlackTree] def blackWithRight[B1 >: B](newRight: Tree[A, B1]): Tree[A, B1] = {
       //assertNotMutable(this)
       //assertNotMutable(newRight)
@@ -731,6 +748,15 @@ private[collection] object NewRedBlackTree {
         new Tree(key, value.asInstanceOf[AnyRef], _left, newRight, initialBlackCount | size)
       }
     }
+    private[NewRedBlackTree] def redWithRight[B1 >: B](newRight: Tree[A, B1]): Tree[A, B1] = {
+      //assertNotMutable(this)
+      //assertNotMutable(newLeft)
+      if ((newRight eq _right) && isRed) this
+      else {
+        val size = sizeOf(_left) + sizeOf(newRight) + 1
+        new Tree(key, value.asInstanceOf[AnyRef], _left, newRight, initialRedCount | size)
+      }
+    }
     private[NewRedBlackTree] def withLeftRight[B1 >: B](newLeft: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = {
       //assertNotMutable(this)
       //assertNotMutable(newLeft)
@@ -741,6 +767,26 @@ private[collection] object NewRedBlackTree {
         new Tree(key, value.asInstanceOf[AnyRef], newLeft, newRight, (_count & colourBit) | size)
       }
     }
+    private[NewRedBlackTree] def redWithLeftRight[B1 >: B](newLeft: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = {
+      //assertNotMutable(this)
+      //assertNotMutable(newLeft)
+      //assertNotMutable(newRight)
+      if ((newLeft eq _left) && (newRight eq _right) && isRed) this
+      else {
+        val size = sizeOf(newLeft) + sizeOf(newRight) + 1
+        new Tree(key, value.asInstanceOf[AnyRef], newLeft, newRight, initialRedCount | size)
+      }
+    }
+    private[NewRedBlackTree] def blackWithLeftRight[B1 >: B](newLeft: Tree[A, B1], newRight: Tree[A, B1]): Tree[A, B1] = {
+      //assertNotMutable(this)
+      //assertNotMutable(newLeft)
+      //assertNotMutable(newRight)
+      if ((newLeft eq _left) && (newRight eq _right) && isBlack) this
+      else {
+        val size = sizeOf(newLeft) + sizeOf(newRight) + 1
+        new Tree(key, value.asInstanceOf[AnyRef], newLeft, newRight, initialBlackCount | size)
+      }
+    }
   }
   //see #Tree docs "Colour, mutablity and size encoding"
   //we make these final vals because the optimiser inlines them, without reference to the enclosing module
@@ -956,7 +1002,7 @@ private[collection] object NewRedBlackTree {
       if((v2.asInstanceOf[AnyRef] eq v.asInstanceOf[AnyRef])
         && (l2 eq l)
         && (r2 eq r)) t.asInstanceOf[Tree[A, C]]
-      else mkTree(isBlackTree(t), k, v2, l2, r2)
+      else mkTree(t.isBlack, k, v2, l2, r2)
     }
 
   def filterEntries[A, B](t: Tree[A, B], f: (A, B) => Boolean): Tree[A, B] = if(t eq null) null else {
@@ -1064,67 +1110,72 @@ private[collection] object NewRedBlackTree {
   // Red-Black Trees in a Functional Setting, Chris Okasaki: [[https://wiki.rice.edu/confluence/download/attachments/2761212/Okasaki-Red-Black.pdf]] */
 
   private[this] def del[A, B](tree: Tree[A, B], k: A)(implicit ordering: Ordering[A]): Tree[A, B] = if (tree eq null) null else {
-    def delLeft =
-      if (isBlackTree(tree.left)) balLeft(tree.key, tree.value, del(tree.left, k), tree.right)
-      else RedTree(tree.key, tree.value, del(tree.left, k), tree.right)
-    def delRight =
-      if (isBlackTree(tree.right)) balRight(tree.key, tree.value, tree.left, del(tree.right, k))
-      else RedTree(tree.key, tree.value, tree.left, del(tree.right, k))
     val cmp = ordering.compare(k, tree.key)
-    if (cmp < 0) delLeft
-    else if (cmp > 0) delRight
-    else append(tree.left, tree.right)
+    if (cmp < 0) {
+      val newLeft = del(tree.left, k)
+      if (newLeft eq tree.left) tree
+      else if (isBlackTree(tree.left)) balLeft(tree, newLeft, tree.right)
+      else tree.redWithLeft(newLeft)
+    } else if (cmp > 0) {
+      val newRight = del(tree.right, k)
+      if (newRight eq tree.right) tree
+      else if (isBlackTree(tree.right)) balRight(tree, tree.left, newRight)
+      else tree.redWithRight(newRight)
+    } else append(tree.left, tree.right)
   }
 
-  private[this] def balance[A, B](x: A, xv: B, tl: Tree[A, B], tr: Tree[A, B]) =
+  private[this] def balance[A, B](tree: Tree[A,B], tl: Tree[A, B], tr: Tree[A, B]): Tree[A, B] =
     if (isRedTree(tl)) {
-      if (isRedTree(tr)) RedTree(x, xv, tl.black, tr.black)
-      else if (isRedTree(tl.left)) RedTree(tl.key, tl.value, tl.left.black, BlackTree(x, xv, tl.right, tr))
-      else if (isRedTree(tl.right))
-        RedTree(tl.right.key, tl.right.value, BlackTree(tl.key, tl.value, tl.left, tl.right.left), BlackTree(x, xv, tl.right.right, tr))
-      else BlackTree(x, xv, tl, tr)
+      if (isRedTree(tr)) tree.redWithLeftRight(tl.black, tr.black)
+      else if (isRedTree(tl.left))  tl.withLeftRight(tl.left.black, tree.blackWithLeftRight(tl.right, tr))
+      else if (isRedTree(tl.right)) tl.right.withLeftRight(tl.blackWithRight(tl.right.left), tree.blackWithLeftRight(tl.right.right, tr))
+      else tree.blackWithLeftRight(tl, tr)
     } else if (isRedTree(tr)) {
-      if (isRedTree(tr.right)) RedTree(tr.key, tr.value, BlackTree(x, xv, tl, tr.left), tr.right.black)
-      else if (isRedTree(tr.left))
-        RedTree(tr.left.key, tr.left.value, BlackTree(x, xv, tl, tr.left.left), BlackTree(tr.key, tr.value, tr.left.right, tr.right))
-      else BlackTree(x, xv, tl, tr)
-    } else BlackTree(x, xv, tl, tr)
-
-  private[this] def balLeft[A, B](x: A, xv: B, tl: Tree[A, B], tr: Tree[A, B]) =
-    if (isRedTree(tl)) RedTree(x, xv, tl.black, tr)
-    else if (isBlackTree(tr)) balance(x, xv, tl, tr.red)
+      if (isRedTree(tr.right))     tr.withLeftRight(tree.blackWithLeftRight(tl, tr.left), tr.right.black)
+      else if (isRedTree(tr.left)) tr.left.withLeftRight(tree.blackWithLeftRight(tl, tr.left.left), tr.blackWithLeftRight(tr.left.right, tr.right))
+      else tree.blackWithLeftRight(tl, tr)
+    } else tree.blackWithLeftRight(tl, tr)
+
+  private[this] def balLeft[A, B](tree: Tree[A,B], tl: Tree[A, B], tr: Tree[A, B]): Tree[A, B] =
+    if (isRedTree(tl)) tree.redWithLeftRight(tl.black, tr)
+    else if (isBlackTree(tr)) balance(tree, tl, tr.red)
     else if (isRedTree(tr) && isBlackTree(tr.left))
-      RedTree(tr.left.key, tr.left.value, BlackTree(x, xv, tl, tr.left.left), balance(tr.key, tr.value, tr.left.right, tr.right.red))
+         tr.left.redWithLeftRight(tree.blackWithLeftRight(tl, tr.left.left), balance(tr, tr.left.right, tr.right.red))
     else sys.error("Defect: invariance violation")
 
-  private[this] def balRight[A, B](x: A, xv: B, tl: Tree[A, B], tr: Tree[A, B]) =
-    if (isRedTree(tr)) RedTree(x, xv, tl, tr.black)
-    else if (isBlackTree(tl)) balance(x, xv, tl.red, tr)
+  private[this] def balRight[A, B](tree: Tree[A,B], tl: Tree[A, B], tr: Tree[A, B]): Tree[A, B] =
+    if (isRedTree(tr)) tree.redWithLeftRight(tl, tr.black)
+    else if (isBlackTree(tl)) balance(tree, tl.red, tr)
     else if (isRedTree(tl) && isBlackTree(tl.right))
-      RedTree(tl.right.key, tl.right.value, balance(tl.key, tl.value, tl.left.red, tl.right.left), BlackTree(x, xv, tl.right.right, tr))
+         tl.right.redWithLeftRight(balance(tl, tl.left.red, tl.right.left), tree.blackWithLeftRight(tl.right.right, tr))
     else sys.error("Defect: invariance violation")
 
   /** `append` is similar to `join2` but requires that both subtrees have the same black height */
-  private[this] def append[A, B](tl: Tree[A, B], tr: Tree[A, B]): Tree[A, B] =
+  private[this] def append[A, B](tl: Tree[A, B], tr: Tree[A, B]): Tree[A, B] = {
     if (tl eq null) tr
     else if (tr eq null) tl
-    else if (isRedTree(tl) && isRedTree(tr)) {
-      val bc = append(tl.right, tr.left)
-      if (isRedTree(bc)) {
-        RedTree(bc.key, bc.value, RedTree(tl.key, tl.value, tl.left, bc.left), RedTree(tr.key, tr.value, bc.right, tr.right))
-      } else {
-        RedTree(tl.key, tl.value, tl.left, RedTree(tr.key, tr.value, bc, tr.right))
-      }
-    } else if (isBlackTree(tl) && isBlackTree(tr)) {
-      val bc = append(tl.right, tr.left)
-      if (isRedTree(bc)) {
-        RedTree(bc.key, bc.value, BlackTree(tl.key, tl.value, tl.left, bc.left), BlackTree(tr.key, tr.value, bc.right, tr.right))
-      } else {
-        balLeft(tl.key, tl.value, tl.left, BlackTree(tr.key, tr.value, bc, tr.right))
-      }
-    } else if (isRedTree(tr)) RedTree(tr.key, tr.value, append(tl, tr.left), tr.right)
-    else if (isRedTree(tl)) RedTree(tl.key, tl.value, tl.left, append(tl.right, tr))
-    else sys.error("unmatched tree on append: " + tl + ", " + tr)
+    else if (tl.isRed) {
+           if (tr.isRed) {
+             //tl is red, tr is red
+             val bc = append(tl.right, tr.left)
+             if (isRedTree(bc)) bc.withLeftRight(tl.withRight(bc.left), tr.withLeft(bc.right))
+             else tl.withRight(tr.withLeft(bc))
+           } else {
+             //tl is red, tr is black
+             tl.withRight(append(tl.right, tr))
+           }
+         } else {
+           if (tr.isBlack) {
+             //tl is black tr is black
+             val bc = append(tl.right, tr.left)
+             if (isRedTree(bc)) bc.withLeftRight(tl.withRight(bc.left), tr.withLeft(bc.right))
+             else balLeft(tl, tl.left, tr.withLeft(bc))
+           } else {
+             //tl is black tr is red
+             tr.withLeft(append(tl, tr.left))
+           }
+         }
+  }
 
 
   // Bulk operations based on "Just Join for Parallel Ordered Sets" (https://www.cs.cmu.edu/~guyb/papers/BFS16.pdf)
@@ -1142,7 +1193,7 @@ private[collection] object NewRedBlackTree {
   /** Compute the rank from a tree and its black height */
   @`inline` private[this] def rank(t: Tree[_, _], bh: Int): Int = {
     if(t eq null) 0
-    else if(isBlackTree(t)) 2*(bh-1)
+    else if(t.isBlack) 2*(bh-1)
     else 2*bh-1
   }
 
@@ -1178,7 +1229,7 @@ private[collection] object NewRedBlackTree {
 
   private[this] def join[A, B](tl: Tree[A, B], k: A, v: B, tr: Tree[A, B]): Tree[A, B] = {
     @tailrec def h(t: Tree[_, _], i: Int): Int =
-      if(t eq null) i+1 else h(t.left, if(isBlackTree(t)) i+1 else i)
+      if(t eq null) i+1 else h(t.left, if(t.isBlack) i+1 else i)
     val bhtl = h(tl, 0)
     val bhtr = h(tr, 0)
     if(bhtl > bhtr) {