fix sameCBF in TreeSet and TreeMap

Author: mkeskells

library/src/scala/collection/generic/SortedMapFactory.scala

@@ -33,6 +33,7 @@ abstract class SortedMapFactory[CC[A, B] <: SortedMap[A, B] with SortedMapLike[A
     new MapBuilder[A, B, CC[A, B]](empty(ord))
 
   class SortedMapCanBuildFrom[A, B](implicit ord: Ordering[A]) extends CanBuildFrom[Coll, (A, B), CC[A, B]] {
+    private[collection] def factory = SortedMapFactory.this
     private[collection] def ordering = ord
     def apply(from: Coll) = newBuilder[A, B](ord)
     def apply() = newBuilder[A, B]

library/src/scala/collection/generic/SortedSetFactory.scala

@@ -33,6 +33,8 @@ abstract class SortedSetFactory[CC[A] <: SortedSet[A] with SortedSetLike[A, CC[A
   implicit def newCanBuildFrom[A](implicit ord : Ordering[A]) : CanBuildFrom[Coll, A, CC[A]] = new SortedSetCanBuildFrom()(ord)
 
   class SortedSetCanBuildFrom[A](implicit ord: Ordering[A]) extends CanBuildFrom[Coll, A, CC[A]] {
+    private[collection] def factory = SortedSetFactory.this
+    private[collection] def ordering = ord
     def apply(from: Coll) = newBuilder[A](ord)
     def apply() = newBuilder[A](ord)
   }

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

@@ -267,14 +267,22 @@ final class TreeMap[A, +B] private (tree: RB.Tree[A, B])(implicit val ordering:
     val (l, r) = RB.partitionEntries[A, B](tree, (k, v) => p((k, v)))
     (newMapOrSelf(l), newMapOrSelf(r))
   }
-
-  override def transform[W, That](f: (A, B) => W)(implicit bf: CanBuildFrom[TreeMap[A, B], (A, W), That]): That = {
+  private def sameCBF(bf: CanBuildFrom[_,_,_]): Boolean = {
     bf match {
-      case sm: TreeMap.SortedMapCanBuildFrom[A, B] if sm.ordering == ordering =>
-        val t2 = RB.transform[A, B, W](tree, f)
-        if (t2 eq tree) this.asInstanceOf[That]
-        else new TreeMap(t2).asInstanceOf[That]
-      case _ => super.transform(f)
+      case cbf: SortedMapFactory[_]#SortedMapCanBuildFrom[_,_] => {
+        val factory:AnyRef = cbf.factory
+        ((factory eq TreeMap) || (factory eq immutable.SortedMap) || (factory eq collection.SortedMap)) &&
+          cbf.ordering == ordering
+      }
+      case w: WrappedCanBuildFrom[_,_,_] => sameCBF(w.wrapped)
+      case _ => false
     }
   }
+
+
+  override def transform[W, That](f: (A, B) => W)(implicit bf: CanBuildFrom[TreeMap[A, B], (A, W), That]): That = {
+    if (sameCBF(bf))
+      newMapOrSelf(RB.transform[A, B, W](tree, f)).asInstanceOf[That]
+    else super.transform(f)
+  }
 }

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

@@ -205,10 +205,13 @@ final class TreeSet[A] private[immutable] (private[immutable] val tree: RB.Tree[
   override def firstKey = head
   override def lastKey = last
 
-
   private def sameCBF(bf: CanBuildFrom[_,_,_]): Boolean = {
     bf match {
-      case tsb: TreeSet.TreeSetBuilder[_] if tsb.ordering == ordering => true
+      case cbf: SortedSetFactory[_]#SortedSetCanBuildFrom[_] => {
+        val factory:AnyRef = cbf.factory
+        ((factory eq TreeSet) || (factory eq immutable.SortedSet) || (factory eq collection.SortedSet)) &&
+          cbf.ordering == ordering
+      }
       case w: WrappedCanBuildFrom[_,_,_] => sameCBF(w.wrapped)
       case _ => false
     }
@@ -221,7 +224,7 @@ final class TreeSet[A] private[immutable] (private[immutable] val tree: RB.Tree[
 
   private[scala] def addAllImpl[B >: A, That](that: GenTraversableOnce[B])(implicit bf: CanBuildFrom[TreeSet[A], B, That]): That = {
     that match {
-      case ts: TreeSet[A] if ordering == ts.ordering && sameCBF(bf) =>
+      case ts: TreeSet[A] if sameCBF(bf) =>
         newSetOrSelf(RB.union(tree, ts.tree)).asInstanceOf[That]
       case _ =>
         val b = bf(repr)