Build red-black trees from ordered sets in linear time

- Optimize building TreeMap from SortedMap
- Optimize building TreeSet from SortedSet and Range
- Static value for Ordering.Int.reverse so we can check for it with
  eq and avoid allocating a new instance for each check
- Bug fix: Check for ordering when reusing existing TreeMap/TreeSet
- Compare Orderings with == instead of eq

Author: szeiger

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

@@ -572,6 +572,40 @@ private[collection] object RedBlackTree {
     override def nextResult(tree: Tree[A, B]) = tree.value
   }
 
+  /** Build a Tree suitable for a TreeSet from an ordered sequence of keys */
+  def fromOrderedKeys[A](xs: Iterator[A], size: Int): Tree[A, Null] = {
+    val maxUsedDepth = 32 - Integer.numberOfLeadingZeros(size) // maximum depth of non-leaf nodes
+    def f(level: Int, size: Int): Tree[A, Null] = size match {
+      case 0 => null
+      case 1 => mkTree(level != maxUsedDepth || level == 1, xs.next(), null, null, null)
+      case n =>
+        val leftSize = (size-1)/2
+        val left = f(level+1, leftSize)
+        val x = xs.next()
+        val right = f(level+1, size-1-leftSize)
+        BlackTree(x, null, left, right)
+    }
+    f(1, size)
+  }
+
+  /** Build a Tree suitable for a TreeMap from an ordered sequence of key/value pairs */
+  def fromOrderedEntries[A, B](xs: Iterator[(A, B)], size: Int): Tree[A, B] = {
+    val maxUsedDepth = 32 - Integer.numberOfLeadingZeros(size) // maximum depth of non-leaf nodes
+    def f(level: Int, size: Int): Tree[A, B] = size match {
+      case 0 => null
+      case 1 =>
+        val (k, v) = xs.next()
+        mkTree(level != maxUsedDepth || level == 1, k, v, null, null)
+      case n =>
+        val leftSize = (size-1)/2
+        val left = f(level+1, leftSize)
+        val (k, v) = xs.next()
+        val right = f(level+1, size-1-leftSize)
+        BlackTree(k, v, left, right)
+    }
+    f(1, size)
+  }
+
   // Bulk operations based on "Just Join for Parallel Ordered Sets" (https://www.cs.cmu.edu/~guyb/papers/BFS16.pdf)
   // We don't store the black height in the tree so we pass it down into the join methods and derive the black height
   // of child nodes from it. Where possible the black height is used directly instead of deriving the rank from it.

library/src/scala/collection/immutable/TreeMap.scala

@@ -59,7 +59,7 @@ final class TreeMap[K, +V] private (private val tree: RB.Tree[K, V])(implicit va
 
   override def concat[V1 >: V](that: collection.IterableOnce[(K, V1)]): TreeMap[K, V1] = {
     val t = that match {
-      case tm: TreeMap[K, V] if ordering eq tm.ordering =>
+      case tm: TreeMap[K, V] if ordering == tm.ordering =>
         RB.union(tree, tm.tree)
       case _ =>
         val it = that.iterator
@@ -75,7 +75,7 @@ final class TreeMap[K, +V] private (private val tree: RB.Tree[K, V])(implicit va
   }
 
   override def removeAll(keys: IterableOnce[K]): TreeMap[K, V] = keys match {
-    case ts: TreeSet[K] if ordering eq ts.ordering =>
+    case ts: TreeSet[K] if ordering == ts.ordering =>
       val t = RB.difference(tree, ts.tree)
       if(t eq tree) this else new TreeMap(t)
     case _ => super.removeAll(keys)
@@ -182,9 +182,11 @@ object TreeMap extends SortedMapFactory[TreeMap] {
 
   def empty[K : Ordering, V]: TreeMap[K, V] = new TreeMap()
 
-  def from[K : Ordering, V](it: IterableOnce[(K, V)]): TreeMap[K, V] =
+  def from[K, V](it: IterableOnce[(K, V)])(implicit ordering: Ordering[K]): TreeMap[K, V] =
     it match {
-      case tm: TreeMap[K, V] => tm
+      case tm: TreeMap[K, V] if ordering == tm.ordering => tm
+      case sm: scala.collection.SortedMap[K, V] if ordering == sm.ordering =>
+        new TreeMap[K, V](RB.fromOrderedEntries(sm.iterator, sm.size))
       case _ =>
         var t: RB.Tree[K, V] = null
         val i = it.iterator
@@ -200,7 +202,7 @@ object TreeMap extends SortedMapFactory[TreeMap] {
     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 =>
+        case tm: TreeMap[K, V] if ordering == tm.ordering =>
           tree = RB.union(tree, tm.tree)
         case _ =>
           val it = xs.iterator

library/src/scala/collection/immutable/TreeSet.scala

@@ -134,7 +134,7 @@ final class TreeSet[A] private (private[immutable] val tree: RB.Tree[A, Null])(i
 
   override def concat(that: collection.IterableOnce[A]): TreeSet[A] = {
     val t = that match {
-      case ts: TreeSet[A] if ordering eq ts.ordering =>
+      case ts: TreeSet[A] if ordering == ts.ordering =>
         RB.union(tree, ts.tree)
       case _ =>
         val it = that.iterator
@@ -146,15 +146,15 @@ final class TreeSet[A] private (private[immutable] val tree: RB.Tree[A, Null])(i
   }
 
   override def intersect(that: collection.Set[A]): TreeSet[A] = that match {
-    case ts: TreeSet[A] if ordering eq ts.ordering =>
+    case ts: TreeSet[A] if ordering == 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 =>
+    case ts: TreeSet[A] if ordering == ts.ordering =>
       val t = RB.difference(tree, ts.tree)
       if(t eq tree) this else newSet(t)
     case _ =>
@@ -175,9 +175,14 @@ object TreeSet extends SortedIterableFactory[TreeSet] {
 
   def empty[A: Ordering]: TreeSet[A] = new TreeSet[A]
 
-  def from[E: Ordering](it: scala.collection.IterableOnce[E]): TreeSet[E] =
+  def from[E](it: scala.collection.IterableOnce[E])(implicit ordering: Ordering[E]): TreeSet[E] =
     it match {
-      case ts: TreeSet[E] => ts
+      case ts: TreeSet[E] if ordering == ts.ordering => ts
+      case ss: scala.collection.SortedSet[E] if ordering == ss.ordering =>
+        new TreeSet[E](RB.fromOrderedKeys(ss.iterator, ss.size))
+      case r: Range if (ordering eq Ordering.Int) || (ordering eq Ordering.Int.reverse) =>
+        val it = if((ordering eq Ordering.Int) == (r.step > 0)) r.iterator else r.reverseIterator
+        new TreeSet[E](RB.fromOrderedKeys(it, r.size))
       case _ =>
         var t: RB.Tree[E, Null] = null
         val i = it.iterator
@@ -190,7 +195,7 @@ object TreeSet extends SortedIterableFactory[TreeSet] {
     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 =>
+        case ts: TreeSet[A] if ordering == ts.ordering =>
           tree = RB.union(tree, ts.tree)
         case _ =>
           val it = xs.iterator

library/src/scala/math/Ordering.scala

@@ -230,6 +230,8 @@ object Ordering extends LowPriorityOrderingImplicits {
     override def hashCode(): Int = outer.hashCode() * reverseSeed
   }
 
+  private final val IntReverse: Ordering[Int] = new Reverse(Ordering.Int)
+
   private final class IterableOrdering[CC[X] <: Iterable[X], T](private val ord: Ordering[T]) extends Ordering[CC[T]] {
     def compare(x: CC[T], y: CC[T]): Int = {
       val xe = x.iterator
@@ -327,6 +329,7 @@ object Ordering extends LowPriorityOrderingImplicits {
 
   trait IntOrdering extends Ordering[Int] {
     def compare(x: Int, y: Int) = java.lang.Integer.compare(x, y)
+    override def reverse: Ordering[Int] = Ordering.IntReverse
   }
   implicit object Int extends IntOrdering