Merge pull request scala#6003 from szeiger/fix/pf-overloads2

Two fixes for overload resolution for PartialFunctions

Author: szeiger

src/compiler/scala/tools/nsc/typechecker/Typers.scala

@@ -3341,7 +3341,7 @@ trait Typers extends Adaptations with Tags with TypersTracking with PatternTyper
 
             def funArgTypes(tpAlts: List[(Type, Symbol)]) = tpAlts.map { case (tp, alt) =>
               val relTp = tp.asSeenFrom(pre, alt.owner)
-              val argTps = functionOrSamArgTypes(relTp)
+              val argTps = functionOrPfOrSamArgTypes(relTp)
               //println(s"funArgTypes $argTps from $relTp")
               argTps.map(approximateAbstracts)
             }
@@ -3350,6 +3350,10 @@ trait Typers extends Adaptations with Tags with TypersTracking with PatternTyper
               try functionType(funArgTypes(argTpWithAlt).transpose.map(lub), WildcardType)
               catch { case _: IllegalArgumentException => WildcardType }
 
+            def partialFunctionProto(argTpWithAlt: List[(Type, Symbol)]): Type =
+              try appliedType(PartialFunctionClass, funArgTypes(argTpWithAlt).transpose.map(lub) :+ WildcardType)
+              catch { case _: IllegalArgumentException => WildcardType }
+
             // To propagate as much information as possible to typedFunction, which uses the expected type to
             // infer missing parameter types for Function trees that we're typing as arguments here,
             // we expand the parameter types for all alternatives to the expected argument length,
@@ -3360,7 +3364,7 @@ trait Typers extends Adaptations with Tags with TypersTracking with PatternTyper
             // and lubbing the argument types (we treat SAM and FunctionN types equally, but non-function arguments
             // do not receive special treatment: they are typed under WildcardType.)
             val altArgPts =
-              if (settings.isScala212 && args.exists(treeInfo.isFunctionMissingParamType))
+              if (settings.isScala212 && args.exists(t => treeInfo.isFunctionMissingParamType(t) || treeInfo.isPartialFunctionMissingParamType(t)))
                 try alts.map(alt => formalTypes(alt.info.paramTypes, argslen).map(ft => (ft, alt))).transpose // do least amount of work up front
                 catch { case _: IllegalArgumentException => args.map(_ => Nil) } // fail safe in case formalTypes fails to align to argslen
               else args.map(_ => Nil) // will type under argPt == WildcardType
@@ -3375,8 +3379,12 @@ trait Typers extends Adaptations with Tags with TypersTracking with PatternTyper
                   // the overloaded type into a single function type from which `typedFunction`
                   // can derive the argument type for `x` in the function literal above
                   val argPt =
-                    if (argPtAlts.nonEmpty && treeInfo.isFunctionMissingParamType(tree)) functionProto(argPtAlts)
-                    else WildcardType
+                    if (argPtAlts.isEmpty) WildcardType
+                    else if (treeInfo.isFunctionMissingParamType(tree)) functionProto(argPtAlts)
+                    else if (treeInfo.isPartialFunctionMissingParamType(tree)) {
+                      if (argPtAlts.exists(ts => isPartialFunctionType(ts._1))) partialFunctionProto(argPtAlts)
+                      else functionProto(argPtAlts)
+                    } else WildcardType
 
                   val argTyped = typedArg(tree, amode, BYVALmode, argPt)
                   (argTyped, argTyped.tpe.deconst)

src/reflect/scala/reflect/internal/Definitions.scala

@@ -691,11 +691,11 @@ trait Definitions extends api.StandardDefinitions {
       }
     }
 
-    // the argument types expected by the function described by `tp` (a FunctionN or SAM type),
-    // or `Nil` if `tp` does not represent a function type or SAM (or if it happens to be Function0...)
-    def functionOrSamArgTypes(tp: Type): List[Type] = {
+    // the argument types expected by the function described by `tp` (a FunctionN or PartialFunction or SAM type),
+    // or `Nil` if `tp` does not represent a function type or PartialFunction or SAM (or if it happens to be Function0...)
+    def functionOrPfOrSamArgTypes(tp: Type): List[Type] = {
       val dealiased = tp.dealiasWiden
-      if (isFunctionTypeDirect(dealiased)) dealiased.typeArgs.init
+      if (isFunctionTypeDirect(dealiased) || isPartialFunctionType(dealiased)) dealiased.typeArgs.init
       else samOf(tp) match {
         case samSym if samSym.exists => tp.memberInfo(samSym).paramTypes
         case _ => Nil

src/reflect/scala/reflect/internal/TreeInfo.scala

@@ -265,6 +265,10 @@ abstract class TreeInfo {
 
   def isFunctionMissingParamType(tree: Tree): Boolean = tree match {
     case Function(vparams, _) => vparams.exists(_.tpt.isEmpty)
+    case _ => false
+  }
+
+  def isPartialFunctionMissingParamType(tree: Tree): Boolean = tree match {
     case Match(EmptyTree, _) => true
     case _ => false
   }

test/files/run/InferOverloadedPartialFunction.scala

@@ -44,4 +44,30 @@ object Test extends App {
   def h[R](pf: Function2[Int, String, R]): Int = 1
   def h[R](pf: PartialFunction[(Double, Double), R]): Int = 2
   assert(h { case (a: Double, b: Double) => 42: Int } == 2)
+
+  val xs = new SortedMap
+  assert(xs.collectF { kv => 1 } == 0)
+  assert(xs.collectF { case (k, v) => 1 } == 0)
+  assert(xs.collectF { case (k, v) => (1, 1) } == 2)
+  assert(xs.collect { case (k, v) => 1 } == 0)
+  assert(xs.collect { case (k, v) => (1, 1) } == 1)
+
+  val ys = new SortedMapMixed
+  assert(ys.collect { kv => 1 } == 0)
+  assert(ys.collect { kv => (1, 1) } == 0)
+  assert(ys.collect { case (k, v) => 1 } == 1) // could be 0 with the extra work in https://github.com/scala/scala/pull/5975/commits/3c95dac0dcbb0c8eb4686264026ad9c86b2022de
+  assert(ys.collect { case (k, v) => (1, 1) } == 2)
+}
+
+class SortedMap {
+  def collect[B](pf: PartialFunction[(String, Int), B]): Int = 0
+  def collect[K2 : Ordering, V2](pf: PartialFunction[(String, Int), (K2, V2)]): Int = 1
+  def collectF[B](pf: Function1[(String, Int), B]): Int = if(pf.isInstanceOf[PartialFunction[_, _]]) 1 else 0
+  def collectF[K2 : Ordering, V2](pf: Function1[(String, Int), (K2, V2)]): Int = if(pf.isInstanceOf[PartialFunction[_, _]]) 3 else 2
+}
+
+class SortedMapMixed {
+  type PF[-A, +B] = PartialFunction[A, B]
+  def collect[B](pf: Function1[(String, Int), B]): Int = if(pf.isInstanceOf[PartialFunction[_, _]]) 1 else 0
+  def collect[K2 : Ordering, V2](pf: PF[(String, Int), (K2, V2)]): Int = 2
 }