Improve type inference for dependent function types

Given a dependently typed function value like this one:

    def f: (x: C) => D => x.T => E

we did not propagate information about the subsequent types `D` and `x.T` to
the result type of the closure with parameter `(x: C)`. Doing so is a bit
tricky because of the dependency. But it's necessary to infer the
types of subsequent parameters.

Test case: eff-dependent.scala

Author: odersky

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

@@ -86,6 +86,7 @@ object untpd extends Trees.Instance[Untyped] with UntypedTreeInfo {
   case class GenAlias(pat: Tree, expr: Tree) extends Tree
   case class ContextBounds(bounds: TypeBoundsTree, cxBounds: List[Tree]) extends TypTree
   case class PatDef(mods: Modifiers, pats: List[Tree], tpt: Tree, rhs: Tree) extends DefTree
+  case class DependentTypeTree(tp: List[Symbol] => Type) extends Tree
 
   @sharable object EmptyTypeIdent extends Ident(tpnme.EMPTY) with WithoutTypeOrPos[Untyped] {
     override def isEmpty = true

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

@@ -1096,6 +1096,8 @@ class Namer { typer: Typer =>
         WildcardType
       case TypeTree() =>
         inferredType
+      case DependentTypeTree(tpFun) =>
+        tpFun(paramss.head)
       case TypedSplice(tpt: TypeTree) if !isFullyDefined(tpt.tpe, ForceDegree.none) =>
         val rhsType = typedAheadExpr(mdef.rhs, tpt.tpe).tpe
         mdef match {

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

@@ -683,18 +683,37 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
     assignType(cpy.If(tree)(cond1, thenp2, elsep2), thenp2, elsep2)
   }
 
-  private def decomposeProtoFunction(pt: Type, defaultArity: Int)(implicit ctx: Context): (List[Type], Type) = pt match {
-    case _ if defn.isNonDepFunctionType(pt) =>
-      // if expected parameter type(s) are wildcards, approximate from below.
-      // if expected result type is a wildcard, approximate from above.
-      // this can type the greatest set of admissible closures.
-      val funType = pt.dealias
-      (funType.argTypesLo.init, funType.argTypesHi.last)
-    case SAMType(meth) =>
-      val mt @ MethodTpe(_, formals, restpe) = meth.info
-      (formals, if (mt.isDependent) WildcardType else restpe)
-    case _ =>
-      (List.tabulate(defaultArity)(alwaysWildcardType), WildcardType)
+  /** Decompose function prototype into a list of parameter prototypes and a result prototype
+   *  tree, using WildcardTypes where a type is not known.
+   *  For the result type we do this even if the expected type is not fully
+   *  defined, which is a bit of a hack. But it's needed to make the following work
+   *  (see typers.scala and printers/PlainPrinter.scala for examples).
+   *
+   *     def double(x: Char): String = s"$x$x"
+   *     "abc" flatMap double
+   */
+  private def decomposeProtoFunction(pt: Type, defaultArity: Int)(implicit ctx: Context): (List[Type], untpd.Tree) = {
+    def typeTree(tp: Type) = tp match {
+      case _: WildcardType => untpd.TypeTree()
+      case _ => untpd.TypeTree(tp)
+    }
+    pt match {
+      case _ if defn.isNonDepFunctionType(pt) =>
+        // if expected parameter type(s) are wildcards, approximate from below.
+        // if expected result type is a wildcard, approximate from above.
+        // this can type the greatest set of admissible closures.
+        val funType = pt.dealias
+        (funType.argTypesLo.init, typeTree(funType.argTypesHi.last))
+      case SAMType(meth) =>
+        val mt @ MethodTpe(_, formals, restpe) = meth.info
+        (formals,
+         if (mt.isDependent)
+           untpd.DependentTypeTree(syms => restpe.substParams(mt, syms.map(_.termRef)))
+         else
+           typeTree(restpe))
+      case _ =>
+        (List.tabulate(defaultArity)(alwaysWildcardType), untpd.TypeTree())
+    }
   }
 
   def typedFunction(tree: untpd.Function, pt: Type)(implicit ctx: Context) = track("typedFunction") {
@@ -756,7 +775,7 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
       case _ =>
     }
 
-    val (protoFormals, protoResult) = decomposeProtoFunction(pt, params.length)
+    val (protoFormals, resultTpt) = decomposeProtoFunction(pt, params.length)
 
     def refersTo(arg: untpd.Tree, param: untpd.ValDef): Boolean = arg match {
       case Ident(name) => name == param.name
@@ -865,19 +884,6 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
             else cpy.ValDef(param)(
               tpt = untpd.TypeTree(
                 inferredParamType(param, protoFormal(i)).underlyingIfRepeated(isJava = false)))
-
-        // Define result type of closure as the expected type, thereby pushing
-        // down any implicit searches. We do this even if the expected type is not fully
-        // defined, which is a bit of a hack. But it's needed to make the following work
-        // (see typers.scala and printers/PlainPrinter.scala for examples).
-        //
-        //     def double(x: Char): String = s"$x$x"
-        //     "abc" flatMap double
-        //
-        val resultTpt = protoResult match {
-          case WildcardType(_) => untpd.TypeTree()
-          case _ => untpd.TypeTree(protoResult)
-        }
         val inlineable = pt.hasAnnotation(defn.InlineParamAnnot)
         desugar.makeClosure(inferredParams, fnBody, resultTpt, inlineable)
       }
@@ -1700,7 +1706,8 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
           case tree: untpd.PackageDef => typedPackageDef(tree)
           case tree: untpd.Annotated => typedAnnotated(tree, pt)
           case tree: untpd.TypedSplice => typedTypedSplice(tree)
-          case tree:  untpd.UnApply => typedUnApply(tree, pt)
+          case tree: untpd.UnApply => typedUnApply(tree, pt)
+          case tree: untpd.DependentTypeTree => typed(untpd.TypeTree().withPos(tree.pos), pt)
           case tree @ untpd.PostfixOp(qual, Ident(nme.WILDCARD)) => typedAsFunction(tree, pt)
           case untpd.EmptyTree => tpd.EmptyTree
           case _ => typedUnadapted(desugar(tree), pt)

tests/run/eff-dependent.scala

@@ -0,0 +1,38 @@
+object Test extends App {
+
+  trait Effect
+
+  // Type X => Y
+  abstract class Fun[-X, +Y] {
+    type Eff <: Effect
+    def apply(x: X): implicit Eff => Y
+  }
+
+  class CanThrow extends Effect
+  class CanIO extends Effect
+
+  val i2s = new Fun[Int, String] { type Eff = CanThrow; def apply(x: Int) = x.toString }
+  val s2i = new Fun[String, Int] { type Eff = CanIO; def apply(x: String) = x.length }
+
+  implicit val ct: CanThrow = new CanThrow
+  implicit val ci: CanIO = new CanIO
+
+  // def map(f: A => B)(xs: List[A]): List[B]
+  def map[A, B](f: Fun[A, B])(xs: List[A]): implicit f.Eff => List[B] =
+    xs.map(f.apply)
+
+  // def mapFn[A, B]: (A => B) -> List[A] -> List[B]
+  def mapFn[A, B]: (f: Fun[A, B]) => List[A] => implicit f.Eff => List[B] =
+    f => xs => map(f)(xs)
+
+  // def compose(f: A => B)(g: B => C)(x: A): C
+  def compose[A, B, C](f: Fun[A, B])(g: Fun[B, C])(x: A): implicit f.Eff => implicit g.Eff => C = g(f(x))
+
+  // def composeFn: (A => B) -> (B => C) -> A -> C
+  def composeFn[A, B, C]: (f: Fun[A, B]) => (g: Fun[B, C]) => A => implicit f.Eff => implicit g.Eff => C =
+    f => g => x => compose(f)(g)(x)
+
+  assert(mapFn(i2s)(List(1, 2, 3)).mkString == "123")
+  assert(composeFn(i2s)(s2i)(22) == 2)
+
+}