Support polymorphic function values

A polymorphic function value can be written as:
    new PolyFunction {
      def apply[T_1, ..., T_M](x_1: P_1, ..., x_N: P_N): R = body
    }
This is erased to:
    new FunctionN {
      def apply(x_1: Object, ..., x_N: Object): Object = body
    }

Getting everything to erase correctly was tricky, the current
implementation is a bit messy currently.

Author: smarter

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -99,6 +99,7 @@ class Compiler {
     List(new Erasure) ::             // Rewrite types to JVM model, erasing all type parameters, abstract types and refinements.
     List(new ElimErasedValueType,    // Expand erased value types to their underlying implmementation types
          new VCElideAllocations,     // Peep-hole optimization to eliminate unnecessary value class allocations
+         new ElimPolyFunction,       // Rewrite PolyFunction subclasses to FunctionN subclasses
          new Mixin,                  // Expand trait fields and trait initializers
          new LazyVals,               // Expand lazy vals
          new Memoize,                // Add private fields to getters and setters

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

@@ -184,10 +184,24 @@ object TypeErasure {
           MethodType(Nil, defn.BoxedUnitType)
         else if (sym.isAnonymousFunction && einfo.paramInfos.length > MaxImplementedFunctionArity)
           MethodType(nme.ALLARGS :: Nil, JavaArrayType(defn.ObjectType) :: Nil, einfo.resultType)
+        else if (sym.name == nme.apply && sym.owner.derivesFrom(defn.PolyFunctionClass)) {
+          // The erasure of `apply` in subclasses of PolyFunction has to match
+          // the erasure of FunctionN#apply, since after `ElimPolyFunction` we replace
+          // a `PolyFunction` parent by a `FunctionN` parent.
+          einfo.derivedLambdaType(
+            paramInfos = einfo.paramInfos.map(_ => defn.ObjectType),
+            resType = defn.ObjectType
+          )
+        }
         else
           einfo
       case einfo =>
-        einfo
+        // Erase the parameters of `apply` in subclasses of PolyFunction
+        if (sym.is(TermParam) && sym.owner.name == nme.apply
+            && sym.owner.owner.derivesFrom(defn.PolyFunctionClass))
+          defn.ObjectType
+        else
+          einfo
     }
   }
 

compiler/src/dotty/tools/dotc/transform/ElimErasedValueType.scala

@@ -89,7 +89,10 @@ class ElimErasedValueType extends MiniPhase with InfoTransformer {
       val info1 = site.memberInfo(sym1)
       val info2 = site.memberInfo(sym2)
       def isDefined(sym: Symbol) = sym.originDenotation.validFor.firstPhaseId <= ctx.phaseId
-      if (isDefined(sym1) && isDefined(sym2) && !info1.matchesLoosely(info2))
+      if (isDefined(sym1) && isDefined(sym2) && !info1.matchesLoosely(info2)
+          && !(sym1.name == nme.apply &&
+               (sym1.owner.derivesFrom(defn.PolyFunctionClass)
+                || sym2.owner.derivesFrom(defn.PolyFunctionClass))))
         // The reason for the `isDefined` condition is that we need to exclude mixin forwarders
         // from the tests. For instance, in compileStdLib, compiling scala.immutable.SetProxy, line 29:
         //    new AbstractSet[B] with SetProxy[B] { val self = newSelf }

compiler/src/dotty/tools/dotc/transform/ElimPolyFunction.scala

@@ -0,0 +1,68 @@
+package dotty.tools.dotc
+package transform
+
+import ast.{Trees, tpd}
+import core._, core.Decorators._
+import MegaPhase._, Phases.Phase
+import Types._, Contexts._, Constants._, Names._, NameOps._, Flags._, DenotTransformers._
+import SymDenotations._, Symbols._, StdNames._, Annotations._, Trees._, Scopes._, Denotations._
+import TypeErasure.ErasedValueType, ValueClasses._
+
+/** This phase rewrite PolyFunction subclasses to FunctionN subclasses
+ *
+ *      class Foo extends PolyFunction {
+ *          def apply(x_1: P_1, ..., x_N: P_N): R = rhs
+ *      }
+ *   becomes:
+ *      class Foo extends FunctionN {
+ *          def apply(x_1: P_1, ..., x_N: P_N): R = rhs
+ *      }
+ */
+class ElimPolyFunction extends MiniPhase with DenotTransformer {
+
+  import tpd._
+
+  override def phaseName: String = ElimPolyFunction.name
+
+  override def runsAfter = Set(Erasure.name)
+
+  override def changesParents: Boolean = true // Replaces PolyFunction by FunctionN
+
+  override def transform(ref: SingleDenotation)(implicit ctx: Context) = ref match {
+    case ref: ClassDenotation if ref.symbol != defn.PolyFunctionClass && ref.derivesFrom(defn.PolyFunctionClass) =>
+      val cinfo = ref.classInfo
+      val newParent = functionTypeOfPoly(cinfo)
+      val newParents = cinfo.classParents.map(parent =>
+        if (parent.typeSymbol == defn.PolyFunctionClass)
+          newParent
+        else
+          parent
+      )
+      ref.copySymDenotation(info = cinfo.derivedClassInfo(classParents = newParents))
+    case _ =>
+      ref
+  }
+
+  def functionTypeOfPoly(cinfo: ClassInfo)(implicit ctx: Context): Type = {
+    val applyMeth = cinfo.decls.lookup(nme.apply).info
+    val arity = applyMeth.paramNamess.head.length
+    defn.FunctionType(arity)
+  }
+
+  override def transformTemplate(tree: Template)(implicit ctx: Context): Tree = {
+    val newParents = tree.parents.mapconserve(parent =>
+      if (parent.tpe.typeSymbol == defn.PolyFunctionClass) {
+        val cinfo = tree.symbol.owner.asClass.classInfo
+        tpd.TypeTree(functionTypeOfPoly(cinfo))
+      }
+      else
+        parent
+    )
+    cpy.Template(tree)(parents = newParents)
+  }
+}
+
+object ElimPolyFunction {
+  val name = "elimPolyFunction"
+}
+

compiler/src/dotty/tools/dotc/typer/TypeAssigner.scala

@@ -53,7 +53,8 @@ trait TypeAssigner {
         parentType.findMember(decl.name, cls.thisType, excluded = Private)
           .suchThat(decl.matches(_))
       val inheritedInfo = inherited.info
-      if (inheritedInfo.exists && decl.info <:< inheritedInfo && !(inheritedInfo <:< decl.info)) {
+      val isPolyFunctionApply = decl.name == nme.apply && (parent <:< defn.PolyFunctionType)
+      if (isPolyFunctionApply || inheritedInfo.exists && decl.info <:< inheritedInfo && !(inheritedInfo <:< decl.info)) {
         val r = RefinedType(parent, decl.name, decl.info)
         typr.println(i"add ref $parent $decl --> " + r)
         r

tests/run/polymorphic-functions.scala

@@ -4,6 +4,10 @@ object Test {
   }
 
   def main(args: Array[String]): Unit = {
-    //test1(...)
+    val fun = new PolyFunction {
+      def apply[T <: AnyVal](x: List[T]): List[(T, T)] = x.map(e => (e, e))
+    }
+
+    assert(test1(fun) == List((1, 1), (2, 2), (3, 3)))
   }
 }