Value parameter inference for polymorphic lambdas

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
 }