Merge pull request #2868 from dotty-staging/fix-#2867

Fix #2866: Handle partial functions in preTypeArgs

Author: odersky

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

@@ -273,17 +273,24 @@ trait UntypedTreeInfo extends TreeInfo[Untyped] { self: Trees.Instance[Untyped]
     case _ => false
   }
 
-  def isFunctionWithUnknownParamType(tree: Tree) = tree match {
+  def functionWithUnknownParamType(tree: Tree): Option[Tree] = tree match {
     case Function(args, _) =>
-      args.exists {
+      if (args.exists {
         case ValDef(_, tpt, _) => tpt.isEmpty
         case _ => false
-      }
+      }) Some(tree)
+      else None
     case Match(EmptyTree, _) =>
-      true
-    case _ => false
+      Some(tree)
+    case Block(Nil, expr) =>
+      functionWithUnknownParamType(expr)
+    case _ =>
+      None
   }
 
+  def isFunctionWithUnknownParamType(tree: Tree): Boolean =
+    functionWithUnknownParamType(tree).isDefined
+
   /** Is `tree` an implicit function or closure, possibly nested in a block? */
   def isImplicitClosure(tree: Tree)(implicit ctx: Context): Boolean = unsplice(tree) match {
     case Function((param: untpd.ValDef) :: _, _) => param.mods.is(Implicit)

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

@@ -704,6 +704,18 @@ class Definitions {
     }
   }
 
+  object PartialFunctionOf {
+    def apply(arg: Type, result: Type)(implicit ctx: Context) =
+      PartialFunctionType.appliedTo(arg :: result :: Nil)
+    def unapply(pft: Type)(implicit ctx: Context) = {
+      if (pft.isRef(PartialFunctionClass)) {
+        val targs = pft.dealias.argInfos
+        if (targs.length == 2) Some((targs.head, targs.tail)) else None
+      }
+      else None
+    }
+  }
+
   object ArrayOf {
     def apply(elem: Type)(implicit ctx: Context) =
       if (ctx.erasedTypes) JavaArrayType(elem)

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

@@ -1379,21 +1379,19 @@ trait Applications extends Compatibility { self: Typer with Dynamic =>
   }
 
   /** Try to typecheck any arguments in `pt` that are function values missing a
-   *  parameter type. The expected type for these arguments is the lub of the
-   *  corresponding formal parameter types of all alternatives. Type variables
-   *  in formal parameter types are replaced by wildcards. The result of the
-   *  typecheck is stored in `pt`, to be retrieved when its `typedArgs` are selected.
+   *  parameter type. If the formal parameter types corresponding to a closure argument
+   *  all agree on their argument types, typecheck the argument with an expected
+   *  function or partial function type that contains these argument types,
+   *  The result of the typecheck is stored in `pt`, to be retrieved when its `typedArgs` are selected.
    *  The benefit of doing this is to allow idioms like this:
    *
    *     def map(f: Char => Char): String = ???
    *     def map[U](f: Char => U): Seq[U] = ???
    *     map(x => x.toUpper)
    *
    *  Without `pretypeArgs` we'd get a "missing parameter type" error for `x`.
-   *  With `pretypeArgs`, we use the union of the two formal parameter types
-   *  `Char => Char` and `Char => ?` as the expected type of the closure `x => x.toUpper`.
-   *  That union is `Char => Char`, so we have an expected parameter type `Char`
-   *  for `x`, and the code typechecks.
+   *  With `pretypeArgs`, we use the `Char => ?` as the expected type of the
+   *  closure `x => x.toUpper`, which makes the code typecheck.
    */
   private def pretypeArgs(alts: List[TermRef], pt: FunProto)(implicit ctx: Context): Unit = {
     def recur(altFormals: List[List[Type]], args: List[untpd.Tree]): Unit = args match {
@@ -1402,30 +1400,35 @@ trait Applications extends Compatibility { self: Typer with Dynamic =>
           case ValDef(_, tpt, _) => tpt.isEmpty
           case _ => false
         }
-        if (untpd.isFunctionWithUnknownParamType(arg)) {
+        val fn = untpd.functionWithUnknownParamType(arg)
+        if (fn.isDefined) {
           def isUniform[T](xs: List[T])(p: (T, T) => Boolean) = xs.forall(p(_, xs.head))
           val formalsForArg: List[Type] = altFormals.map(_.head)
-          // For alternatives alt_1, ..., alt_n, test whether formal types for current argument are of the form
-          //   (p_1_1, ..., p_m_1) => r_1
-          //   ...
-          //   (p_1_n, ..., p_m_n) => r_n
-          val decomposedFormalsForArg: List[Option[(List[Type], Type, Boolean)]] =
-            formalsForArg.map(defn.FunctionOf.unapply)
-          if (decomposedFormalsForArg.forall(_.isDefined)) {
-            val formalParamTypessForArg: List[List[Type]] =
-              decomposedFormalsForArg.map(_.get._1)
-            if (isUniform(formalParamTypessForArg)((x, y) => x.length == y.length)) {
-              val commonParamTypes = formalParamTypessForArg.transpose.map(ps =>
-                // Given definitions above, for i = 1,...,m,
-                //   ps(i) = List(p_i_1, ..., p_i_n)  -- i.e. a column
-                // If all p_i_k's are the same, assume the type as formal parameter
-                // type of the i'th parameter of the closure.
-                if (isUniform(ps)(ctx.typeComparer.isSameTypeWhenFrozen(_, _))) ps.head
-                else WildcardType)
-              val commonFormal = defn.FunctionOf(commonParamTypes, WildcardType)
-              overload.println(i"pretype arg $arg with expected type $commonFormal")
-              pt.typedArg(arg, commonFormal)(ctx.addMode(Mode.ImplicitsEnabled))
+          def argTypesOfFormal(formal: Type): List[Type] =
+            formal match {
+              case defn.FunctionOf(args, result, isImplicit) => args
+              case defn.PartialFunctionOf(arg, result) => arg :: Nil
+              case _ => Nil
             }
+          val formalParamTypessForArg: List[List[Type]] =
+            formalsForArg.map(argTypesOfFormal)
+          if (formalParamTypessForArg.forall(_.nonEmpty) &&
+              isUniform(formalParamTypessForArg)((x, y) => x.length == y.length)) {
+            val commonParamTypes = formalParamTypessForArg.transpose.map(ps =>
+              // Given definitions above, for i = 1,...,m,
+              //   ps(i) = List(p_i_1, ..., p_i_n)  -- i.e. a column
+              // If all p_i_k's are the same, assume the type as formal parameter
+              // type of the i'th parameter of the closure.
+              if (isUniform(ps)(ctx.typeComparer.isSameTypeWhenFrozen(_, _))) ps.head
+              else WildcardType)
+            def isPartial = // we should generate a partial function for the arg
+              fn.get.isInstanceOf[untpd.Match] &&
+              formalsForArg.exists(_.isRef(defn.PartialFunctionClass))
+            val commonFormal =
+              if (isPartial) defn.PartialFunctionOf(commonParamTypes.head, newTypeVar(TypeBounds.empty))
+              else defn.FunctionOf(commonParamTypes, WildcardType)
+            overload.println(i"pretype arg $arg with expected type $commonFormal")
+            pt.typedArg(arg, commonFormal)(ctx.addMode(Mode.ImplicitsEnabled))
           }
         }
         recur(altFormals.map(_.tail), args1)

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

@@ -401,15 +401,18 @@ object ProtoTypes {
   /**  Same as `constrained(tl, EmptyTree)`, but returns just the created type lambda */
   def constrained(tl: TypeLambda)(implicit ctx: Context): TypeLambda = constrained(tl, EmptyTree)._1
 
-  /** Create a new TypeVar that represents a dependent method parameter singleton */
-  def newDepTypeVar(tp: Type)(implicit ctx: Context): TypeVar = {
+  def newTypeVar(bounds: TypeBounds)(implicit ctx: Context): TypeVar = {
     val poly = PolyType(DepParamName.fresh().toTypeName :: Nil)(
-        pt => TypeBounds.upper(AndType(tp, defn.SingletonClass.typeRef)) :: Nil,
+        pt => bounds :: Nil,
         pt => defn.AnyType)
     constrained(poly, untpd.EmptyTree, alwaysAddTypeVars = true)
       ._2.head.tpe.asInstanceOf[TypeVar]
   }
 
+  /** Create a new TypeVar that represents a dependent method parameter singleton */
+  def newDepTypeVar(tp: Type)(implicit ctx: Context): TypeVar =
+    newTypeVar(TypeBounds.upper(AndType(tp, defn.SingletonClass.typeRef)))
+
   /** The result type of `mt`, where all references to parameters of `mt` are
    *  replaced by either wildcards (if typevarsMissContext) or TypeParamRefs.
    */