Second version of function specialization

Duhemm · liufengyun · commit 8c75922e8a35 · 2020-11-24T21:55:11.000+01:00
Fix compilation errors

Don't change parents in `specializeFunctions`

Cleanup

Move type to JVM tag conversion to Definitions

Add more tests for function specialization

Pass by name should not introduce boxing

Cleanup

No specialization for Function3

Adapt to recent changes in transformers and phases

Address review comments

Optimize phase `SpecializeFunctions`

 - Stop immediately if the type doesn't derive from `Function{0,1,2}`
 - We don't need to check if any of the parents derives from
   `Function{0,1,2}`, we can just check if the type derives from it.
diff --git a/compiler/src/dotty/tools/dotc/core/NameOps.scala b/compiler/src/dotty/tools/dotc/core/NameOps.scala
@@ -279,9 +279,9 @@ object NameOps {
     /** This method is to be used on **type parameters** from a class, since
      *  this method does sorting based on their names
      */
-    def specializedFor(classTargs: List[Types.Type], classTargsNames: List[Name], methodTargs: List[Types.Type], methodTarsNames: List[Name])(implicit ctx: Context): name.ThisName = {
-      val methodTags: Seq[Name] = (methodTargs zip methodTarsNames).sortBy(_._2).map(x => typeToTag(x._1))
-      val classTags: Seq[Name] = (classTargs zip classTargsNames).sortBy(_._2).map(x => typeToTag(x._1))
+    def specializedFor(classTargs: List[Types.Type], classTargsNames: List[Name], methodTargs: List[Types.Type], methodTarsNames: List[Name])(implicit ctx: Context): Name = {
+      val methodTags: Seq[Name] = (methodTargs zip methodTarsNames).sortBy(_._2).map(x => defn.typeTag(x._1))
+      val classTags: Seq[Name] = (classTargs zip classTargsNames).sortBy(_._2).map(x => defn.typeTag(x._1))
 
       likeSpacedN(name ++ nme.specializedTypeNames.prefix ++
         methodTags.fold(nme.EMPTY)(_ ++ _) ++ nme.specializedTypeNames.separator ++
@@ -294,10 +294,10 @@ object NameOps {
      *
      *  `<return type><first type><second type><...>`
      */
-    def specializedFunction(ret: Types.Type, args: List[Types.Type])(implicit ctx: Context): name.ThisName =
+    def specializedFunction(ret: Types.Type, args: List[Types.Type])(implicit ctx: Context): Name =
       name ++ nme.specializedTypeNames.prefix ++
-      nme.specializedTypeNames.separator ++ typeToTag(ret) ++
-      args.map(typeToTag).fold(nme.EMPTY)(_ ++ _) ++ nme.specializedTypeNames.suffix
+      nme.specializedTypeNames.separator ++ defn.typeTag(ret) ++
+      args.map(defn.typeTag).fold(nme.EMPTY)(_ ++ _) ++ nme.specializedTypeNames.suffix
 
     /** If name length exceeds allowable limit, replace part of it by hash */
     def compactified(using Context): TermName = termName(compactify(name.toString))
diff --git a/compiler/src/dotty/tools/dotc/core/Names.scala b/compiler/src/dotty/tools/dotc/core/Names.scala
@@ -35,7 +35,7 @@ object Names {
   abstract class Name extends Designator, Showable derives CanEqual {
 
     /** A type for names of the same kind as this name */
-    type ThisName <: Name { type ThisName = self.ThisName }
+    type ThisName <: Name
 
     /** Is this name a type name? */
     def isTypeName: Boolean
diff --git a/compiler/src/dotty/tools/dotc/transform/SpecializeFunctions.scala b/compiler/src/dotty/tools/dotc/transform/SpecializeFunctions.scala
@@ -1,190 +1,156 @@
 package dotty.tools.dotc
 package transform
 
-import TreeTransforms.{ MiniPhaseTransform, TransformerInfo }
 import ast.Trees._, ast.tpd, core._
 import Contexts.Context, Types._, Decorators._, Symbols._, DenotTransformers._
 import SymDenotations._, Scopes._, StdNames._, NameOps._, Names._
+import MegaPhase.MiniPhase
 
 import scala.collection.mutable
 
 /** Specializes classes that inherit from `FunctionN` where there exists a
  *  specialized form.
  */
-class SpecializeFunctions extends MiniPhaseTransform with InfoTransformer {
+class SpecializeFunctions extends MiniPhase with InfoTransformer {
   import ast.tpd._
   val phaseName = "specializeFunctions"
+  override def runsAfter = Set(classOf[ElimByName])
 
-  private[this] var _blacklistedSymbols: List[Symbol] = _
+  private val jFunction = "scala.compat.java8.JFunction".toTermName
 
-  private def blacklistedSymbols(implicit ctx: Context): List[Symbol] = {
-    if (_blacklistedSymbols eq null) _blacklistedSymbols = List(
-      ctx.getClassIfDefined("scala.math.Ordering").asClass.membersNamed("Ops".toTypeName).first.symbol
-    )
-
-    _blacklistedSymbols
-  }
-
-  /** Transforms the type to include decls for specialized applys and replace
-   *  the class parents with specialized versions.
-   */
-  def transformInfo(tp: Type, sym: Symbol)(implicit ctx: Context) = tp match {
-    case tp: ClassInfo if !sym.is(Flags.Package) && (tp.decls ne EmptyScope) => {
+  /** Transforms the type to include decls for specialized applys  */
+  override def transformInfo(tp: Type, sym: Symbol)(implicit ctx: Context) = tp match {
+    case tp: ClassInfo if !sym.is(Flags.Package) && (tp.decls ne EmptyScope) && derivesFromFn012(sym) =>
       var newApplys = Map.empty[Name, Symbol]
 
-      val newParents = tp.parents.mapConserve { parent =>
-        List(0, 1, 2, 3).flatMap { arity =>
-          val func = defn.FunctionClass(arity)
-          if (!parent.derivesFrom(func)) Nil
-          else {
-            val typeParams = tp.typeRef.baseArgInfos(func)
-            val interface = specInterface(typeParams)
-
-            if (interface.exists) {
-              if (tp.decls.lookup(nme.apply).exists) {
-                val specializedMethodName = nme.apply.specializedFunction(typeParams.last, typeParams.init)
-                newApplys = newApplys + (specializedMethodName -> interface)
-              }
+      var arity = 0
+      while (arity < 3) {
+        val func = defn.FunctionClass(arity)
+        if (tp.derivesFrom(func)) {
+          val typeParams = tp.cls.typeRef.baseType(func).argInfos
+          val isSpecializable =
+            defn.isSpecializableFunction(
+              sym.asClass,
+              typeParams.init,
+              typeParams.last
+            )
 
-              if (parent.isRef(func)) List(interface.typeRef)
-              else Nil
-            }
-            else Nil
+          if (isSpecializable && tp.decls.lookup(nme.apply).exists) {
+            val interface = specInterface(typeParams)
+            val specializedMethodName = nme.apply.specializedFunction(typeParams.last, typeParams.init)
+            newApplys += (specializedMethodName -> interface)
           }
         }
-        .headOption
-        .getOrElse(parent)
+        arity += 1
       }
 
       def newDecls =
-        if (newApplys.isEmpty) tp.decls
-        else
-          newApplys.toList.map { case (name, interface) =>
-            ctx.newSymbol(
-              sym,
-              name,
-              Flags.Override | Flags.Method,
-              interface.info.decls.lookup(name).info
-            )
-          }
-          .foldLeft(tp.decls.cloneScope) {
-            (scope, sym) => scope.enter(sym); scope
-          }
+        newApplys.toList.map { case (name, interface) =>
+          ctx.newSymbol(
+            sym,
+            name,
+            Flags.Override | Flags.Method | Flags.Synthetic,
+            interface.info.decls.lookup(name).info
+          )
+        }
+        .foldLeft(tp.decls.cloneScope) {
+          (scope, sym) => scope.enter(sym); scope
+        }
 
-      tp.derivedClassInfo(
-        classParents = newParents,
-        decls = newDecls
-      )
-    }
+      if (newApplys.isEmpty) tp
+      else tp.derivedClassInfo(decls = newDecls)
 
     case _ => tp
   }
 
   /** Transforms the `Template` of the classes to contain forwarders from the
-   *  generic applys to the specialized ones. Also replaces parents of the
-   *  class on the tree level and inserts the specialized applys in the
-   *  template body.
+   *  generic applys to the specialized ones. Also inserts the specialized applys
+   *  in the template body.
    */
-  override def transformTemplate(tree: Template)(implicit ctx: Context, info: TransformerInfo) = {
-    val applyBuf = new mutable.ListBuffer[Tree]
-    val newBody = tree.body.mapConserve {
-      case dt: DefDef if dt.name == nme.apply && dt.vparamss.length == 1 => {
-        val specName = nme.apply.specializedFunction(
-          dt.tpe.widen.finalResultType,
-          dt.vparamss.head.map(_.symbol.info)
-        )
-
-        val specializedApply = tree.symbol.enclosingClass.info.decls.lookup(specName)//member(specName).symbol
-        //val specializedApply = tree.symbol.enclosingClass.info.member(specName).symbol
-
-        if (false) {
-          println(tree.symbol.enclosingClass.show)
-          println("'" + specName.show + "'")
-          println(specializedApply)
-          println(specializedApply.exists)
-        }
-
-
-        if (specializedApply.exists) {
-          val apply = specializedApply.asTerm
-          val specializedDecl =
-            polyDefDef(apply, trefs => vrefss => {
-              dt.rhs
-                .changeOwner(dt.symbol, apply)
-                .subst(dt.vparamss.flatten.map(_.symbol), vrefss.flatten.map(_.symbol))
+  override def transformTemplate(tree: Template)(implicit ctx: Context) = {
+    val cls = tree.symbol.enclosingClass.asClass
+    if (derivesFromFn012(cls)) {
+      val applyBuf = new mutable.ListBuffer[Tree]
+      val newBody = tree.body.mapConserve {
+        case dt: DefDef if dt.name == nme.apply && dt.vparamss.length == 1 =>
+          val typeParams = dt.vparamss.head.map(_.symbol.info)
+          val retType = dt.tpe.widen.finalResultType
+
+          val specName = specializedName(nme.apply, typeParams :+ retType)
+          val specializedApply = cls.info.decls.lookup(specName)
+          if (specializedApply.exists) {
+            val apply = specializedApply.asTerm
+            val specializedDecl =
+              polyDefDef(apply, trefs => vrefss => {
+                dt.rhs
+                  .changeOwner(dt.symbol, apply)
+                  .subst(dt.vparamss.flatten.map(_.symbol), vrefss.flatten.map(_.symbol))
+              })
+            applyBuf += specializedDecl
+
+            // create a forwarding to the specialized apply
+            cpy.DefDef(dt)(rhs = {
+              tpd
+                .ref(apply)
+                .appliedToArgs(dt.vparamss.head.map(vparam => ref(vparam.symbol)))
             })
-          applyBuf += specializedDecl
-
-          // create a forwarding to the specialized apply
-          cpy.DefDef(dt)(rhs = {
-            tpd
-            .ref(apply)
-            .appliedToArgs(dt.vparamss.head.map(vparam => ref(vparam.symbol)))
-          })
-        } else dt
-      }
-      case x => x
-    }
-
-    val missing: List[TypeTree] = List(0, 1, 2, 3).flatMap { arity =>
-      val func = defn.FunctionClass(arity)
-      val tr = tree.symbol.enclosingClass.typeRef
+          } else dt
 
-      if (!tr.parents.exists(_.isRef(func))) Nil
-      else {
-        val typeParams = tr.baseArgInfos(func)
-        val interface = specInterface(typeParams)
-
-        if (interface.exists) List(interface.info)
-        else Nil
+        case x => x
       }
-    }.map(TypeTree)
 
-    cpy.Template(tree)(
-      parents = tree.parents ++ missing,
-      body = applyBuf.toList ++ newBody
-    )
+      cpy.Template(tree)(
+        body = applyBuf.toList ::: newBody
+      )
+    } else tree
   }
 
   /** Dispatch to specialized `apply`s in user code when available */
-  override def transformApply(tree: Apply)(implicit ctx: Context, info: TransformerInfo) =
+  override def transformApply(tree: Apply)(implicit ctx: Context) =
     tree match {
-      case app @ Apply(fun, args)
+      case Apply(fun, args)
         if fun.symbol.name == nme.apply &&
         fun.symbol.owner.derivesFrom(defn.FunctionClass(args.length))
-      => {
+      =>
         val params = (fun.tpe.widen.firstParamTypes :+ tree.tpe).map(_.widenSingleton.dealias)
-        val specializedApply = specializedName(nme.apply, params)
-
-        if (!params.exists(_.isInstanceOf[ExprType]) && fun.symbol.owner.info.decls.lookup(specializedApply).exists) {
+        val isSpecializable =
+          defn.isSpecializableFunction(
+            fun.symbol.owner.asClass,
+            params.init,
+            params.last)
+
+        if (isSpecializable && !params.exists(_.isInstanceOf[ExprType])) {
+          val specializedApply = specializedName(nme.apply, params)
           val newSel = fun match {
             case Select(qual, _) =>
               qual.select(specializedApply)
-            case _ => {
+            case _ =>
               (fun.tpe: @unchecked) match {
                 case TermRef(prefix: ThisType, name) =>
                   tpd.This(prefix.cls).select(specializedApply)
                 case TermRef(prefix: NamedType, name) =>
                   tpd.ref(prefix).select(specializedApply)
               }
-            }
           }
 
           newSel.appliedToArgs(args)
         }
         else tree
-      }
+
       case _ => tree
     }
 
-  @inline private def specializedName(name: Name, args: List[Type])(implicit ctx: Context) =
-    name.specializedFor(args, args.map(_.typeSymbol.name), Nil, Nil)
+  private def specializedName(name: Name, args: List[Type])(implicit ctx: Context) =
+    name.specializedFunction(args.last, args.init)
 
-  @inline private def specInterface(typeParams: List[Type])(implicit ctx: Context) = {
-    val specName =
-      ("JFunction" + (typeParams.length - 1)).toTermName
-      .specializedFunction(typeParams.last, typeParams.init)
+  private def functionName(typeParams: List[Type])(implicit ctx: Context) =
+    jFunction ++ (typeParams.length - 1).toString
 
-    ctx.getClassIfDefined("scala.compat.java8.".toTermName ++ specName)
-  }
+  private def specInterface(typeParams: List[Type])(implicit ctx: Context) =
+    ctx.getClassIfDefined(functionName(typeParams).specializedFunction(typeParams.last, typeParams.init))
+
+  private def derivesFromFn012(sym: Symbol)(implicit ctx: Context): Boolean =
+    sym.derivesFrom(defn.FunctionClass(0)) ||
+      sym.derivesFrom(defn.FunctionClass(1)) ||
+      sym.derivesFrom(defn.FunctionClass(2))
 }
diff --git a/compiler/src/dotty/tools/dotc/transform/SpecializedApplyMethods.scala b/compiler/src/dotty/tools/dotc/transform/SpecializedApplyMethods.scala
@@ -1,10 +1,10 @@
 package dotty.tools.dotc
 package transform
 
-import TreeTransforms.{ MiniPhaseTransform, TransformerInfo }
 import ast.Trees._, ast.tpd, core._
 import Contexts.Context, Types._, Decorators._, Symbols._, DenotTransformers._
 import SymDenotations._, Scopes._, StdNames._, NameOps._, Names._
+import MegaPhase.MiniPhase
 
 /** This phase synthesizes specialized methods for FunctionN, this is done
  *  since there are no scala signatures in the bytecode for the specialized
@@ -14,14 +14,14 @@ import SymDenotations._, Scopes._, StdNames._, NameOps._, Names._
  *  can hardcode them. This should, however be removed once we're using a
  *  different standard library.
  */
-class SpecializedApplyMethods extends MiniPhaseTransform with InfoTransformer {
+class SpecializedApplyMethods extends MiniPhase with InfoTransformer {
   import ast.tpd._
 
   val phaseName = "specializedApplyMethods"
 
-  private[this] var func0Applys: List[Symbol] = _
-  private[this] var func1Applys: List[Symbol] = _
-  private[this] var func2Applys: List[Symbol] = _
+  private[this] var func0Applys: collection.Set[Symbol] = _
+  private[this] var func1Applys: collection.Set[Symbol] = _
+  private[this] var func2Applys: collection.Set[Symbol] = _
   private[this] var func0: Symbol = _
   private[this] var func1: Symbol = _
   private[this] var func2: Symbol = _
@@ -30,30 +30,30 @@ class SpecializedApplyMethods extends MiniPhaseTransform with InfoTransformer {
     val definitions = ctx.definitions
     import definitions._
 
-    def specApply(sym: Symbol, args: List[Type], ret: Type)(implicit ctx: Context) = {
+    def specApply(sym: Symbol, args: List[Type], ret: Type)(implicit ctx: Context): Symbol = {
       val name = nme.apply.specializedFunction(ret, args)
       ctx.newSymbol(sym, name, Flags.Method, MethodType(args, ret))
     }
 
     func0 = FunctionClass(0)
-    func0Applys = for (r <- ScalaValueTypes.toList) yield specApply(func0, Nil, r)
+    func0Applys = for (r <- defn.Function0SpecializedReturns) yield specApply(func0, Nil, r)
 
     func1 = FunctionClass(1)
     func1Applys = for {
-      r  <- List(UnitType, BooleanType, IntType, FloatType, LongType, DoubleType)
-      t1 <- List(IntType, LongType, FloatType, DoubleType)
+      r  <- defn.Function1SpecializedReturns
+      t1 <- defn.Function1SpecializedParams
     } yield specApply(func1, List(t1), r)
 
     func2 = FunctionClass(2)
     func2Applys = for {
-      r  <- List(UnitType, BooleanType, IntType, FloatType, LongType, DoubleType)
-      t1 <- List(IntType, LongType, DoubleType)
-      t2 <- List(IntType, LongType, DoubleType)
+      r  <- Function2SpecializedReturns
+      t1 <- Function2SpecializedParams
+      t2 <- Function2SpecializedReturns
     } yield specApply(func2, List(t1, t2), r)
   }
 
   /** Add symbols for specialized methods to FunctionN */
-  def transformInfo(tp: Type, sym: Symbol)(implicit ctx: Context) = tp match {
+  override def transformInfo(tp: Type, sym: Symbol)(implicit ctx: Context) = tp match {
     case tp: ClassInfo if defn.isPlainFunctionClass(sym) => {
       init()
       val newDecls = sym.name.functionArity match {
@@ -75,7 +75,7 @@ class SpecializedApplyMethods extends MiniPhaseTransform with InfoTransformer {
   }
 
   /** Create bridge methods for FunctionN with specialized applys */
-  override def transformTemplate(tree: Template)(implicit ctx: Context, info: TransformerInfo) = {
+  override def transformTemplate(tree: Template)(implicit ctx: Context) = {
     val owner = tree.symbol.owner
     val additionalSymbols =
       if (owner eq func0) func0Applys
diff --git a/compiler/test/dotty/tools/dotc/transform/SpecializeFunctionsTests.scala b/compiler/test/dotty/tools/dotc/transform/SpecializeFunctionsTests.scala
