Merge pull request scala/scala#9315 from mkeskells/2.12.x_RedBlack_reuse_fp

reuse common code for filter/partition for RedBlackTree based Set/Map

Author: dwijnand

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

@@ -1021,41 +1021,48 @@ private[collection] object NewRedBlackTree {
     blacken(fk(t))
   }
 
-  def partitionEntries[A, B](t: Tree[A, B], p: (A, B) => Boolean): (Tree[A, B], Tree[A, B]) = if(t eq null) (null, null) else {
-    var tmpk, tmpd = null: Tree[A, B] // shared vars to avoid returning tuples from fk
-    def fk(t: Tree[A, B]): Unit = {
-      val k = t.key
-      val v = t.value
-      val l = t.left
-      val r = t.right
-      var l2k, l2d, r2k, r2d = null: Tree[A, B]
-      if(l ne null) {
-        fk(l)
-        l2k = tmpk
-        l2d = tmpd
-      }
-      val keep = p(k, v)
-      if(r ne null) {
-        fk(r)
-        r2k = tmpk
-        r2d = tmpd
+  private[this] val null2 = (null, null)
+
+  def partitionEntries[A, B](t: Tree[A, B], p: (A, B) => Boolean): (Tree[A, B], Tree[A, B]) = {
+    if (t eq null) null2
+    else {
+      object partitioner {
+        var tmpk, tmpd = null: Tree[A, B] // shared vars to avoid returning tuples from fk
+        def fk(t: Tree[A, B]): Unit = {
+          val k                  = t.key
+          val v                  = t.value
+          val l                  = t.left
+          val r                  = t.right
+          var l2k, l2d, r2k, r2d = null: Tree[A, B]
+          if (l ne null) {
+            fk(l)
+            l2k = tmpk
+            l2d = tmpd
+          }
+          val keep = p(k, v)
+          if (r ne null) {
+            fk(r)
+            r2k = tmpk
+            r2d = tmpd
+          }
+          val jk =
+            if (!keep) join2(l2k, r2k)
+            else if ((l2k eq l) && (r2k eq r)) t
+                 else join(l2k, k, v, r2k)
+          val jd =
+            if (keep) join2(l2d, r2d)
+            else if ((l2d eq l) && (r2d eq r)) t
+                 else join(l2d, k, v, r2d)
+          tmpk = jk
+          tmpd = jd
+        }
       }
-      val jk =
-        if(!keep) join2(l2k, r2k)
-        else if((l2k eq l) && (r2k eq r)) t
-        else join(l2k, k, v, r2k)
-      val jd =
-        if(keep) join2(l2d, r2d)
-        else if((l2d eq l) && (r2d eq r)) t
-        else join(l2d, k, v, r2d)
-      tmpk = jk
-      tmpd = jd
+
+      partitioner.fk(t)
+      (blacken(partitioner.tmpk), blacken(partitioner.tmpd))
     }
-    fk(t)
-    (blacken(tmpk), blacken(tmpd))
   }
 
-
   // Based on Stefan Kahrs' Haskell version of Okasaki's Red&Black Trees
   // Constructing Red-Black Trees, Ralf Hinze: [[http://www.cs.ox.ac.uk/ralf.hinze/publications/WAAAPL99b.ps.gz]]
   // Red-Black Trees in a Functional Setting, Chris Okasaki: [[https://wiki.rice.edu/confluence/download/attachments/2761212/Okasaki-Red-Black.pdf]] */
@@ -1253,4 +1260,4 @@ private[collection] object NewRedBlackTree {
       val tr = _difference(r1, t2.right)
       join2(tl, tr)
     }
-}
+}