1 files changed, 580 insertions, 0 deletions
diff --git a/src/library/scala/collection/mutable/RedBlackTree.scala b/src/library/scala/collection/mutable/RedBlackTree.scala
new file mode 100644
index 0000000000..e4793242bf
--- /dev/null
+++ b/src/library/scala/collection/mutable/RedBlackTree.scala
@@ -0,0 +1,580 @@
+package scala.collection.mutable
+
+import scala.annotation.tailrec
+import scala.collection.Iterator
+
+/**
+ * An object containing the red-black tree implementation used by mutable `TreeMaps`.
+ *
+ * The trees implemented in this object are *not* thread safe.
+ *
+ * @author Rui Gonçalves
+ * @version 2.12
+ * @since 2.12
+ */
+private[collection] object RedBlackTree {
+
+  // ---- class structure ----
+
+  // For performance reasons, this implementation uses `null` references to represent leaves instead of a sentinel node.
+  // Currently, the internal nodes do not store their subtree size - only the tree object keeps track of their size.
+  // Therefore, while obtaining the size of the whole tree is O(1), knowing the number of entries inside a range is O(n)
+  // on the size of the range.
+
+  @SerialVersionUID(21575944040195605L)
+  final class Tree[A, B](var root: Node[A, B], var size: Int) extends Serializable
+
+  @SerialVersionUID(1950599696441054720L)
+  final class Node[A, B](var key: A, var value: B, var red: Boolean,
+                         var left: Node[A, B], var right: Node[A, B], var parent: Node[A, B]) extends Serializable {
+
+    override def toString: String = "Node(" + key + ", " + value + ", " + red + ", " + left + ", " + right + ")"
+  }
+
+  object Tree {
+    def empty[A, B]: Tree[A, B] = new Tree(null, 0)
+  }
+
+  object Node {
+
+    @inline def apply[A, B](key: A, value: B, red: Boolean,
+                            left: Node[A, B], right: Node[A, B], parent: Node[A, B]): Node[A, B] =
+      new Node(key, value, red, left, right, parent)
+
+    @inline def leaf[A, B](key: A, value: B, red: Boolean, parent: Node[A, B]): Node[A, B] =
+      new Node(key, value, red, null, null, parent)
+
+    def unapply[A, B](t: Node[A, B]) = Some((t.key, t.value, t.left, t.right, t.parent))
+  }
+
+  // ---- getters ----
+
+  def isRed(node: Node[_, _]) = (node ne null) && node.red
+  def isBlack(node: Node[_, _]) = (node eq null) || !node.red
+
+  // ---- size ----
+
+  def size(node: Node[_, _]): Int = if (node eq null) 0 else 1 + size(node.left) + size(node.right)
+  def size(tree: Tree[_, _]): Int = tree.size
+  def isEmpty(tree: Tree[_, _]) = tree.root eq null
+  def clear(tree: Tree[_, _]): Unit = { tree.root = null; tree.size = 0 }
+
+  // ---- search ----
+
+  def get[A: Ordering, B](tree: Tree[A, B], key: A): Option[B] = getNode(tree.root, key) match {
+    case null => None
+    case node => Some(node.value)
+  }
+
+  @tailrec private[this] def getNode[A, B](node: Node[A, B], key: A)(implicit ord: Ordering[A]): Node[A, B] =
+    if (node eq null) null
+    else {
+      val cmp = ord.compare(key, node.key)
+      if (cmp < 0) getNode(node.left, key)
+      else if (cmp > 0) getNode(node.right, key)
+      else node
+    }
+
+  def contains[A: Ordering](tree: Tree[A, _], key: A) = getNode(tree.root, key) ne null
+
+  def min[A, B](tree: Tree[A, B]): Option[(A, B)] = minNode(tree.root) match {
+    case null => None
+    case node => Some((node.key, node.value))
+  }
+
+  def minKey[A](tree: Tree[A, _]): Option[A] = minNode(tree.root) match {
+    case null => None
+    case node => Some(node.key)
+  }
+
+  private def minNode[A, B](node: Node[A, B]): Node[A, B] =
+    if (node eq null) null else minNodeNonNull(node)
+
+  @tailrec def minNodeNonNull[A, B](node: Node[A, B]): Node[A, B] =
+    if (node.left eq null) node else minNodeNonNull(node.left)
+
+  def max[A, B](tree: Tree[A, B]): Option[(A, B)] = maxNode(tree.root) match {
+    case null => None
+    case node => Some((node.key, node.value))
+  }
+
+  def maxKey[A](tree: Tree[A, _]): Option[A] = maxNode(tree.root) match {
+    case null => None
+    case node => Some(node.key)
+  }
+
+  private def maxNode[A, B](node: Node[A, B]): Node[A, B] =
+    if (node eq null) null else maxNodeNonNull(node)
+
+  @tailrec def maxNodeNonNull[A, B](node: Node[A, B]): Node[A, B] =
+    if (node.right eq null) node else maxNodeNonNull(node.right)
+
+  /**
+   * Returns the first (lowest) map entry with a key equal or greater than `key`. Returns `None` if there is no such
+   * node.
+   */
+  def minAfter[A, B](tree: Tree[A, B], key: A)(implicit ord: Ordering[A]): Option[(A, B)] =
+    minNodeAfter(tree.root, key) match {
+      case null => None
+      case node => Some((node.key, node.value))
+    }
+
+  def minKeyAfter[A](tree: Tree[A, _], key: A)(implicit ord: Ordering[A]): Option[A] =
+    minNodeAfter(tree.root, key) match {
+      case null => None
+      case node => Some(node.key)
+    }
+
+  private[this] def minNodeAfter[A, B](node: Node[A, B], key: A)(implicit ord: Ordering[A]): Node[A, B] = {
+    if (node eq null) null
+    else {
+      var y: Node[A, B] = null
+      var x = node
+      var cmp = 1
+      while ((x ne null) && cmp != 0) {
+        y = x
+        cmp = ord.compare(key, x.key)
+        x = if (cmp < 0) x.left else x.right
+      }
+      if (cmp <= 0) y else successor(y)
+    }
+  }
+
+  /**
+   * Returns the last (highest) map entry with a key smaller than `key`. Returns `None` if there is no such node.
+   */
+  def maxBefore[A, B](tree: Tree[A, B], key: A)(implicit ord: Ordering[A]): Option[(A, B)] =
+    maxNodeBefore(tree.root, key) match {
+      case null => None
+      case node => Some((node.key, node.value))
+    }
+
+  def maxKeyBefore[A](tree: Tree[A, _], key: A)(implicit ord: Ordering[A]): Option[A] =
+    maxNodeBefore(tree.root, key) match {
+      case null => None
+      case node => Some(node.key)
+    }
+
+  private[this] def maxNodeBefore[A, B](node: Node[A, B], key: A)(implicit ord: Ordering[A]): Node[A, B] = {
+    if (node eq null) null
+    else {
+      var y: Node[A, B] = null
+      var x = node
+      var cmp = 1
+      while ((x ne null) && cmp != 0) {
+        y = x
+        cmp = ord.compare(key, x.key)
+        x = if (cmp < 0) x.left else x.right
+      }
+      if (cmp > 0) y else predecessor(y)
+    }
+  }
+
+  // ---- insertion ----
+
+  def insert[A, B](tree: Tree[A, B], key: A, value: B)(implicit ord: Ordering[A]): Unit = {
+    var y: Node[A, B] = null
+    var x = tree.root
+    var cmp = 1
+    while ((x ne null) && cmp != 0) {
+      y = x
+      cmp = ord.compare(key, x.key)
+      x = if (cmp < 0) x.left else x.right
+    }
+
+    if (cmp == 0) y.value = value
+    else {
+      val z = Node.leaf(key, value, red = true, y)
+
+      if (y eq null) tree.root = z
+      else if (cmp < 0) y.left = z
+      else y.right = z
+
+      fixAfterInsert(tree, z)
+      tree.size += 1
+    }
+  }
+
+  private[this] def fixAfterInsert[A, B](tree: Tree[A, B], node: Node[A, B]): Unit = {
+    var z = node
+    while (isRed(z.parent)) {
+      if (z.parent eq z.parent.parent.left) {
+        val y = z.parent.parent.right
+        if (isRed(y)) {
+          z.parent.red = false
+          y.red = false
+          z.parent.parent.red = true
+          z = z.parent.parent
+        } else {
+          if (z eq z.parent.right) {
+            z = z.parent
+            rotateLeft(tree, z)
+          }
+          z.parent.red = false
+          z.parent.parent.red = true
+          rotateRight(tree, z.parent.parent)
+        }
+      } else { // symmetric cases
+        val y = z.parent.parent.left
+        if (isRed(y)) {
+          z.parent.red = false
+          y.red = false
+          z.parent.parent.red = true
+          z = z.parent.parent
+        } else {
+          if (z eq z.parent.left) {
+            z = z.parent
+            rotateRight(tree, z)
+          }
+          z.parent.red = false
+          z.parent.parent.red = true
+          rotateLeft(tree, z.parent.parent)
+        }
+      }
+    }
+    tree.root.red = false
+  }
+
+  // ---- deletion ----
+
+  def delete[A, B](tree: Tree[A, B], key: A)(implicit ord: Ordering[A]): Unit = {
+    val z = getNode(tree.root, key)
+    if (z ne null) {
+      var y = z
+      var yIsRed = y.red
+      var x: Node[A, B] = null
+      var xParent: Node[A, B] = null
+
+      if (z.left eq null) {
+        x = z.right
+        transplant(tree, z, z.right)
+        xParent = z.parent
+      }
+      else if (z.right eq null) {
+        x = z.left
+        transplant(tree, z, z.left)
+        xParent = z.parent
+      }
+      else {
+        y = minNodeNonNull(z.right)
+        yIsRed = y.red
+        x = y.right
+
+        if (y.parent eq z) xParent = y
+        else {
+          xParent = y.parent
+          transplant(tree, y, y.right)
+          y.right = z.right
+          y.right.parent = y
+        }
+        transplant(tree, z, y)
+        y.left = z.left
+        y.left.parent = y
+        y.red = z.red
+      }
+
+      if (!yIsRed) fixAfterDelete(tree, x, xParent)
+      tree.size -= 1
+    }
+  }
+
+  private[this] def fixAfterDelete[A, B](tree: Tree[A, B], node: Node[A, B], parent: Node[A, B]): Unit = {
+    var x = node
+    var xParent = parent
+    while ((x ne tree.root) && isBlack(x)) {
+      if (x eq xParent.left) {
+        var w = xParent.right
+        // assert(w ne null)
+
+        if (w.red) {
+          w.red = false
+          xParent.red = true
+          rotateLeft(tree, xParent)
+          w = xParent.right
+        }
+        if (isBlack(w.left) && isBlack(w.right)) {
+          w.red = true
+          x = xParent
+        } else {
+          if (isBlack(w.right)) {
+            w.left.red = false
+            w.red = true
+            rotateRight(tree, w)
+            w = xParent.right
+          }
+          w.red = xParent.red
+          xParent.red = false
+          w.right.red = false
+          rotateLeft(tree, xParent)
+          x = tree.root
+        }
+      } else { // symmetric cases
+        var w = xParent.left
+        // assert(w ne null)
+
+        if (w.red) {
+          w.red = false
+          xParent.red = true
+          rotateRight(tree, xParent)
+          w = xParent.left
+        }
+        if (isBlack(w.right) && isBlack(w.left)) {
+          w.red = true
+          x = xParent
+        } else {
+          if (isBlack(w.left)) {
+            w.right.red = false
+            w.red = true
+            rotateLeft(tree, w)
+            w = xParent.left
+          }
+          w.red = xParent.red
+          xParent.red = false
+          w.left.red = false
+          rotateRight(tree, xParent)
+          x = tree.root
+        }
+      }
+      xParent = x.parent
+    }
+    if (x ne null) x.red = false
+  }
+
+  // ---- helpers ----
+
+  /**
+   * Returns the node that follows `node` in an in-order tree traversal. If `node` has the maximum key (and is,
+   * therefore, the last node), this method returns `null`.
+   */
+  private[this] def successor[A, B](node: Node[A, B]): Node[A, B] = {
+    if (node.right ne null) minNodeNonNull(node.right)
+    else {
+      var x = node
+      var y = x.parent
+      while ((y ne null) && (x eq y.right)) {
+        x = y
+        y = y.parent
+      }
+      y
+    }
+  }
+
+  /**
+   * Returns the node that precedes `node` in an in-order tree traversal. If `node` has the minimum key (and is,
+   * therefore, the first node), this method returns `null`.
+   */
+  private[this] def predecessor[A, B](node: Node[A, B]): Node[A, B] = {
+    if (node.left ne null) maxNodeNonNull(node.left)
+    else {
+      var x = node
+      var y = x.parent
+      while ((y ne null) && (x eq y.left)) {
+        x = y
+        y = y.parent
+      }
+      y
+    }
+  }
+
+  private[this] def rotateLeft[A, B](tree: Tree[A, B], x: Node[A, B]): Unit = if (x ne null) {
+    // assert(x.right ne null)
+    val y = x.right
+    x.right = y.left
+
+    if (y.left ne null) y.left.parent = x
+    y.parent = x.parent
+
+    if (x.parent eq null) tree.root = y
+    else if (x eq x.parent.left) x.parent.left = y
+    else x.parent.right = y
+
+    y.left = x
+    x.parent = y
+  }
+
+  private[this] def rotateRight[A, B](tree: Tree[A, B], x: Node[A, B]): Unit = if (x ne null) {
+    // assert(x.left ne null)
+    val y = x.left
+    x.left = y.right
+
+    if (y.right ne null) y.right.parent = x
+    y.parent = x.parent
+
+    if (x.parent eq null) tree.root = y
+    else if (x eq x.parent.right) x.parent.right = y
+    else x.parent.left = y
+
+    y.right = x
+    x.parent = y
+  }
+
+  /**
+   * Transplant the node `from` to the place of node `to`. This is done by setting `from` as a child of `to`'s previous
+   * parent and setting `from`'s parent to the `to`'s previous parent. The children of `from` are left unchanged.
+   */
+  private[this] def transplant[A, B](tree: Tree[A, B], to: Node[A, B], from: Node[A, B]): Unit = {
+    if (to.parent eq null) tree.root = from
+    else if (to eq to.parent.left) to.parent.left = from
+    else to.parent.right = from
+
+    if (from ne null) from.parent = to.parent
+  }
+
+  // ---- tree traversal ----
+
+  def foreach[A, B, U](tree: Tree[A, B], f: ((A, B)) => U): Unit = foreachNode(tree.root, f)
+
+  private[this] def foreachNode[A, B, U](node: Node[A, B], f: ((A, B)) => U): Unit =
+    if (node ne null) foreachNodeNonNull(node, f)
+
+  private[this] def foreachNodeNonNull[A, B, U](node: Node[A, B], f: ((A, B)) => U): Unit = {
+    if (node.left ne null) foreachNodeNonNull(node.left, f)
+    f((node.key, node.value))
+    if (node.right ne null) foreachNodeNonNull(node.right, f)
+  }
+
+  def foreachKey[A, U](tree: Tree[A, _], f: A => U): Unit = foreachNodeKey(tree.root, f)
+
+  private[this] def foreachNodeKey[A, U](node: Node[A, _], f: A => U): Unit =
+    if (node ne null) foreachNodeKeyNonNull(node, f)
+
+  private[this] def foreachNodeKeyNonNull[A, U](node: Node[A, _], f: A => U): Unit = {
+    if (node.left ne null) foreachNodeKeyNonNull(node.left, f)
+    f(node.key)
+    if (node.right ne null) foreachNodeKeyNonNull(node.right, f)
+  }
+
+  def transform[A, B](tree: Tree[A, B], f: (A, B) => B): Unit = transformNode(tree.root, f)
+
+  private[this] def transformNode[A, B, U](node: Node[A, B], f: (A, B) => B): Unit =
+    if (node ne null) transformNodeNonNull(node, f)
+
+  private[this] def transformNodeNonNull[A, B, U](node: Node[A, B], f: (A, B) => B): Unit = {
+    if (node.left ne null) transformNodeNonNull(node.left, f)
+    node.value = f(node.key, node.value)
+    if (node.right ne null) transformNodeNonNull(node.right, f)
+  }
+
+  def iterator[A: Ordering, B](tree: Tree[A, B], start: Option[A] = None, end: Option[A] = None): Iterator[(A, B)] =
+    new EntriesIterator(tree, start, end)
+
+  def keysIterator[A: Ordering](tree: Tree[A, _], start: Option[A] = None, end: Option[A] = None): Iterator[A] =
+    new KeysIterator(tree, start, end)
+
+  def valuesIterator[A: Ordering, B](tree: Tree[A, B], start: Option[A] = None, end: Option[A] = None): Iterator[B] =
+    new ValuesIterator(tree, start, end)
+
+  private[this] abstract class TreeIterator[A, B, R](tree: Tree[A, B], start: Option[A], end: Option[A])
+                                                    (implicit ord: Ordering[A]) extends Iterator[R] {
+
+    protected[this] def nextResult(node: Node[A, B]): R
+
+    def hasNext: Boolean = nextNode ne null
+
+    def next(): R = nextNode match {
+      case null => throw new NoSuchElementException("next on empty iterator")
+      case node =>
+        nextNode = successor(node)
+        setNullIfAfterEnd()
+        nextResult(node)
+    }
+
+    private[this] var nextNode: Node[A, B] = start match {
+      case None => minNode(tree.root)
+      case Some(from) => minNodeAfter(tree.root, from)
+    }
+
+    private[this] def setNullIfAfterEnd(): Unit =
+      if (end.isDefined && (nextNode ne null) && ord.compare(nextNode.key, end.get) >= 0)
+        nextNode = null
+
+    setNullIfAfterEnd()
+  }
+
+  private[this] final class EntriesIterator[A: Ordering, B](tree: Tree[A, B], start: Option[A], end: Option[A])
+    extends TreeIterator[A, B, (A, B)](tree, start, end) {
+
+    def nextResult(node: Node[A, B]) = (node.key, node.value)
+  }
+
+  private[this] final class KeysIterator[A: Ordering, B](tree: Tree[A, B], start: Option[A], end: Option[A])
+    extends TreeIterator[A, B, A](tree, start, end) {
+
+    def nextResult(node: Node[A, B]) = node.key
+  }
+
+  private[this] final class ValuesIterator[A: Ordering, B](tree: Tree[A, B], start: Option[A], end: Option[A])
+    extends TreeIterator[A, B, B](tree, start, end) {
+
+    def nextResult(node: Node[A, B]) = node.value
+  }
+
+  // ---- debugging ----
+
+  /**
+   * Checks if the tree is in a valid state. That happens if:
+   * - It is a valid binary search tree;
+   * - All red-black properties are satisfied;
+   * - All non-null nodes have their `parent` reference correct;
+   * - The size variable in `tree` corresponds to the actual size of the tree.
+   */
+  def isValid[A: Ordering, B](tree: Tree[A, B]): Boolean =
+    isValidBST(tree.root) && hasProperParentRefs(tree) && isValidRedBlackTree(tree) && size(tree.root) == tree.size
+
+  /**
+   * Returns true if all non-null nodes have their `parent` reference correct.
+   */
+  private[this] def hasProperParentRefs[A, B](tree: Tree[A, B]): Boolean = {
+
+    def hasProperParentRefs(node: Node[A, B]): Boolean = {
+      if (node eq null) true
+      else {
+        if ((node.left ne null) && (node.left.parent ne node) ||
+          (node.right ne null) && (node.right.parent ne node)) false
+        else hasProperParentRefs(node.left) && hasProperParentRefs(node.right)
+      }
+    }
+
+    if(tree.root eq null) true
+    else (tree.root.parent eq null) && hasProperParentRefs(tree.root)
+  }
+
+  /**
+   * Returns true if this node follows the properties of a binary search tree.
+   */
+  private[this] def isValidBST[A, B](node: Node[A, B])(implicit ord: Ordering[A]): Boolean = {
+    if (node eq null) true
+    else {
+      if ((node.left ne null) && (ord.compare(node.key, node.left.key) <= 0) ||
+        (node.right ne null) && (ord.compare(node.key, node.right.key) >= 0)) false
+      else isValidBST(node.left) && isValidBST(node.right)
+    }
+  }
+
+  /**
+   * Returns true if the tree has all the red-black tree properties: if the root node is black, if all children of red
+   * nodes are black and if the path from any node to any of its null children has the same number of black nodes.
+   */
+  private[this] def isValidRedBlackTree[A, B](tree: Tree[A, B]): Boolean = {
+
+    def noRedAfterRed(node: Node[A, B]): Boolean = {
+      if (node eq null) true
+      else if (node.red && (isRed(node.left) || isRed(node.right))) false
+      else noRedAfterRed(node.left) && noRedAfterRed(node.right)
+    }
+
+    def blackHeight(node: Node[A, B]): Int = {
+      if (node eq null) 1
+      else {
+        val lh = blackHeight(node.left)
+        val rh = blackHeight(node.right)
+
+        if (lh == -1 || lh != rh) -1
+        else if (isRed(node)) lh
+        else lh + 1
+      }
+    }
+
+    isBlack(tree.root) && noRedAfterRed(tree.root) && blackHeight(tree.root) >= 0
+  }
+}