[Proof of concept] Polymorphic function types

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

@@ -90,6 +90,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 TailRec,                // Rewrite tail recursion to loops
          new Mixin,                  // Expand trait fields and trait initializers
          new LazyVals,               // Expand lazy vals

compiler/src/dotty/tools/dotc/ast/Desugar.scala

@@ -1430,6 +1430,35 @@ object desugar {
     }
 
     val desugared = tree match {
+      case PolyFunction(targs, body) =>
+        val Function(vargs, res) = body
+        // TODO: Figure out if we need a `PolyFunctionWithMods` instead.
+        val mods = body match {
+          case body: FunctionWithMods => body.mods
+          case _ => untpd.EmptyModifiers
+        }
+        val polyFunctionTpt = ref(defn.PolyFunctionType)
+        val applyTParams = targs.asInstanceOf[List[TypeDef]]
+        if (ctx.mode.is(Mode.Type)) {
+          // Desugar [T_1, ..., T_M] -> (P_1, ..., P_N) => R
+          // Into    scala.PolyFunction { def apply[T_1, ..., T_M](x$1: P_1, ..., x$N: P_N): R }
+
+          val applyVParams = vargs.zipWithIndex.map { case (p, n) =>
+            makeSyntheticParameter(n + 1, p).withAddedFlags(mods.flags)
+          }
+          RefinedTypeTree(polyFunctionTpt, List(
+            DefDef(nme.apply, applyTParams, List(applyVParams), res, EmptyTree)
+          ))
+        } else {
+          // Desugar [T_1, ..., T_M] -> (x_1: P_1, ..., x_N: P_N) => body
+          // Into    new scala.PolyFunction { def apply[T_1, ..., T_M](x_1: P_1, ..., x_N: P_N) = body }
+
+          val applyVParams = vargs.asInstanceOf[List[ValDef]]
+            .map(varg => varg.withAddedFlags(mods.flags | Param))
+          New(Template(emptyConstructor, List(polyFunctionTpt), Nil, EmptyValDef,
+            List(DefDef(nme.apply, applyTParams, List(applyVParams), TypeTree(), res))
+          ))
+        }
       case SymbolLit(str) =>
         Literal(Constant(scala.Symbol(str)))
       case InterpolatedString(id, segments) =>

compiler/src/dotty/tools/dotc/ast/Trees.scala

@@ -331,6 +331,7 @@ object Trees {
     }
 
     def withFlags(flags: FlagSet): ThisTree[Untyped] = withMods(untpd.Modifiers(flags))
+    def withAddedFlags(flags: FlagSet): ThisTree[Untyped] = withMods(rawMods | flags)
 
     def setComment(comment: Option[Comment]): this.type = {
       comment.map(putAttachment(DocComment, _))

compiler/src/dotty/tools/dotc/ast/untpd.scala

@@ -72,6 +72,12 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   class FunctionWithMods(args: List[Tree], body: Tree, val mods: Modifiers)(implicit @constructorOnly src: SourceFile)
     extends Function(args, body)
 
+  /** A polymorphic function type */
+  case class PolyFunction(targs: List[Tree], body: Tree)(implicit @constructorOnly src: SourceFile) extends Tree {
+    override def isTerm = body.isTerm
+    override def isType = body.isType
+  }
+
   /** A function created from a wildcard expression
    *  @param  placeholderParams  a list of definitions of synthetic parameters.
    *  @param  body               the function body where wildcards are replaced by
@@ -491,6 +497,10 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
       case tree: Function if (args eq tree.args) && (body eq tree.body) => tree
       case _ => finalize(tree, untpd.Function(args, body)(tree.source))
     }
+    def PolyFunction(tree: Tree)(targs: List[Tree], body: Tree)(implicit ctx: Context): Tree = tree match {
+      case tree: PolyFunction if (targs eq tree.targs) && (body eq tree.body) => tree
+      case _ => finalize(tree, untpd.PolyFunction(targs, body)(tree.source))
+    }
     def InfixOp(tree: Tree)(left: Tree, op: Ident, right: Tree)(implicit ctx: Context): Tree = tree match {
       case tree: InfixOp if (left eq tree.left) && (op eq tree.op) && (right eq tree.right) => tree
       case _ => finalize(tree, untpd.InfixOp(left, op, right)(tree.source))
@@ -579,6 +589,8 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
         cpy.InterpolatedString(tree)(id, segments.mapConserve(transform))
       case Function(args, body) =>
         cpy.Function(tree)(transform(args), transform(body))
+      case PolyFunction(targs, body) =>
+        cpy.PolyFunction(tree)(transform(targs), transform(body))
       case InfixOp(left, op, right) =>
         cpy.InfixOp(tree)(transform(left), op, transform(right))
       case PostfixOp(od, op) =>
@@ -634,6 +646,8 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
         this(x, segments)
       case Function(args, body) =>
         this(this(x, args), body)
+      case PolyFunction(targs, body) =>
+        this(this(x, targs), body)
       case InfixOp(left, op, right) =>
         this(this(this(x, left), op), right)
       case PostfixOp(od, op) =>

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

@@ -1035,6 +1035,9 @@ class Definitions {
     if (n <= MaxImplementedFunctionArity && (!isContextual || ctx.erasedTypes) && !isErased) ImplementedFunctionType(n)
     else FunctionClass(n, isContextual, isErased).typeRef
 
+  lazy val PolyFunctionClass = ctx.requiredClass("scala.PolyFunction")
+  def PolyFunctionType = PolyFunctionClass.typeRef
+
   /** If `cls` is a class in the scala package, its name, otherwise EmptyTypeName */
   def scalaClassName(cls: Symbol)(implicit ctx: Context): TypeName =
     if (cls.isClass && cls.owner == ScalaPackageClass) cls.asClass.name else EmptyTypeName

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

@@ -196,10 +196,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
     }
   }
 
@@ -383,6 +397,7 @@ class TypeErasure(isJava: Boolean, semiEraseVCs: Boolean, isConstructor: Boolean
    *      - otherwise, if T is a type parameter coming from Java, []Object
    *      - otherwise, Object
    *   - For a term ref p.x, the type <noprefix> # x.
+   *   - For a refined type scala.PolyFunction { def apply[...](x_1, ..., x_N): R }, scala.FunctionN
    *   - For a typeref scala.Any, scala.AnyVal, scala.Singleton, scala.Tuple, or scala.*: : |java.lang.Object|
    *   - For a typeref scala.Unit, |scala.runtime.BoxedUnit|.
    *   - For a typeref scala.FunctionN, where N > MaxImplementedFunctionArity, scala.FunctionXXL
@@ -429,6 +444,12 @@ class TypeErasure(isJava: Boolean, semiEraseVCs: Boolean, isConstructor: Boolean
       SuperType(this(thistpe), this(supertpe))
     case ExprType(rt) =>
       defn.FunctionType(0)
+    case RefinedType(parent, nme.apply, refinedInfo) if parent.typeSymbol eq defn.PolyFunctionClass =>
+      assert(refinedInfo.isInstanceOf[PolyType])
+      val res = refinedInfo.resultType
+      val paramss = res.paramNamess
+      assert(paramss.length == 1)
+      this(defn.FunctionType(paramss.head.length, isContextual = res.isImplicitMethod, isErased = res.isErasedMethod))
     case tp: TypeProxy =>
       this(tp.underlying)
     case AndType(tp1, tp2) =>

compiler/src/dotty/tools/dotc/parsing/Parsers.scala

@@ -924,6 +924,8 @@ object Parsers {
           val tparams = typeParamClause(ParamOwner.TypeParam)
           if (in.token == TLARROW)
             atSpan(start, in.skipToken())(LambdaTypeTree(tparams, toplevelTyp()))
+          else if (isIdent && in.name.toString == "->")
+            atSpan(start, in.skipToken())(PolyFunction(tparams, toplevelTyp()))
           else { accept(TLARROW); typ() }
         }
         else infixType()
@@ -1323,6 +1325,11 @@ object Parsers {
         atSpan(in.skipToken()) { Return(if (isExprIntro) expr() else EmptyTree, EmptyTree) }
       case FOR =>
         forExpr()
+      case LBRACKET =>
+        val start = in.offset
+        val tparams = typeParamClause(ParamOwner.TypeParam)
+        assert(isIdent && in.name.toString == "->", "Expected `->`")
+        atSpan(start, in.skipToken())(PolyFunction(tparams, expr()))
       case _ =>
         if (isIdent(nme.inline) && !in.inModifierPosition() && in.lookaheadIn(canStartExpressionTokens)) {
           val start = in.skipToken()

compiler/src/dotty/tools/dotc/printing/RefinedPrinter.scala

@@ -558,6 +558,11 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
 		  (keywordText("erased ") provided isErased) ~
           argsText ~ " => " ~ toText(body)
         }
+      case PolyFunction(targs, body) =>
+        val targsText = "[" ~ Text(targs.map((arg: Tree) => toText(arg)), ", ") ~ "]"
+        changePrec(GlobalPrec) {
+          targsText ~ " -> " ~ toText(body)
+        }
       case InfixOp(l, op, r) =>
         val opPrec = parsing.precedence(op.name)
         changePrec(opPrec) { toText(l) ~ " " ~ toText(op) ~ " " ~ toText(r) }

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

@@ -87,7 +87,10 @@ class ElimErasedValueType extends MiniPhase with InfoTransformer { thisPhase =>
       val site = root.thisType
       val info1 = site.memberInfo(sym1)
       val info2 = site.memberInfo(sym2)
-      if (!info1.matchesLoosely(info2))
+      if (!info1.matchesLoosely(info2) &&
+          !(sym1.name == nme.apply &&
+            (sym1.owner.derivesFrom(defn.PolyFunctionClass) ||
+             sym2.owner.derivesFrom(defn.PolyFunctionClass))))
         ctx.error(DoubleDefinition(sym1, sym2, root), root.sourcePos)
     }
     val earlyCtx = ctx.withPhase(ctx.elimRepeatedPhase.next)

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

@@ -416,21 +416,28 @@ object Erasure {
      */
     override def typedSelect(tree: untpd.Select, pt: Type)(implicit ctx: Context): Tree = {
 
+      val qual1 = typed(tree.qualifier, AnySelectionProto)
+
       def mapOwner(sym: Symbol): Symbol = {
-        def recur(owner: Symbol): Symbol =
-          if (defn.specialErasure.contains(owner)) {
+        def recur(owner: Symbol): Symbol = {
+          val owner = sym.maybeOwner
+          if (!owner.exists) {
+            // Hack for PolyFunction#apply
+            qual1.tpe.widen.typeSymbol
+          } else if (defn.specialErasure.contains(owner)) {
             assert(sym.isConstructor, s"${sym.showLocated}")
             defn.specialErasure(owner)
           } else if (defn.isSyntheticFunctionClass(owner))
             defn.erasedFunctionClass(owner)
           else
             owner
-        recur(sym.owner)
+        }
+        recur(sym.maybeOwner)
       }
 
       val origSym = tree.symbol
       val owner = mapOwner(origSym)
-      val sym = if (owner eq origSym.owner) origSym else owner.info.decl(origSym.name).symbol
+      val sym = if (owner eq origSym.maybeOwner) origSym else owner.info.decl(tree.name).symbol
       assert(sym.exists, origSym.showLocated)
 
       def select(qual: Tree, sym: Symbol): Tree =
@@ -474,7 +481,7 @@ object Erasure {
         }
       }
 
-      checkNotErased(recur(typed(tree.qualifier, AnySelectionProto)))
+      checkNotErased(recur(qual1))
     }
 
     override def typedThis(tree: untpd.This)(implicit ctx: Context): Tree =

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

@@ -54,7 +54,8 @@ trait TypeAssigner {
           required = EmptyFlagConjunction, excluded = Private)
           .suchThat(decl.matches(_))
       val inheritedInfo = inherited.info
-      if (inheritedInfo.exists &&
+      val isPolyFunctionApply = decl.name == nme.apply && (parent <:< defn.PolyFunctionType)
+      if (isPolyFunctionApply || inheritedInfo.exists &&
           decl.info.widenExpr <:< inheritedInfo.widenExpr &&
           !(inheritedInfo.widenExpr <:< decl.info.widenExpr)) {
         val r = RefinedType(parent, decl.name, decl.info)

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

@@ -1271,7 +1271,9 @@ class Typer extends Namer
       typr.println(s"adding refinement $refinement")
       checkRefinementNonCyclic(refinement, refineCls, seen)
       val rsym = refinement.symbol
-      if (rsym.info.isInstanceOf[PolyType] && rsym.allOverriddenSymbols.isEmpty)
+      val polymorphicRefinementAllowed =
+        tpt1.tpe.typeSymbol == defn.PolyFunctionClass && rsym.name == nme.apply
+      if (!polymorphicRefinementAllowed && rsym.info.isInstanceOf[PolyType] && rsym.allOverriddenSymbols.isEmpty)
         ctx.error(PolymorphicMethodMissingTypeInParent(rsym, tpt1.symbol), refinement.sourcePos)
 
       val member = refineCls.info.member(rsym.name)

library/src/scala/PolyFunction.scala

@@ -0,0 +1,10 @@
+package scala
+
+/** Marker trait for polymorphic function types.
+ *
+ *  This is the only trait that can be refined with a polymorphic method,
+ *  as long as that method is called `apply`, e.g.:
+ *      PolyFunction { def apply[T_1, ..., T_M](x_1: P_1, ..., x_N: P_N): R }
+ *  This type will be erased to FunctionN.
+ */
+trait PolyFunction

tests/neg/bad-selftype.scala

@@ -2,5 +2,5 @@ trait x0[T] { self: x0 => }           // error
 
 trait x1[T] { self: (=> String) => }  // error
 
-trait x2[T] { self: ([X] => X) => }   // error
+trait x2[T] { self: ([X] =>> X) => }   // error
 

tests/neg/i4373.scala

@@ -17,7 +17,7 @@ object Test {
   type T1 = _ // error
   type T2 = _[Int] // error
   type T3 = _ { type S } // error
-  type T4 = [X] => _ // error
+  type T4 = [X] =>> _ // error
 
   // Open questions:
   type T5 = TypeConstr[_ { type S }] // error

tests/neg/i6385a.scala

@@ -7,5 +7,5 @@ object Test {
   def f[F[_]](x: Box[F]) = ???
   def db: Box[D] = ???
   def cb: Box[C] = db  // error
-  f[[X] => C[X]](db)   // error
-}
+  f[[X] =>> C[X]](db)   // error
+}

tests/run/polymorphic-functions.scala

@@ -0,0 +1,11 @@
+object Test {
+  def test1(f: [T <: AnyVal] -> List[T] => List[(T, T)]) = {
+    f(List(1, 2, 3))
+  }
+
+  def main(args: Array[String]): Unit = {
+    val fun = [T <: AnyVal] -> (x: List[T]) => x.map(e => (e, e))
+
+    assert(test1(fun) == List((1, 1), (2, 2), (3, 3)))
+  }
+}