less allocation on deserialisation of RedBlackTree (SortedMap)

reuse deserialisation logic for map and set

Author: mkeskells

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

@@ -955,34 +955,20 @@ private[collection] object NewRedBlackTree {
     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] = {
+  /** Build a Tree from an ordered sequence of key and values */
+  def fromOrderedEntries[A, B](keys: Iterator[A], values: Iterator[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()
+        val k = keys.next()
+        val v = values.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 k = keys.next()
+        val v = values.next()
         val right = f(level+1, size-1-leftSize)
         BlackTree(k, v, left, right)
     }

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

@@ -106,16 +106,26 @@ object TreeMap extends ImmutableSortedMapFactory[TreeMap] {
     private[this] def readObject(in: java.io.ObjectInputStream) = {
       val size = in.readInt()
       ordering = in.readObject().asInstanceOf[Ordering[A]]
-      implicit val ord = ordering
-
-      val data = Array.newBuilder[(A, B)]
-      data.sizeHint(size)
-      for (i <- 0 until size) {
-        val key = in.readObject().asInstanceOf[A]
-        val value = in.readObject().asInstanceOf[B]
-        data += ((key, value))
+      size match {
+        case 0 => //tree is null already
+        case 1 =>
+          val key   = in.readObject().asInstanceOf[A]
+          val value = in.readObject().asInstanceOf[B]
+          tree = RB.update(null, key, value, true)(ordering)
+        case _ =>
+          val keys   = new Array[Any](size)
+          val values = new Array[Any](size)
+          var i      = 0
+          while (i < size) {
+            keys(i) = in.readObject()
+            values(i) = in.readObject()
+            i += 1
+          }
+          tree = RB.fromOrderedEntries(
+            keys.iterator.asInstanceOf[Iterator[A]],
+            values.iterator.asInstanceOf[Iterator[B]],
+            size)
       }
-      tree = RB.fromOrderedEntries(data.result.iterator, size)
     }
     @throws[IOException]
     private[this] def readResolve(): AnyRef =

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

@@ -85,16 +85,32 @@ object TreeSet extends ImmutableSortedSetFactory[TreeSet] {
     private[this] def readObject(in: java.io.ObjectInputStream) = {
       val size = in.readInt()
       ordering = in.readObject().asInstanceOf[Ordering[A]]
-      val data = Iterable.newBuilder[A]
-      data.sizeHint(size)
-      for (i <- 0 until size)
-        data += in.readObject().asInstanceOf[A]
-      tree = RB.fromOrderedKeys(data.result.iterator, size)
+      size match {
+        case 0 => //tree is null already
+        case 1 =>
+          val entry = in.readObject().asInstanceOf[A]
+          tree = RB.update(null, entry, (), false)(ordering)
+        case _ =>
+          val entries = new Array[Any](size)
+          var i       = 0
+          while (i < size) {
+            entries(i) = in.readObject()
+            i += 1
+          }
+          tree = RB.fromOrderedEntries(
+            entries.iterator.asInstanceOf[Iterator[A]],
+            unitsIterator,
+            size)
+      }
     }
     @throws[IOException]
     private[this] def readResolve(): AnyRef =
       new TreeSet(tree)(ordering)
   }
+  private[this] object unitsIterator extends AbstractIterator[Unit] {
+    override def hasNext: Boolean = true
+    override def next(): Unit = ()
+  }
 }
 
 /** This class implements immutable sets using a tree.