Rewrite collection.{Map, Set}.+ and collection.Set.-

This is incomplete since, it's blocked by scalameta/scalameta#1212 (Add support to query type of a term)

iset: immutable.Set[Int]
cset: collecion.Set[Int]

both have +/- implemented via SetLike

given

iset + 1
cset + 1

we know that + is from SetLike, but it's not possible to get the type of iset/cset

Rewrite collection.Set.+ restrict type of lhs

scalafix allow us to match on SetLike.+ but we cannot check the type of the expression on the lhs

We can work arround this limitation if the lhs is a identifier (simple case). This allow us
to document the expected type on the lhs and prepare the implementation once scalameta/scalameta#1212 is resolved

Author: MasseGuillaume

scalafix/input/src/main/scala/fix/SetMapSrc.scala

@@ -3,12 +3,27 @@ rule = "scala:fix.Scalacollectioncompat_newcollections"
  */
 package fix
 
-class SetMapSrc(set: Set[Int], map: Map[Int, Int]) {
-  set + (2, 3)
-  map + (2 -> 3, 3 -> 4)
-  (set + (2, 3)).map(x => x)
-  set + (2, 3) - 4
-  map.mapValues(_ + 1)
-  map.zip(List())
+import scala.collection
+import scala.collection.immutable
+import scala.collection.mutable.{Map, Set} // Challenge to make sure the scoping is correct
+
+class SetMapSrc(iset: immutable.Set[Int], 
+                cset: collection.Set[Int],
+                imap: immutable.Map[Int, Int],
+                cmap: collection.Map[Int, Int]) {
+  iset + (2, 3)
+  imap + (2 -> 3, 3 -> 4)
+  (iset + (2, 3)).toString
+  iset + (2, 3) - 4
+  imap.mapValues(_ + 1)
+  iset + 1
+  iset - 2
+  cset + 1
+  cset - 2
+  cmap + (2 -> 3)
+  cmap + ((4, 5))
+  imap + (2 -> 3)
+  imap + ((4, 5))
+  imap.zip(List())
   List().zip(List())
 }

scalafix/output/src/main/scala/fix/SetMapSrc.scala

@@ -3,12 +3,27 @@
 
 package fix
 
-class SetMapSrc(set: Set[Int], map: Map[Int, Int]) {
-  set + 2 + 3
-  map + (2 -> 3) + (3 -> 4)
-  (set + 2 + 3).map(x => x)
-  set + 2 + 3 - 4
-  map.mapValues(_ + 1).toMap
-  map.zip(List()).toMap
+import scala.collection
+import scala.collection.immutable
+import scala.collection.mutable.{Map, Set} // Challenge to make sure the scoping is correct
+
+class SetMapSrc(iset: immutable.Set[Int], 
+                cset: collection.Set[Int],
+                imap: immutable.Map[Int, Int],
+                cmap: collection.Map[Int, Int]) {
+  iset + 2 + 3
+  imap + (2 -> 3) + (3 -> 4)
+  (iset + 2 + 3).toString
+  iset + 2 + 3 - 4
+  imap.mapValues(_ + 1).toMap
+  iset + 1
+  iset - 2
+  cset ++ _root_.scala.collection.Set(1)
+  cset -- _root_.scala.collection.Set(2)
+  cmap ++ _root_.scala.collection.Map(2 -> 3)
+  cmap ++ _root_.scala.collection.Map((4, 5))
+  imap + (2 -> 3)
+  imap + ((4, 5))
+  imap.zip(List()).toMap
   List().zip(List())
 }

scalafix/rules/src/main/scala/fix/Scalacollectioncompat_newcollections.scala

@@ -30,6 +30,8 @@ case class Scalacollectioncompat_newcollections(index: SemanticdbIndex)
     Symbol("_root_.scala.Predef.Map#")
   )
 
+  val CollectionSet: TypeMatcher = TypeMatcher(Symbol("_root_.scala.collection.Set#"))
+
   def replaceSymbols(ctx: RuleCtx): Patch = {
     ctx.replaceSymbols(
       "scala.collection.LinearSeq" -> "scala.collection.immutable.List",
@@ -52,6 +54,18 @@ case class Scalacollectioncompat_newcollections(index: SemanticdbIndex)
     Symbol("_root_.scala.runtime.Tuple2Zipped.Ops.zipped."),
     Symbol("_root_.scala.runtime.Tuple3Zipped.Ops.zipped.")
   )
+  val setPlus = 
+    SymbolMatcher.exact(
+      Symbol("_root_.scala.collection.SetLike#`+`(Ljava/lang/Object;)Lscala/collection/Set;.")
+    )
+  val setMinus = 
+    SymbolMatcher.exact(
+      Symbol("_root_.scala.collection.SetLike#`-`(Ljava/lang/Object;)Lscala/collection/Set;.")
+    )
+  val mapPlus = 
+    SymbolMatcher.exact(
+      Symbol("_root_.scala.collection.MapLike#`+`(Lscala/Tuple2;)Lscala/collection/Map;.")
+    )
   val setPlus2 = SymbolMatcher.exact(
     Symbol("_root_.scala.collection.SetLike#`+`(Ljava/lang/Object;Ljava/lang/Object;Lscala/collection/Seq;)Lscala/collection/Set;.")
   )
@@ -110,6 +124,9 @@ case class Scalacollectioncompat_newcollections(index: SemanticdbIndex)
       Symbol("_root_.scala.collection.IterableLike#zip(Lscala/collection/GenIterable;Lscala/collection/generic/CanBuildFrom;)Ljava/lang/Object;.")
     )
 
+  def startsWithParens(tree: Tree): Boolean = 
+    tree.tokens.headOption.map(_.is[Token.LeftParen]).getOrElse(false)
+
   def replaceMutableSet(ctx: RuleCtx) =
     ctx.tree.collect {
       case retainSet(n: Name) =>
@@ -220,7 +237,7 @@ case class Scalacollectioncompat_newcollections(index: SemanticdbIndex)
   def replaceSetMapPlus2(ctx: RuleCtx): Patch = {
     def rewritePlus(ap: Term.ApplyInfix, lhs: Term, op: Term.Name, rhs1: Term, rhs2: Term): Patch = {
       val tokensToReplace =
-        if(ap.tokens.headOption.map(_.is[Token.LeftParen]).getOrElse(false)) {
+        if(startsWithParens(ap)) {
           // don't drop surrounding parens
           ap.tokens.slice(1, ap.tokens.size - 1)
         } else ap.tokens
@@ -245,6 +262,34 @@ case class Scalacollectioncompat_newcollections(index: SemanticdbIndex)
     }.asPatch
   }
 
+  def replaceSetMapPlusMinus(ctx: RuleCtx): Patch = {
+    def rewriteOp(op: Tree, rhs: Tree, doubleOp: String, col0: String): Patch = {
+      val col = "_root_.scala.collection." + col0
+      val callSite =
+        if (startsWithParens(rhs)) {
+          ctx.addLeft(rhs, col)          
+        }
+        else {
+          ctx.addLeft(rhs, col + "(") +
+          ctx.addRight(rhs, ")")
+        }
+
+      ctx.addRight(op, doubleOp) + callSite
+    }
+
+    ctx.tree.collect {
+      case Term.ApplyInfix(CollectionSet(), op @ setPlus(_), Nil, List(rhs)) =>
+        rewriteOp(op, rhs, "+", "Set")
+        
+      case Term.ApplyInfix(CollectionSet(), op @ setMinus(_), Nil, List(rhs)) =>
+        rewriteOp(op, rhs, "-", "Set")
+
+      case Term.ApplyInfix(_, op @ mapPlus(_), Nil, List(rhs)) =>
+        rewriteOp(op, rhs, "+", "Map")
+    }.asPatch
+  }
+
+
   def replaceMutSetMapPlus(ctx: RuleCtx): Patch = {
     def rewriteMutPlus(lhs: Term, op: Term.Name): Patch = {
       ctx.addRight(lhs, ".clone()") +
@@ -299,7 +344,7 @@ case class Scalacollectioncompat_newcollections(index: SemanticdbIndex)
         ctx.addRight(ap, ".toMap")
     }.asPatch
   }
-  
+
   override def fix(ctx: RuleCtx): Patch = {
     replaceToList(ctx) +
       replaceSymbols(ctx) +
@@ -310,6 +355,7 @@ case class Scalacollectioncompat_newcollections(index: SemanticdbIndex)
       replaceMutableSet(ctx) +
       replaceSymbolicFold(ctx) +
       replaceSetMapPlus2(ctx) +
+      replaceSetMapPlusMinus(ctx) +
       replaceMutSetMapPlus(ctx) +
       replaceMutMapUpdated(ctx) +
       replaceIterableSameElements(ctx) +