Add lazy versions of sorted and immutable keysets, and deprecate the old versions

Author: natsukagami

library/src/scala/collection/SortedMap.scala

@@ -132,9 +132,22 @@ transparent trait SortedMapOps[K, +V, +CC[X, Y] <: Map[X, Y] & SortedMapOps[X, Y
       rangeUntil(next)
   }
 
-  override def keySet: SortedSet[K] = new KeySortedSet
+  override def keySet: SortedSet[K] = new LazyKeySortedSet
 
   /** The implementation class of the set returned by `keySet` */
+  private[collection] class LazyKeySortedSet extends MapOps.LazyKeySet(this) with SortedSet[K] {
+    implicit def ordering: Ordering[K] = SortedMapOps.this.ordering
+    def iteratorFrom(start: K): Iterator[K] = SortedMapOps.this.keysIteratorFrom(start)
+
+    override def diff(that: Set[K]): SortedSet[K] = fromSpecific(view.filterNot(that))
+    override def rangeImpl(from: Option[K], until: Option[K]): SortedSet[K] = {
+      val map = SortedMapOps.this.rangeImpl(from, until)
+      new map.LazyKeySortedSet
+    }
+  }
+
+  /** The old implementation class of the set returned by `keySet` */
+  @deprecated("KeySortedSet is no longer used in the .keySet implementation", since = "3.8.0")
   protected class KeySortedSet extends SortedSet[K] with GenKeySet with GenKeySortedSet {
     def diff(that: Set[K]): SortedSet[K] = fromSpecific(view.filterNot(that))
     def rangeImpl(from: Option[K], until: Option[K]): SortedSet[K] = {
@@ -144,6 +157,7 @@ transparent trait SortedMapOps[K, +V, +CC[X, Y] <: Map[X, Y] & SortedMapOps[X, Y
   }
 
   /** A generic trait that is reused by sorted keyset implementations */
+  @deprecated("GenKeySortedSet is no longer used in .keySet implementations", since = "3.8.0")
   protected trait GenKeySortedSet extends GenKeySet { this: SortedSet[K] =>
     implicit def ordering: Ordering[K] = SortedMapOps.this.ordering
     def iteratorFrom(start: K): Iterator[K] = SortedMapOps.this.keysIteratorFrom(start)

library/src/scala/collection/immutable/HashMap.scala

@@ -60,7 +60,8 @@ final class HashMap[K, +V] private[immutable] (private[immutable] val rootNode:
 
   override def keySet: Set[K] = if (size == 0) Set.empty else new HashKeySet
 
-  private[immutable] final class HashKeySet extends ImmutableKeySet {
+
+  private[immutable] final class HashKeySet extends LazyImmutableKeySet {
 
     private def newKeySetOrThis(newHashMap: HashMap[K, ?]): Set[K] =
       if (newHashMap eq HashMap.this) this else newHashMap.keySet

library/src/scala/collection/immutable/Map.scala

@@ -143,9 +143,17 @@ transparent trait MapOps[K, +V, +CC[X, +Y] <: MapOps[X, Y, CC, ?], +C <: MapOps[
     */
   def transform[W](f: (K, V) => W): CC[K, W] = map { case (k, v) => (k, f(k, v)) }
 
-  override def keySet: Set[K] = new ImmutableKeySet
+  override def keySet: Set[K] = new LazyImmutableKeySet
 
   /** The implementation class of the set returned by `keySet` */
+  private[immutable] class LazyImmutableKeySet extends MapOps.LazyKeySet(this) with Set[K] {
+    override def diff(that: collection.Set[K]): Set[K] = super.diff(that)
+    override def incl(elem: K): Set[K] = if (this(elem)) this else MapOps.this.updated(elem, ()).keySet
+    override def excl(elem: K): Set[K] = if (this(elem)) MapOps.this.removed(elem).keySet else this
+  }
+
+  /** The implementation class of the set returned by `keySet` */
+  @deprecated("ImmutableKeySet is no longer used in .keySet implementations", since = "3.8.0")
   protected[immutable] class ImmutableKeySet extends AbstractSet[K] with GenKeySet with DefaultSerializable {
     def incl(elem: K): Set[K] = if (this(elem)) this else empty ++ this + elem
     def excl(elem: K): Set[K] = if (this(elem)) empty ++ this - elem else this

library/src/scala/collection/immutable/Set.scala

@@ -110,7 +110,7 @@ object Set extends IterableFactory[Set] {
       case s: Set3[E]    => s
       case s: Set4[E]    => s
       case s: HashMap[E @unchecked, _]#HashKeySet => s
-      case s: MapOps[E, Any, Map, Map[E, Any]]#ImmutableKeySet @unchecked => s
+      case s: MapOps[E, Any, Map, Map[E, Any]]#LazyImmutableKeySet @unchecked => s
       // We also want `SortedSet` (and subclasses, such as `BitSet`)
       // to rebuild themselves, to avoid element type widening issues.
       case _ => newBuilder[E].addAll(it).result()

library/src/scala/collection/immutable/SortedMap.scala

@@ -93,9 +93,21 @@ transparent trait SortedMapOps[K, +V, +CC[X, +Y] <: Map[X, Y] & SortedMapOps[X,
 
   def unsorted: Map[K, V]
 
-  override def keySet: SortedSet[K] = new ImmutableKeySortedSet
+  override def keySet: SortedSet[K] = new LazyImmutableKeySortedSet
 
   /** The implementation class of the set returned by `keySet` */
+  private class LazyImmutableKeySortedSet extends LazyKeySortedSet with SortedSet[K] {
+    override def diff(that: scala.collection.Set[K]): SortedSet[K] = super.diff(that)
+    override def rangeImpl(from: Option[K], until: Option[K]): SortedSet[K] = {
+      val map = self.rangeImpl(from, until)
+      new map.LazyImmutableKeySortedSet
+    }
+    def incl(elem: K): SortedSet[K] = fromSpecific(this).incl(elem)
+    def excl(elem: K): SortedSet[K] = fromSpecific(this).excl(elem)
+  }
+
+  /** The implementation class of the set returned by `keySet` */
+  @deprecated("ImmutableKeySortedSet is no longer used by the .keySet implementation", since = "3.8.0")
   protected class ImmutableKeySortedSet extends AbstractSet[K] with SortedSet[K] with GenKeySet with GenKeySortedSet {
     def rangeImpl(from: Option[K], until: Option[K]): SortedSet[K] = {
       val map = self.rangeImpl(from, until)

library/src/scala/collection/mutable/LinkedHashMap.scala

@@ -234,6 +234,7 @@ class LinkedHashMap[K, V]
       def extract(nd: Entry): (K, V) = (nd.key, nd.value)
     }
 
+  @deprecated("LinkedKeySet is now strict and no longer used in the implementation of .keySet", "3.8.0")
   /** Note that a LinkedKeySet could be strict. */
   protected class LinkedKeySet extends KeySet {
     override def iterableFactory: IterableFactory[collection.Set] = LinkedHashSet