Support inferring dependent polymorphic lambdas from the expected type

Reuse the existing `DependentTypeTree` mechanism already in place for
monomorphic lambdas to compute the result type of polymorphic lambdas based on
their expected type.

Author: smarter

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

@@ -1761,15 +1761,9 @@ object desugar {
           // 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, PolyType(_, mt)) if parent.typeSymbol eq defn.PolyFunctionClass =>
-              var bail = false
-              def mapper(tp: Type, topLevel: Boolean = false): Tree = tp match
-                case tp: TypeRef              => ref(tp)
-                case tp: TypeParamRef         => Ident(applyTParams(tp.paramNum).name)
-                case AppliedType(tycon, args) => AppliedTypeTree(mapper(tycon), args.map(mapper(_)))
-                case _                        => if topLevel then TypeTree() else { bail = true; genericEmptyTree }
-              val mapped = mapper(mt.resultType, topLevel = true)
-              if bail then TypeTree() else mapped
+            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]]

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

@@ -151,7 +151,7 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   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[Symbol] => Type)(implicit @constructorOnly src: SourceFile) extends Tree
+  case class DependentTypeTree(tp: (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

@@ -1692,15 +1692,17 @@ class Namer { typer: Typer =>
   def valOrDefDefSig(mdef: ValOrDefDef, sym: Symbol, paramss: List[List[Symbol]], paramFn: Type => Type)(using Context): Type = {
 
     def inferredType = inferredResultType(mdef, sym, paramss, paramFn, WildcardType)
-    lazy val termParamss = paramss.collect { case TermSymbols(vparams) => vparams }
 
     val tptProto = mdef.tpt match {
       case _: untpd.DerivedTypeTree =>
         WildcardType
       case TypeTree() =>
         checkMembersOK(inferredType, mdef.srcPos)
       case DependentTypeTree(tpFun) =>
-        val tpe = tpFun(termParamss.head)
+        // 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) =>
@@ -1724,7 +1726,8 @@ class Namer { typer: Typer =>
             // So fixing levels at instantiation avoids the soundness problem but apparently leads
             // to type inference problems since it comes too late.
             if !Config.checkLevelsOnConstraints then
-              val hygienicType = TypeOps.avoid(rhsType, termParamss.flatten)
+              val termParams = paramss.collect { case TermSymbols(vparams) => vparams }.flatten
+              val hygienicType = TypeOps.avoid(rhsType, termParams)
               if (!hygienicType.isValueType || !(hygienicType <:< tpt.tpe))
                 report.error(
                   em"""return type ${tpt.tpe} of lambda cannot be made hygienic

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.DependentTypeTree((_, 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.DependentTypeTree((_, syms) => restpe.substParams(mt, syms.map(_.termRef)))
            else
              typeTree(restpe))
         case _ =>

tests/neg/polymorphic-functions1.check

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

tests/pos/i16756.scala

@@ -0,0 +1,16 @@
+class DependentPoly {
+
+  sealed trait Col[V] {
+
+    trait Wrapper
+    val wrapper: Wrapper = ???
+  }
+
+  object Col1 extends Col[Int]
+
+  object Col2 extends Col[Double]
+
+  val polyFn: [C <: DependentPoly.this.Col[?]] => (x: C) => x.Wrapper =
+    [C <: Col[?]] => (x: C) => (x.wrapper: x.Wrapper)
+}
+

tests/run/polymorphic-functions.scala

@@ -85,6 +85,13 @@ object Test extends App {
   val v0a: String = v0
   assert(v0 == "foo")
 
+  // Used to fail with:    Found: ... => List[T]
+  //                    Expected: ... => List[x.type]
+  val md2: [T] => (x: T) => List[x.type] = [T] => (x: T) => List(x)
+  val x = 1
+  val v1 = md2(x)
+  val v1a: List[x.type] = v1
+
   // Contextual
   trait Show[T] { def show(t: T): String }
   implicit val si: Show[Int] =