Fix soundness of Map.from and Set.from

som-snytt · som-snytt · commit 7a67d21db166 · 2023-08-13T17:27:27.000-07:00
Map.from and Set.from restrict which collections
can be safely returned.

Map#toMap leverages Map.from for this purpose.

Inline SeqMap.from per review

Delete instanceof kludge

Move tests to junit
diff --git a/library/src/scala/collection/immutable/HashMap.scala b/library/src/scala/collection/immutable/HashMap.scala
@@ -57,7 +57,7 @@ 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 final class HashKeySet extends ImmutableKeySet {
+  private[immutable] final class HashKeySet extends ImmutableKeySet {
 
     private[this] def newKeySetOrThis(newHashMap: HashMap[K, _]): Set[K] =
       if (newHashMap eq HashMap.this) this else newHashMap.keySet
diff --git a/library/src/scala/collection/immutable/Map.scala b/library/src/scala/collection/immutable/Map.scala
@@ -18,6 +18,7 @@ import scala.annotation.unchecked.uncheckedVariance
 import scala.collection.generic.DefaultSerializable
 import scala.collection.immutable.Map.Map4
 import scala.collection.mutable.{Builder, ReusableBuilder}
+import SeqMap.{SeqMap1, SeqMap2, SeqMap3, SeqMap4}
 
 /** Base type of immutable Maps */
 trait Map[K, +V]
@@ -28,7 +29,7 @@ trait Map[K, +V]
 
   override def mapFactory: scala.collection.MapFactory[Map] = Map
 
-  override final def toMap[K2, V2](implicit ev: (K, V) <:< (K2, V2)): Map[K2, V2] = this.asInstanceOf[Map[K2, V2]]
+  override final def toMap[K2, V2](implicit ev: (K, V) <:< (K2, V2)): Map[K2, V2] = Map.from(this.asInstanceOf[Map[K2, V2]])
 
   /** The same map with a given default function.
     *  Note: The default is only used for `apply`. Other methods like `get`, `contains`, `iterator`, `keys`, etc.
@@ -206,11 +207,30 @@ object Map extends MapFactory[Map] {
 
   def empty[K, V]: Map[K, V] = EmptyMap.asInstanceOf[Map[K, V]]
 
-  def from[K, V](it: collection.IterableOnce[(K, V)]): Map[K, V] =
+  def from[K, V](it: IterableOnce[(K, V)]): Map[K, V] =
     it match {
       case it: Iterable[_] if it.isEmpty => empty[K, V]
-      case m: Map[K, V] => m
-      case _ => (newBuilder[K, V] ++= it).result()
+      // Since IterableOnce[(K, V)] launders the variance of K,
+      // identify only our implementations which can be soundly substituted.
+      // For example, the ordering used by sorted maps would fail on widened key type. (scala/bug#12745)
+      // The following type test is not sufficient: case m: Map[K, V] => m
+      case m: HashMap[K, V]    => m
+      case m: Map1[K, V]       => m
+      case m: Map2[K, V]       => m
+      case m: Map3[K, V]       => m
+      case m: Map4[K, V]       => m
+      //case m: WithDefault[K, V] => m    // cf SortedMap.WithDefault
+      //case m: SeqMap[K, V]     => SeqMap.from(it) // inlined here to avoid hard dependency
+      case m: ListMap[K, V]    => m
+      case m: TreeSeqMap[K, V] => m
+      case m: VectorMap[K, V]  => m
+      case m: SeqMap1[K, V]    => m
+      case m: SeqMap2[K, V]    => m
+      case m: SeqMap3[K, V]    => m
+      case m: SeqMap4[K, V]    => m
+
+      // Maps with a reified key type must be rebuilt, such as `SortedMap` and `IntMap`.
+      case _ => newBuilder[K, V].addAll(it).result()
     }
 
   def newBuilder[K, V]: Builder[(K, V), Map[K, V]] = new MapBuilderImpl
diff --git a/library/src/scala/collection/immutable/SeqMap.scala b/library/src/scala/collection/immutable/SeqMap.scala
@@ -46,7 +46,15 @@ object SeqMap extends MapFactory[SeqMap] {
 
   def from[K, V](it: collection.IterableOnce[(K, V)]): SeqMap[K, V] =
     it match {
-      case sm: SeqMap[K, V] => sm
+      //case sm: SeqMap[K, V] => sm
+      case m: ListMap[K, V]    => m
+      case m: TreeSeqMap[K, V] => m
+      case m: VectorMap[K, V]  => m
+      case m: SeqMap1[K, V]    => m
+      case m: SeqMap2[K, V]    => m
+      case m: SeqMap3[K, V]    => m
+      case m: SeqMap4[K, V]    => m
+      case it: Iterable[_] if it.isEmpty => empty[K, V]
       case _ => (newBuilder[K, V] ++= it).result()
     }
 
@@ -66,7 +74,7 @@ object SeqMap extends MapFactory[SeqMap] {
   }
 
   @SerialVersionUID(3L)
-  private final class SeqMap1[K, +V](key1: K, value1: V) extends SeqMap[K,V] with Serializable {
+  private[immutable] final class SeqMap1[K, +V](key1: K, value1: V) extends SeqMap[K,V] with Serializable {
     override def size: Int = 1
     override def knownSize: Int = 1
     override def apply(key: K) = if (key == key1) value1 else throw new NoSuchElementException("key not found: " + key)
@@ -90,7 +98,7 @@ object SeqMap extends MapFactory[SeqMap] {
   }
 
   @SerialVersionUID(3L)
-  private final class SeqMap2[K, +V](key1: K, value1: V, key2: K, value2: V) extends SeqMap[K,V] with Serializable {
+  private[immutable] final class SeqMap2[K, +V](key1: K, value1: V, key2: K, value2: V) extends SeqMap[K,V] with Serializable {
     override def size: Int = 2
     override def knownSize: Int = 2
     override def apply(key: K) =
@@ -125,7 +133,7 @@ object SeqMap extends MapFactory[SeqMap] {
   }
 
   @SerialVersionUID(3L)
-  private class SeqMap3[K, +V](key1: K, value1: V, key2: K, value2: V, key3: K, value3: V) extends SeqMap[K,V] with Serializable {
+  private[immutable] class SeqMap3[K, +V](key1: K, value1: V, key2: K, value2: V, key3: K, value3: V) extends SeqMap[K,V] with Serializable {
     override def size: Int = 3
     override def knownSize: Int = 3
     override def apply(key: K) =
@@ -166,7 +174,7 @@ object SeqMap extends MapFactory[SeqMap] {
   }
 
   @SerialVersionUID(3L)
-  private final class SeqMap4[K, +V](key1: K, value1: V, key2: K, value2: V, key3: K, value3: V, key4: K, value4: V) extends SeqMap[K,V] with Serializable {
+  private[immutable] final class SeqMap4[K, +V](key1: K, value1: V, key2: K, value2: V, key3: K, value3: V, key4: K, value4: V) extends SeqMap[K,V] with Serializable {
     override def size: Int = 4
     override def knownSize: Int = 4
     override def apply(key: K) =
diff --git a/library/src/scala/collection/immutable/Set.scala b/library/src/scala/collection/immutable/Set.scala
@@ -96,12 +96,21 @@ object Set extends IterableFactory[Set] {
 
   def from[E](it: collection.IterableOnce[E]): Set[E] =
     it match {
-      // We want `SortedSet` (and subclasses, such as `BitSet`) to
-      // rebuild themselves to avoid element type widening issues
-      case _: SortedSet[E]         => (newBuilder[E] ++= it).result()
-      case _ if it.knownSize == 0  => empty[E]
-      case s: Set[E]               => s
-      case _                       => (newBuilder[E] ++= it).result()
+      case _ if it.knownSize == 0 => empty[E]
+      // Since IterableOnce[E] launders the variance of E,
+      // identify only our implementations which can be soundly substituted.
+      // It's not sufficient to match `SortedSet[E]` to rebuild and `Set[E]` to retain.
+      case s: HashSet[E] => s
+      case s: ListSet[E] => s
+      case s: Set1[E]    => s
+      case s: Set2[E]    => s
+      case s: Set3[E]    => s
+      case s: Set4[E]    => s
+      //case s: MapOps[E @unchecked, _, _, _]#ImmutableKeySet => s
+      case s: HashMap[E @unchecked, _]#HashKeySet => 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()
     }
 
   def newBuilder[A]: Builder[A, Set[A]] = new SetBuilderImpl[A]
