Second version of function specialization

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.

Author: Duhemm

compiler/src/dotty/tools/dotc/core/NameOps.scala

@@ -274,30 +274,12 @@ object NameOps {
       case nme.clone_ => nme.clone_
     }
 
-<<<<<<< HEAD
-    def specializedFor(classTargs: List[Type], classTargsNames: List[Name], methodTargs: List[Type], methodTarsNames: List[Name])(using Context): N = {
-=======
-    private def typeToTag(tp: Types.Type)(implicit ctx: Context): Name =
-      tp.classSymbol match {
-        case t if t eq defn.IntClass     => nme.specializedTypeNames.Int
-        case t if t eq defn.BooleanClass => nme.specializedTypeNames.Boolean
-        case t if t eq defn.ByteClass    => nme.specializedTypeNames.Byte
-        case t if t eq defn.LongClass    => nme.specializedTypeNames.Long
-        case t if t eq defn.ShortClass   => nme.specializedTypeNames.Short
-        case t if t eq defn.FloatClass   => nme.specializedTypeNames.Float
-        case t if t eq defn.UnitClass    => nme.specializedTypeNames.Void
-        case t if t eq defn.DoubleClass  => nme.specializedTypeNames.Double
-        case t if t eq defn.CharClass    => nme.specializedTypeNames.Char
-        case _                           => nme.specializedTypeNames.Object
-      }
->>>>>>> Fix ordering of specialized names and type parameterized apply
-
     /** 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 ++
@@ -310,10 +292,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))

compiler/src/dotty/tools/dotc/core/Names.scala

@@ -32,10 +32,10 @@ object Names {
    *  in a name table. A derived term name adds a tag, and possibly a number
    *  or a further simple name to some other name.
    */
-  abstract class Name extends Designator, Showable derives Eql { self =>
+  abstract class Name extends Designator with PreName {
 
     /** 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

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))
 }