Implement polymorphic lambdas using Closure nodes for efficiency

Previously, we desugared them manually into anonymous class instances,
but by using a Closure node instead, we ensure that they get translated
into indy lambdas on the JVM.

Also cleaned up and added a TODO in the desugaring of polymorphic function types
into refinement types since I realized that purity wasn't taken into account.

Author: smarter

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

@@ -1061,6 +1061,40 @@ object desugar {
     name
   }
 
+  /** Strip parens and empty blocks around the body of `tree`. */
+  def normalizePolyFunction(tree: PolyFunction)(using Context): PolyFunction =
+    def stripped(body: Tree): Tree = body match
+      case Parens(body1) =>
+        stripped(body1)
+      case Block(Nil, body1) =>
+        stripped(body1)
+      case _ => body
+    cpy.PolyFunction(tree)(tree.targs, stripped(tree.body)).asInstanceOf[PolyFunction]
+
+  /** 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 }
+   */
+  def makePolyFunctionType(tree: PolyFunction)(using Context): RefinedTypeTree =
+    val PolyFunction(tparams: List[untpd.TypeDef] @unchecked, fun @ untpd.Function(vparamTypes, res)) = tree: @unchecked
+    val funFlags = fun match
+      case fun: FunctionWithMods =>
+        fun.mods.flags
+      case _ => EmptyFlags
+
+    // TODO: make use of this in the desugaring when pureFuns is enabled.
+    // val isImpure = funFlags.is(Impure)
+
+    // Function flags to be propagated to each parameter in the desugared method type.
+    val paramFlags = funFlags.toTermFlags & Given
+    val vparams = vparamTypes.zipWithIndex.map:
+      case (p: ValDef, _) => p.withAddedFlags(paramFlags)
+      case (p, n) => makeSyntheticParameter(n + 1, p).withAddedFlags(paramFlags)
+
+    RefinedTypeTree(ref(defn.PolyFunctionType), List(
+       DefDef(nme.apply, tparams :: vparams :: Nil, res, EmptyTree).withFlags(Synthetic)
+    )).withSpan(tree.span)
+  end makePolyFunctionType
+
   /** Invent a name for an anonympus given of type or template `impl`. */
   def inventGivenOrExtensionName(impl: Tree)(using Context): SimpleName =
     val str = impl match
@@ -1454,14 +1488,17 @@ object desugar {
   }
 
   /** Make closure corresponding to function.
-   *      params => body
+   *      [tparams] => params => body
    *  ==>
-   *      def $anonfun(params) = body
+   *      def $anonfun[tparams](params) = body
    *      Closure($anonfun)
    */
-  def makeClosure(params: List[ValDef], body: Tree, tpt: Tree | Null = null, span: Span)(using Context): Block =
+  def makeClosure(tparams: List[TypeDef], vparams: List[ValDef], body: Tree, tpt: Tree | Null = null, span: Span)(using Context): Block =
+    val paramss: List[ParamClause] =
+      if tparams.isEmpty then vparams :: Nil
+      else tparams :: vparams :: Nil
     Block(
-      DefDef(nme.ANON_FUN, params :: Nil, if (tpt == null) TypeTree() else tpt, body)
+      DefDef(nme.ANON_FUN, paramss, if (tpt == null) TypeTree() else tpt, body)
         .withSpan(span)
         .withMods(synthetic | Artifact),
       Closure(Nil, Ident(nme.ANON_FUN), EmptyTree))
@@ -1753,56 +1790,6 @@ object desugar {
       }
     }
 
-    def makePolyFunction(targs: List[Tree], body: Tree, pt: Type): Tree = body match {
-      case Parens(body1) =>
-        makePolyFunction(targs, body1, pt)
-      case Block(Nil, body1) =>
-        makePolyFunction(targs, body1, pt)
-      case Function(vargs, res) =>
-        assert(targs.nonEmpty)
-        // 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: ValDef, _) => p.withAddedFlags(mods.flags)
-            case (p, n) => makeSyntheticParameter(n + 1, p).withAddedFlags(mods.flags.toTermFlags)
-          }
-          RefinedTypeTree(polyFunctionTpt, List(
-            DefDef(nme.apply, applyTParams :: applyVParams :: Nil, res, EmptyTree).withFlags(Synthetic)
-          ))
-        }
-        else {
-          // Desugar [T_1, ..., T_M] -> (x_1: P_1, ..., x_N: P_N) => body
-          // with pt [S_1, ..., S_M] -> (O_1, ..., O_N) => R
-          // Into    new scala.PolyFunction { def apply[T_1, ..., T_M](x_1: P_1, ..., x_N: P_N): R2 = body }
-          // where R2 is R, with all references to S_1..S_M replaced with T1..T_M.
-
-          def typeTree(tp: Type) = tp match
-            case RefinedType(parent, nme.apply, poly @ PolyType(_, mt: MethodType)) if parent.classSymbol eq defn.PolyFunctionClass =>
-              untpd.DependentTypeTree((tsyms, vsyms) =>
-                mt.resultType.substParams(mt, vsyms.map(_.termRef)).substParams(poly, tsyms.map(_.typeRef)))
-            case _ => TypeTree()
-
-          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 :: applyVParams :: Nil, typeTree(pt), 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,
@@ -1815,8 +1802,6 @@ object desugar {
     }
 
     val desugared = tree match {
-      case PolyFunction(targs, body) =>
-        makePolyFunction(targs, body, pt) orElse tree
       case SymbolLit(str) =>
         Apply(
           ref(defn.ScalaSymbolClass.companionModule.termRef),

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

@@ -1872,6 +1872,8 @@ object Types {
         if alwaysDependent || mt.isResultDependent then
           RefinedType(funType, nme.apply, mt)
         else funType
+      case poly @ PolyType(_, mt: MethodType) if !mt.isParamDependent =>
+        RefinedType(defn.PolyFunctionType, nme.apply, poly)
     }
 
     /** The signature of this type. This is by default NotAMethod,

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

@@ -1625,12 +1625,32 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
                 )
               cpy.ValDef(param)(tpt = paramTpt)
           if isErased then param0.withAddedFlags(Flags.Erased) else param0
-      desugared = desugar.makeClosure(inferredParams, fnBody, resultTpt, tree.span)
+      desugared = desugar.makeClosure(Nil, inferredParams, fnBody, resultTpt, tree.span)
 
     typed(desugared, pt)
       .showing(i"desugared fun $tree --> $desugared with pt = $pt", typr)
   }
 
+
+  def typedPolyFunction(tree: untpd.PolyFunction, pt: Type)(using Context): Tree =
+    val tree1 = desugar.normalizePolyFunction(tree)
+    if (ctx.mode is Mode.Type) typed(desugar.makePolyFunctionType(tree1), pt)
+    else typedPolyFunctionValue(tree1, pt)
+
+  def typedPolyFunctionValue(tree: untpd.PolyFunction, pt: Type)(using Context): Tree =
+    val untpd.PolyFunction(tparams: List[untpd.TypeDef] @unchecked, fun) = tree: @unchecked
+    val untpd.Function(vparams: List[untpd.ValDef] @unchecked, body) = fun: @unchecked
+
+    val resultTpt = pt.dealias match
+      case RefinedType(parent, nme.apply, poly @ PolyType(_, mt: MethodType)) if parent.classSymbol eq defn.PolyFunctionClass =>
+        untpd.DependentTypeTree((tsyms, vsyms) =>
+          mt.resultType.substParams(mt, vsyms.map(_.termRef)).substParams(poly, tsyms.map(_.typeRef)))
+      case _ => untpd.TypeTree()
+
+    val desugared = desugar.makeClosure(tparams, vparams, body, resultTpt, tree.span)
+    typed(desugared, pt)
+  end typedPolyFunctionValue
+
   def typedClosure(tree: untpd.Closure, pt: Type)(using Context): Tree = {
     val env1 = tree.env mapconserve (typed(_))
     val meth1 = typedUnadapted(tree.meth)
@@ -1668,6 +1688,9 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
                 else
                   EmptyTree
             }
+          case _: PolyType =>
+            // Polymorphic SAMs are not currently supported (#6904).
+            EmptyTree
           case tp =>
             if !tp.isErroneous then
               throw new java.lang.Error(i"internal error: closing over non-method $tp, pos = ${tree.span}")
@@ -2425,7 +2448,7 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
       case rhs => typedExpr(rhs, tpt1.tpe.widenExpr)
     }
     val vdef1 = assignType(cpy.ValDef(vdef)(name, tpt1, rhs1), sym)
-    postProcessInfo(sym)
+    postProcessInfo(vdef1, sym)
     vdef1.setDefTree
   }
 
@@ -2534,19 +2557,31 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
 
     val ddef2 = assignType(cpy.DefDef(ddef)(name, paramss1, tpt1, rhs1), sym)
 
-    postProcessInfo(sym)
+    postProcessInfo(ddef2, sym)
     ddef2.setDefTree
       //todo: make sure dependent method types do not depend on implicits or by-name params
   }
 
   /** (1) Check that the signature of the class member does not return a repeated parameter type
    *  (2) If info is an erased class, set erased flag of member
+   *  (3) Check that erased classes are not parameters of polymorphic functions.
    */
-  private def postProcessInfo(sym: Symbol)(using Context): Unit =
+  private def postProcessInfo(mdef: MemberDef, sym: Symbol)(using Context): Unit =
     if (!sym.isOneOf(Synthetic | InlineProxy | Param) && sym.info.finalResultType.isRepeatedParam)
       report.error(em"Cannot return repeated parameter type ${sym.info.finalResultType}", sym.srcPos)
     if !sym.is(Module) && !sym.isConstructor && sym.info.finalResultType.isErasedClass then
       sym.setFlag(Erased)
+    if
+      sym.info.isInstanceOf[PolyType] &&
+      ((sym.name eq nme.ANON_FUN) ||
+       (sym.name eq nme.apply) && sym.owner.derivesFrom(defn.PolyFunctionClass))
+    then
+      mdef match
+        case DefDef(_, _ :: vparams :: Nil, _, _) =>
+          vparams.foreach: vparam =>
+            if vparam.symbol.is(Erased) then
+              report.error(em"Implementation restriction: erased classes are not allowed in a poly function definition", vparam.srcPos)
+        case _ =>
 
   def typedTypeDef(tdef: untpd.TypeDef, sym: Symbol)(using Context): Tree = {
     val TypeDef(name, rhs) = tdef
@@ -2693,19 +2728,6 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
       // check value class constraints
       checkDerivedValueClass(cls, body1)
 
-      // check PolyFunction constraints (no erased functions!)
-      if parents1.exists(_.tpe.classSymbol eq defn.PolyFunctionClass) then
-        body1.foreach {
-          case ddef: DefDef =>
-            ddef.paramss.foreach { params =>
-              val erasedParam = params.collectFirst { case vdef: ValDef if vdef.symbol.is(Erased) => vdef }
-              erasedParam.foreach { p =>
-                report.error(em"Implementation restriction: erased classes are not allowed in a poly function definition", p.srcPos)
-              }
-            }
-          case _ =>
-        }
-
       val effectiveOwner = cls.owner.skipWeakOwner
       if !cls.isRefinementClass
          && !cls.isAllOf(PrivateLocal)
@@ -3057,6 +3079,7 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
           case tree: untpd.Block => typedBlock(desugar.block(tree), pt)(using ctx.fresh.setNewScope)
           case tree: untpd.If => typedIf(tree, pt)
           case tree: untpd.Function => typedFunction(tree, pt)
+          case tree: untpd.PolyFunction => typedPolyFunction(tree, pt)
           case tree: untpd.Closure => typedClosure(tree, pt)
           case tree: untpd.Import => typedImport(tree)
           case tree: untpd.Export => typedExport(tree)

tests/neg/polymorphic-functions1.check

@@ -1,7 +1,7 @@
--- [E007] Type Mismatch Error: tests/neg/polymorphic-functions1.scala:1:53 ---------------------------------------------
+-- [E007] Type Mismatch Error: tests/neg/polymorphic-functions1.scala:1:33 ---------------------------------------------
 1 |val f: [T] => (x: T) => x.type = [T] => (x: Int) => x // error
-  |                                                     ^
-  |                                                     Found:    [T] => (x: Int) => x.type
-  |                                                     Required: [T] => (x: T) => x.type
+  |                                 ^^^^^^^^^^^^^^^^^^^^
+  |                                 Found:    [T] => (x: Int) => x.type
+  |                                 Required: [T] => (x: T) => x.type
   |
   | longer explanation available when compiling with `-explain`