Value parameter inference for polymorphic lambdas

Just like

    val f: Int => Int = x => x

is inferred to

    val f: Int => Int = (x: Int) => x

we now accept

    val g: [T] => T => T = [S] => x => x

which is inferred to

    val g: [T] => T => T = [S] => (x: S) => x

This requires a substitution step which is tricky to do since we're operating
with untyped trees at this point. We implement this by generalizing the existing
`DependentTypeTree` mechanism (already used for computing dependent result types
of lambdas) to also be usable in any other position inside the lambda.

We rename the tree to `InLambdaTypeTree` at the same time for clarity.

Note that this mechanism could also probably be used to allow closures with
internal value parameter dependencies, but we don't attempt to support this
here.

Author: smarter

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

@@ -150,8 +150,15 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   /** {x1, ..., xN} T   (only relevant under captureChecking) */
   case class CapturesAndResult(refs: List[Tree], parent: Tree)(implicit @constructorOnly src: SourceFile) extends TypTree
 
-  /** Short-lived usage in typer, does not need copy/transform/fold infrastructure */
-  case class DependentTypeTree(tp: (List[TypeSymbol], List[TermSymbol]) => Type)(implicit @constructorOnly src: SourceFile) extends Tree
+  /** A type tree appearing somewhere in the untyped DefDef of a lambda, it will be typed using `tpFun`.
+   *
+   *  @param isResult  Is this the result type of the lambda? This is handled specially in `Namer#valOrDefDefSig`.
+   *  @param tpFun     Compute the type of the type tree given the parameters of the lambda.
+   *                   A lambda has at most one type parameter list followed by exactly one term parameter list.
+   *
+   *  Note: This is only used briefly in Typer and does not need the copy/transform/fold infrastructure.
+   */
+  case class InLambdaTypeTree(isResult: Boolean, tpFun: (List[TypeSymbol], List[TermSymbol]) => Type)(implicit @constructorOnly src: SourceFile) extends Tree
 
   @sharable object EmptyTypeIdent extends Ident(tpnme.EMPTY)(NoSource) with WithoutTypeOrPos[Untyped] {
     override def isEmpty: Boolean = true

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

@@ -1698,13 +1698,16 @@ class Namer { typer: Typer =>
         WildcardType
       case TypeTree() =>
         checkMembersOK(inferredType, mdef.srcPos)
-      case DependentTypeTree(tpFun) =>
+
+      // We cannot rely on `typedInLambdaTypeTree` since the computed type might not be fully-defined.
+      case InLambdaTypeTree(/*isResult =*/ true, tpFun) =>
         // A lambda has at most one type parameter list followed by exactly one term parameter list.
         val tpe = (paramss: @unchecked) match
           case TypeSymbols(tparams) :: TermSymbols(vparams) :: Nil => tpFun(tparams, vparams)
           case TermSymbols(vparams) :: Nil => tpFun(Nil, vparams)
         if (isFullyDefined(tpe, ForceDegree.none)) tpe
         else typedAheadExpr(mdef.rhs, tpe).tpe
+
       case TypedSplice(tpt: TypeTree) if !isFullyDefined(tpt.tpe, ForceDegree.none) =>
         mdef match {
           case mdef: DefDef if mdef.name == nme.ANON_FUN =>

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

@@ -1323,14 +1323,14 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
           (pt1.argInfos.init, typeTree(interpolateWildcards(pt1.argInfos.last.hiBound)))
         case RefinedType(parent, nme.apply, mt @ MethodTpe(_, formals, restpe))
         if (defn.isNonRefinedFunction(parent) || defn.isErasedFunctionType(parent)) && formals.length == defaultArity =>
-          (formals, untpd.DependentTypeTree((_, syms) => restpe.substParams(mt, syms.map(_.termRef))))
+          (formals, untpd.InLambdaTypeTree(isResult = true, (_, syms) => restpe.substParams(mt, syms.map(_.termRef))))
         case pt1 @ SAMType(mt @ MethodTpe(_, formals, _)) if !SAMType.isParamDependentRec(mt) =>
           val restpe = mt.resultType match
             case mt: MethodType => mt.toFunctionType(isJava = pt1.classSymbol.is(JavaDefined))
             case tp => tp
           (formals,
            if (mt.isResultDependent)
-             untpd.DependentTypeTree((_, syms) => restpe.substParams(mt, syms.map(_.termRef)))
+             untpd.InLambdaTypeTree(isResult = true, (_, syms) => restpe.substParams(mt, syms.map(_.termRef)))
            else
              typeTree(restpe))
         case _ =>
@@ -1641,13 +1641,34 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
     val untpd.PolyFunction(tparams: List[untpd.TypeDef] @unchecked, fun) = tree: @unchecked
     val untpd.Function(vparams: List[untpd.ValDef] @unchecked, body) = fun: @unchecked
 
+    // If the expected type is a polymorphic function with the same number of
+    // type and value parameters, then infer the types of value parameters from the expected type.
+    val inferredVParams = pt match
+      case RefinedType(parent, nme.apply, poly @ PolyType(_, mt: MethodType))
+      if (parent.typeSymbol eq defn.PolyFunctionClass)
+      && tparams.lengthCompare(poly.paramNames) == 0
+      && vparams.lengthCompare(mt.paramNames) == 0
+      =>
+        vparams.zipWithConserve(mt.paramInfos): (vparam, formal) =>
+          // Unlike in typedFunctionValue, `formal` cannot be a TypeBounds since
+          // it must be a valid method parameter type.
+          if vparam.tpt.isEmpty && isFullyDefined(formal, ForceDegree.failBottom) then
+            cpy.ValDef(vparam)(tpt = new untpd.InLambdaTypeTree(isResult = false, (tsyms, vsyms) =>
+              // We don't need to substitute `mt` by `vsyms` because we currently disallow
+              // dependencies between value parameters of a closure.
+              formal.substParams(poly, tsyms.map(_.typeRef)))
+            )
+          else vparam
+      case _ =>
+        vparams
+
     val resultTpt = pt.dealias match
       case RefinedType(parent, nme.apply, poly @ PolyType(_, mt: MethodType)) if parent.classSymbol eq defn.PolyFunctionClass =>
-        untpd.DependentTypeTree((tsyms, vsyms) =>
+        untpd.InLambdaTypeTree(isResult = true, (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)
+    val desugared = desugar.makeClosure(tparams, inferredVParams, body, resultTpt, tree.span)
     typed(desugared, pt)
   end typedPolyFunctionValue
 
@@ -2098,6 +2119,18 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
       case _ =>
         completeTypeTree(InferredTypeTree(), pt, tree)
 
+  def typedInLambdaTypeTree(tree: untpd.InLambdaTypeTree, pt: Type)(using Context): Tree =
+    val tp =
+      if tree.isResult then pt // See InLambdaTypeTree logic in Namer#valOrDefDefSig.
+      else
+        val lambdaCtx = ctx.outersIterator.dropWhile(_.owner.name ne nme.ANON_FUN).next()
+        // A lambda has at most one type parameter list followed by exactly one term parameter list.
+        // Parameters are entered in order in the scope of the lambda.
+        val (tsyms: List[TypeSymbol @unchecked], vsyms: List[TermSymbol @unchecked]) =
+          lambdaCtx.scope.toList.partition(_.isType): @unchecked
+        tree.tpFun(tsyms, vsyms)
+    completeTypeTree(InferredTypeTree(), tp, tree)
+
   def typedSingletonTypeTree(tree: untpd.SingletonTypeTree)(using Context): SingletonTypeTree = {
     val ref1 = typedExpr(tree.ref, SingletonTypeProto)
     checkStable(ref1.tpe, tree.srcPos, "singleton type")
@@ -3109,7 +3142,7 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
           case tree: untpd.TypedSplice => typedTypedSplice(tree)
           case tree: untpd.UnApply => typedUnApply(tree, pt)
           case tree: untpd.Tuple => typedTuple(tree, pt)
-          case tree: untpd.DependentTypeTree => completeTypeTree(untpd.InferredTypeTree(), pt, tree)
+          case tree: untpd.InLambdaTypeTree => typedInLambdaTypeTree(tree, pt)
           case tree: untpd.InfixOp => typedInfixOp(tree, pt)
           case tree: untpd.ParsedTry => typedTry(tree, pt)
           case tree @ untpd.PostfixOp(qual, Ident(nme.WILDCARD)) => typedAsFunction(tree, pt)

tests/neg/polymorphic-functions2.scala

@@ -0,0 +1,6 @@
+val wrongLength1: [T, S] => (T, S) => T = [T] => (x, y) => x // error
+val wrongLength2: [T] => T => T = [T] => (x, x) => x // error
+
+val notSubType: [T] => T => T = [T <: Int] => x => x // error
+
+val notInScope: [T] => T => T = [S] => x => (x: T) // error

tests/run/polymorphic-functions.scala

@@ -106,4 +106,16 @@ object Test extends App {
   val tt2: [T] =>  T => T =  [T] => ((x: T) => x)
   val tt3: [T] =>  T => T =  [T] => { (x: T) => x }
   val tt4: [T] =>  T => T =  [T] => (x: T) => { x }
+
+  // Inferred parameter type
+  val i1a: [T] => T => T = [T] => x => x
+  val i2b: [T] => T => T = [S] => x => x
+  /// This does not work currently because subtyping of polymorphic functions is not implemented.
+  /// val i2c: [T <: Int] => T => T = [T] => x => x
+  val i3a: [T, S <: List[T]] => (T, S) => List[T] =
+    [T, S <: List[T]] => (x, y) => x :: y
+  val i3b: [T, S <: List[T]] => (T, S) => List[T] =
+    [S, T <: List[S]] => (x, y) => x :: y
+  val i4: [T, S <: List[T]] => (T, S) => List[T] =
+    [T, S <: List[T]] => (x, y: S) => x :: y
 }