Specialize byname functions

I tried to extend specialization to all context functions, not just ones of 0 arity. But that runs into problems for dependent context functions, since the necessary casts get complicated. Since context functions over primitive types are an anti-pattern anyway I don't think we need to optimize this case, after all.
scala · Jan 20, 2022 · 177896a · 177896a
1 parent 6c43bed
commit 177896a
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 19 deletions.
diff --git a/compiler/src/dotty/tools/dotc/core/Definitions.scala b/compiler/src/dotty/tools/dotc/core/Definitions.scala
@@ -1558,7 +1558,8 @@ class Definitions {
     new PerRun(Function2SpecializedReturnTypes.map(_.symbol))
 
   def isSpecializableFunction(cls: ClassSymbol, paramTypes: List[Type], retType: Type)(using Context): Boolean =
-    paramTypes.length <= 2 && cls.derivesFrom(FunctionClass(paramTypes.length))
+    paramTypes.length <= 2
+    && (cls.derivesFrom(FunctionClass(paramTypes.length)) || isByNameFunctionClass(cls))
     && isSpecializableFunctionSAM(paramTypes, retType)
 
   /** If the Single Abstract Method of a Function class has this type, is it specializable? */

diff --git a/compiler/src/dotty/tools/dotc/transform/SpecializeFunctions.scala b/compiler/src/dotty/tools/dotc/transform/SpecializeFunctions.scala
@@ -76,24 +76,27 @@ class SpecializeFunctions extends MiniPhase {
             argTypes,
             retType
           )
-
-        if (!isSpecializable || argTypes.exists(_.isInstanceOf[ExprType])) return tree
-
-        val specializedApply = nme.apply.specializedFunction(retType, argTypes)
-        val newSel = fun match {
-          case Select(qual, _) =>
-            qual.select(specializedApply)
-          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.appliedToTermArgs(args)
-
+        if isSpecializable then
+          val specializedApply = nme.apply.specializedFunction(retType, argTypes)
+          val newSel = fun match
+            case Select(qual, _) =>
+              val qual1 = qual.tpe.widen match
+                case defn.ByNameFunction(res) =>
+                  // Need to cast to regular function, since specialied apply methods
+                  // are not members of ContextFunction0. The cast will be eliminated in
+                  // erasure.
+                  qual.cast(defn.FunctionOf(Nil, res))
+                case _ =>
+                  qual
+              qual1.select(specializedApply)
+            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.appliedToTermArgs(args)
+        else tree
       case _ => tree
     }