Merge pull request #4672 from dotty-staging/poly-functions

[Proof of concept] Polymorphic function types

Author: milessabin

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

@@ -1418,6 +1418,41 @@ object desugar {
       }
     }
 
+    def makePolyFunction(targs: List[Tree], body: Tree): Tree = body match {
+      case  Function(vargs, res) =>
+        // 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 _ =>
+        // may happen for erroneous input. An error will already have been reported.
+        assert(ctx.reporter.errorsReported)
+        EmptyTree
+    }
+
     // begin desugar
 
     // Special case for `Parens` desugaring: unlike all the desugarings below,
@@ -1430,6 +1465,8 @@ object desugar {
     }
 
     val desugared = tree match {
+      case PolyFunction(targs, body) =>
+        makePolyFunction(targs, body) orElse tree
       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
@@ -489,6 +495,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))
@@ -577,6 +587,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) =>
@@ -632,6 +644,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

@@ -1044,6 +1044,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,27 @@ 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
+        // Preserve PolyFunction argument types to support PolyFunctions with
+        // PolyFunction arguments
+        if (sym.is(TermParam) && sym.owner.name == nme.apply
+            && sym.owner.owner.derivesFrom(defn.PolyFunctionClass)
+            && !(tp <:< defn.PolyFunctionType)) {
+          defn.ObjectType
+        } else
+          einfo
     }
   }
 
@@ -383,6 +400,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 +447,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) =>
@@ -581,6 +605,11 @@ class TypeErasure(isJava: Boolean, semiEraseVCs: Boolean, isConstructor: Boolean
       case tp: TypeVar =>
         val inst = tp.instanceOpt
         if (inst.exists) sigName(inst) else tpnme.Uninstantiated
+      case tp @ RefinedType(parent, nme.apply, _) if parent.typeSymbol eq defn.PolyFunctionClass => 
+        // we need this case rather than falling through to the default
+        // because RefinedTypes <: TypeProxy and it would be caught by
+        // the case immediately below
+        sigName(this(tp))
       case tp: TypeProxy =>
         sigName(tp.underlying)
       case _: ErrorType | WildcardType | NoType =>

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

@@ -3194,6 +3194,7 @@ object Types {
       companion.eq(ContextualMethodType) ||
       companion.eq(ErasedContextualMethodType)
 
+
     def computeSignature(implicit ctx: Context): Signature = {
       val params = if (isErasedMethod) Nil else paramInfos
       resultSignature.prepend(params, isJavaMethod)

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

@@ -850,9 +850,13 @@ object Parsers {
      */
     def toplevelTyp(): Tree = rejectWildcardType(typ())
 
-    /** Type        ::=  FunTypeMods FunArgTypes `=>' Type
-     *                |  HkTypeParamClause `=>>' Type
+    /** Type        ::=  FunType
+     *                |  HkTypeParamClause ‘=>>’ Type
+     *                |  MatchType
      *                |  InfixType
+     *  FunType     ::=  { 'erased' | 'given' } (MonoFunType | PolyFunType)
+     *  MonoFunType ::=  FunArgTypes ‘=>’ Type
+     *  PolyFunType ::=  HKTypeParamClause '=>' Type
      *  FunArgTypes ::=  InfixType
      *                |  `(' [ FunArgType {`,' FunArgType } ] `)'
      *                |  '(' TypedFunParam {',' TypedFunParam } ')'
@@ -924,7 +928,18 @@ object Parsers {
           val tparams = typeParamClause(ParamOwner.TypeParam)
           if (in.token == TLARROW)
             atSpan(start, in.skipToken())(LambdaTypeTree(tparams, toplevelTyp()))
-          else { accept(TLARROW); typ() }
+          else if (in.token == ARROW) {
+            val arrowOffset = in.skipToken()
+            val body = toplevelTyp()
+            atSpan(start, arrowOffset) {
+              body match {
+                case _: Function => PolyFunction(tparams, body)
+                case _ =>
+                  syntaxError("Implementation restriction: polymorphic function types must have a value parameter", arrowOffset)
+                  Ident(nme.ERROR.toTypeName)
+              }
+            }
+          } else { accept(TLARROW); typ() }
         }
         else infixType()
 
@@ -1223,6 +1238,7 @@ object Parsers {
      *                      |  `throw' Expr
      *                      |  `return' [Expr]
      *                      |  ForExpr
+     *                      |  HkTypeParamClause ‘=>’ Expr
      *                      |  [SimpleExpr `.'] id `=' Expr
      *                      |  SimpleExpr1 ArgumentExprs `=' Expr
      *                      |  Expr2
@@ -1323,6 +1339,19 @@ 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)
+        val arrowOffset = accept(ARROW)
+        val body = expr()
+        atSpan(start, arrowOffset) {
+          body match {
+            case _: Function => PolyFunction(tparams, body)
+            case _ =>
+              syntaxError("Implementation restriction: polymorphic function literals must have a value parameter", arrowOffset)
+              errorTermTree
+          }
+        }
       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

@@ -83,11 +83,19 @@ class ElimErasedValueType extends MiniPhase with InfoTransformer { thisPhase =>
       override def matches(sym1: Symbol, sym2: Symbol) =
         sym1.signature == sym2.signature
     }
+
     def checkNoConflict(sym1: Symbol, sym2: Symbol, info: Type)(implicit ctx: Context): Unit = {
       val site = root.thisType
       val info1 = site.memberInfo(sym1)
       val info2 = site.memberInfo(sym2)
-      if (!info1.matchesLoosely(info2))
+      // PolyFunction apply methods will be eliminated later during
+      // ElimPolyFunction, so we let them pass here.
+      def bothPolyApply =
+        sym1.name == nme.apply &&
+        (sym1.owner.derivesFrom(defn.PolyFunctionClass) ||
+         sym2.owner.derivesFrom(defn.PolyFunctionClass))
+
+      if (!info1.matchesLoosely(info2) && !bothPolyApply)
         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"
+}
+