Specialized keySets for immutable and sorted ma...

Specialized keySets for immutable and sorted maps are supposed to be
views, not actual sets. Closes #4616, no review.

2 files changed, 25 insertions, 3 deletions

diff --git a/src/library/scala/collection/immutable/MapLike.scala b/src/library/scala/collection/immutable/MapLike.scala
index fb2826b4df..beea72d676 100644
--- a/src/library/scala/collection/immutable/MapLike.scala
+++ b/src/library/scala/collection/immutable/MapLike.scala
@@ -116,7 +116,16 @@ trait MapLike[A, +B, +This <: MapLike[A, B, This] with Map[A, B]]
   /** Collects all keys of this map in a set.
    *  @return  a set containing all keys of this map.
    */
-  override def keySet: immutable.Set[A] = immutable.Set.empty ++ (this map (_._1))
+  override def keySet: immutable.Set[A] = new ImmutableDefaultKeySet
+
+  protected class ImmutableDefaultKeySet extends super.DefaultKeySet with immutable.Set[A] {
+    override def + (elem: A): immutable.Set[A] =
+      if (this(elem)) this
+      else immutable.Set[A]() ++ this + elem
+    override def - (elem: A): immutable.Set[A] =
+      if (this(elem)) immutable.Set[A]() ++ this - elem
+      else this
+  }
 
   /** This function transforms all the values of mappings contained
    *  in this map with function `f`.
diff --git a/src/library/scala/collection/immutable/SortedMap.scala b/src/library/scala/collection/immutable/SortedMap.scala
index f7e05fef69..902a0f8457 100644
--- a/src/library/scala/collection/immutable/SortedMap.scala
+++ b/src/library/scala/collection/immutable/SortedMap.scala
@@ -31,14 +31,27 @@ import annotation.bridge
 trait SortedMap[A, +B] extends Map[A, B]
                          with scala.collection.SortedMap[A, B]
                          with MapLike[A, B, SortedMap[A, B]]
-                         with SortedMapLike[A, B, SortedMap[A, B]] {
+                         with SortedMapLike[A, B, SortedMap[A, B]] { self =>
 
   override protected[this] def newBuilder : Builder[(A, B), SortedMap[A, B]] =
     SortedMap.newBuilder[A, B]
 
   override def empty: SortedMap[A, B] = SortedMap.empty
   override def updated [B1 >: B](key: A, value: B1): SortedMap[A, B1] = this + ((key, value))
-  override def keySet: immutable.SortedSet[A] = SortedSet.empty ++ (this map (_._1))
+  override def keySet: immutable.SortedSet[A] = new DefaultKeySortedSet
+
+  protected class DefaultKeySortedSet extends super.DefaultKeySortedSet with immutable.SortedSet[A] {
+    override def + (elem: A): SortedSet[A] =
+      if (this(elem)) this
+      else SortedSet[A]() ++ this + elem
+    override def - (elem: A): SortedSet[A] =
+      if (this(elem)) SortedSet[A]() ++ this - elem
+      else this
+    override def rangeImpl(from : Option[A], until : Option[A]) : SortedSet[A] = {
+      val map = self.rangeImpl(from, until)
+      new map.DefaultKeySortedSet
+    }
+  }
 
   /** Add a key/value pair to this map.
    *  @param    kv the key/value pair