Adapt instance of TreeTypeMap to map trees recursively

Solves an issue involving return values not being specialized.
Test `return_specialization` illustrates this.

Author: AlexSikia

src/dotty/tools/dotc/ast/Trees.scala

@@ -109,8 +109,8 @@ object Trees {
      */
     val uniqueId = {
       nextId += 1
-      /*if (nextId == 651)
-        println("asdsad")*/
+      if (nextId == 72)
+        println("asdsad")
       //assert(nextId != 214, this)
       nextId
     }

src/dotty/tools/dotc/core/Symbols.scala

@@ -344,9 +344,9 @@ object Symbols {
 
     type ThisName <: Name
     private[this] var _id: Int = nextId
-    /*if (_id == 651) {
+    if (_id == 16651) {
       println("plop")
-    }*/
+    }
     //assert(_id != 30214)
 
     /** The unique id of this symbol */

src/dotty/tools/dotc/transform/TypeSpecializer.scala

@@ -102,7 +102,9 @@ class TypeSpecializer extends MiniPhaseTransform  with InfoTransformer {
                                   (implicit ctx: Context): List[Symbol] = {
       val newSym =
         ctx.newSymbol(decl.owner, (decl.name + names.mkString).toTermName,
-                      decl.flags | Flags.Synthetic, poly.instantiate(instantiations.toList))
+                      decl.flags | Flags.Synthetic,poly.instantiate(instantiations.toList))
+                      //poly.derivedPolyType(poly.paramNames, poly.paramBounds,
+                        //                   poly.instantiate(instantiations.toList)))
       val map = newSymbolMap.getOrElse(decl, mutable.HashMap.empty)
       map.put(instantiations, newSym)
       newSymbolMap.put(decl, map)
@@ -150,26 +152,39 @@ class TypeSpecializer extends MiniPhaseTransform  with InfoTransformer {
         def specialize(decl : Symbol): List[Tree] = {
           if (newSymbolMap.contains(decl)) {
             val declSpecs = newSymbolMap(decl)
+
             val newSyms = declSpecs.values.toList
             val instantiations = declSpecs.keys.toArray
             var index = -1
             println(s"specializing ${tree.symbol} for $origTParams")
-          newSyms.map { newSym =>
-            index += 1
-            polyDefDef(newSym.asTerm, { tparams => vparams => {
-              assert(tparams.isEmpty)
-              new TreeTypeMap(
-                typeMap = _
-                  .substDealias(origTParams, instantiations(index))
-                  .subst(origVParams, vparams.flatten.map(_.tpe)),
-                oldOwners = tree.symbol :: Nil,
-                newOwners = newSym :: Nil
-              ).transform(tree.rhs)
-            }})
-          }
+            /*val attempted = newSyms.zip(*/newSyms.map { newSym =>
+              index += 1
+              polyDefDef(newSym.asTerm, { tparams => vparams => {
+                //assert(tparams.isEmpty)
+
+                val tmap: (Tree => Tree) = _ match {
+                  case Return(t, from) if from.symbol == tree.symbol => Return(t, ref(newSym))
+                  case t: TypeApply => transformTypeApply(t)
+                  case t => t
+                }
+
+                new TreeTypeMap(
+                  treeMap = tmap,
+                  typeMap = _
+                    .substDealias(origTParams, instantiations(index))
+                    .subst(origVParams, vparams.flatten.map(_.tpe))
+                  ,
+                  oldOwners = tree.symbol :: Nil,
+                  newOwners = newSym :: Nil
+                ).transform(tree.rhs)
+              }
+              })
+            }//)
+            //attempted.filterNot(a => failedSpec.contains(a._1)).map(_._2)
         } else Nil
         }
-        Thicket(tree :: specialize(tree.symbol))
+        val specialized_trees = specialize(tree.symbol)
+        Thicket(tree :: specialized_trees)
       case _ => tree
     }
   }
@@ -206,4 +221,16 @@ class TypeSpecializer extends MiniPhaseTransform  with InfoTransformer {
       } else tree
     } else tree
   }
-}
+}
+
+/**
+ * var failedSpec: List[Symbol] = List()
+ *
+ *
+ * match {
+                  case t if !poly.bounds.contains(t) => {
+                    failedSpec = newSym :: failedSpec
+                    t
+                  }
+                  case t => t
+  */

test/dotc/tests.scala

@@ -207,18 +207,19 @@ class tests extends CompilerTest {
 
   val javaDir = "./tests/pos/java-interop/"
   @Test def java_all = compileFiles(javaDir, twice)
-*/
 
-  //@Test def specialization = compileFile(specialDir, "specialization")//, specialise)
-  //@Test def mutual_spec = compileFile(specialDir, "mutual_specialization")
+*/
+  //@Test def specialization = compileFile(specialDir, "specialization")
+  //@Test def mutual_spec = compileFile(specialDir, "mutual_specialization", List("-Xprint:all"))
   //@Test def return_spec = compileFile(specialDir, "return_specialization")
-  //@Test def nothing_spec = compileFile(specialDir, "nothing_specialization")
-  //@Test def method_in_class_spec = compileFile(specialDir, "method_in_class_specialization")
-  //@Test def method_in_method_spec = compileFile(specialDir, "method_in_method_specialization", List("-Xprint:all"))
-  @Test def type_check_spec = compileFile(specialDir, "type_check_specialization")
-  //@Test def bounds_spec = compileFile(specialDir, "bounds_specialization", List("-Xprint:all"))
-  //@Test def multi_spec = compileFile(specialDir, "multi_specialization", List("-Xprint:all"))
-  //@Test def pos_spec_all = compileFiles(specialDir)
+  //  @Test def nothing_spec = compileFile(specialDir, "nothing_specialization")
+  //  @Test def method_in_class_spec = compileFile(specialDir, "method_in_class_specialization")
+  //  @Test def method_in_method_spec = compileFile(specialDir, "method_in_method_specialization")
+  //  @Test def pos_type_check = compileFile(specialDir, "type_test")
+  //  @Test def bounds_spec = compileFile(specialDir, "bounds_specialization")
+  //  @Test def multi_spec = compileFile(specialDir, "multi_specialization")
+  //  @Test def pos_spec_all = compileFiles(specialDir)
+    @Test def pos_this_specialization = compileFile(specialDir, "this_specialization", List("-Xprint:specialize"))
 
   //@Test def mini_method = compileFiles(miniMethodDir)//, List("-Xprint:all"))
   //@Test def mini_more = compileFiles(miniMoreDir)//, List("-Xprint:all"))

tests/pos/specialization/return_specialization.scala

@@ -1,6 +1,6 @@
 object return_specialization {
-  def qwa[@specialized T](a: (String, String) => T, b: T): T = {
-      if(a ne this) return a("1", "2")
+  def qwa[@specialized(Int) T](a: (T, T) => T, b: T): T = {
+      if(a ne this) return a(b, b)
       else b
   }
 }

tests/pos/specialization/this_specialization.scala

@@ -0,0 +1,5 @@
+sealed abstract class Foo[@specialized +A] {
+  def bop[@specialized B >: A]: Foo[B] = new Bar[B](this)
+}
+
+case class Bar[@specialized a](tl: Foo[a]) extends Foo[a]

tests/pos/specialization/type_check_specialization.scala

@@ -0,0 +1,3 @@
+object type_test {
+  def typeTest(i: Char): Unit = i.isInstanceOf[Int]
+}