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/transform/TypeSpecializer.scala

@@ -101,7 +101,23 @@ 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))
+
+      /* The following generated symbols which kept type bounds. It served, as illustrated by the `this_specialization`
+       * test, as a way of keeping type bounds when instantiating a `this` referring to a generic class. However,
+       * because type bounds are not transitive, this did not work out and we introduced casts instead.
+       *
+       * ctx.newSymbol(decl.owner, (decl.name + names.mkString).toTermName,
+       *               decl.flags | Flags.Synthetic,
+       *               poly.derivedPolyType(poly.paramNames,
+       *                                    (poly.paramBounds zip instantiations).map
+       *                                             {case (bounds, instantiation) =>
+       *                                               TypeBounds(bounds.lo, AndType(bounds.hi, instantiation))},
+       *                                    poly.instantiate(indices, instantiations)
+       *                                    )
+       *               )
+       */
+
       val map = newSymbolMap.getOrElse(decl, mutable.HashMap.empty)
       map.put(instantiations, newSym)
       newSymbolMap.put(decl, map)
@@ -153,48 +169,85 @@ class TypeSpecializer extends MiniPhaseTransform  with InfoTransformer {
             val instantiations = declSpecs.keys.toArray
             var index = -1
             println(s"specializing ${tree.symbol} for $origTParams")
-          newSyms.map { newSym =>
+            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: Apply => transformApply(t)
+                case t => t
+              }
+
               new TreeTypeMap(
+                treeMap = tmap,
                 typeMap = _
-                  .substDealias(origTParams, instantiations(index))
-                  .subst(origVParams, vparams.flatten.map(_.tpe)),
+                    .substDealias(origTParams, instantiations(index))
+                    .subst(origVParams, vparams.flatten.map(_.tpe))
+                ,
                 oldOwners = tree.symbol :: Nil,
                 newOwners = newSym :: Nil
               ).transform(tree.rhs)
             }})
           }
         } else Nil
         }
-        Thicket(tree :: specialize(tree.symbol))
+        val specialized_trees = specialize(tree.symbol)
+        Thicket(tree :: specialized_trees)
       case _ => tree
     }
   }
 
   override def transformTypeApply(tree: tpd.TypeApply)(implicit ctx: Context, info: TransformerInfo): Tree = {
+    val TypeApply(fun, _) = tree
+    if (fun.tpe.isParameterless) rewireTree(tree)
+    tree
+  }
 
+  override def transformApply(tree: Apply)(implicit ctx: Context, info: TransformerInfo): Tree = {
+    val Apply(fun, args) = tree
+    fun match {
+      case fun: TypeApply => {
+        println(
+          s"""
+               |args             ->  ${args}
 
-    def allowedToSpecialize(sym: Symbol): Boolean = {
-      sym.name != nme.asInstanceOf_ &&
-        !(sym is Flags.JavaDefined) &&
-        !sym.isConstructor//isPrimaryConstructor
+              |f.fun            ->  ${fun.fun.tree}
+             """.stripMargin)
+
+        val newFun = rewireTree(fun)
+        if (fun ne newFun) {
+        val b = (args zip newFun.tpe.firstParamTypes)
+        val a = b.map{
+          case (arg, tpe) =>
+            arg.ensureConforms(tpe)
+        }
+        Apply(newFun,a)
+        /* zip (instantiations zip paramTypes)).map{
+              case (argType, (specType, castType)) => argType.ensureConforms(specType)})*/
+        } else tree
+      }
+      case _ => tree
     }
+  }
 
+  def rewireTree(tree: Tree)(implicit ctx: Context): Tree = {
+    assert(tree.isInstanceOf[TypeApply])
     val TypeApply(fun,args) = tree
     if (newSymbolMap.contains(fun.symbol)){
       val newSymInfos = newSymbolMap(fun.symbol)
       val betterDefs = newSymInfos.filter(x => (x._1 zip args).forall{a =>
-         val specializedType = a._1
-         val argType = a._2
+        val specializedType = a._1
+        val argType = a._2
         argType.tpe <:< specializedType
       }).toList
 
-      if (betterDefs.length > 1) ctx.debuglog("Several specialized variants fit.")
-      //assert(betterDefs.length < 2) // TODO: How to select the best if there are several ?
+      if (betterDefs.length > 1) {
+        ctx.debuglog("Several specialized variants fit.")
+        tree
+      }
 
-      if (betterDefs.nonEmpty) {
+      else if (betterDefs.nonEmpty) {
         val best = betterDefs.head
         println(s"method ${fun.symbol.name} of ${fun.symbol.owner} rewired to specialized variant with type (${best._1})")
         val prefix = fun match {
@@ -213,4 +266,3 @@ class TypeSpecializer extends MiniPhaseTransform  with InfoTransformer {
     } else tree
   }
 }
-

test/dotc/tests.scala

@@ -240,16 +240,18 @@ 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,8 @@
+sealed abstract class Foo[@specialized +A] {
+  def bop[@specialized B >: A]: Foo[B] = new Bar[B](this)
+  //def bip[@specialized C >: A, @specialized D >: A]: Foo[D] = new Cho[D, C](new Bar[C](this))
+}
+
+case class Bar[@specialized a](tl: Foo[a]) extends Foo[a]
+
+//case class Cho[@specialized c, @specialized d](tl: Bar[d]) extends Foo[c]

tests/pos/specialization/type_check_specialization.scala

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