Fix #4177: Generate optimised applyOrElse implementation for partial function literals

Author: allanrenucci

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

@@ -273,8 +273,12 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
         coord = fns.map(_.pos).reduceLeft(_ union _))
     val constr = ctx.newConstructor(cls, Synthetic, Nil, Nil).entered
     def forwarder(fn: TermSymbol, name: TermName) = {
-      val fwdMeth = fn.copy(cls, name, Synthetic | Method).entered.asTerm
-      DefDef(fwdMeth, prefss => ref(fn).appliedToArgss(prefss))
+      var flags = Synthetic | Method
+      def isOverriden(denot: SingleDenotation) = fn.info.overrides(denot.info, matchLoosely = true)
+      val isOverride = parents.exists(_.member(name).hasAltWith(isOverriden))
+      if (isOverride) flags = flags | Override
+      val fwdMeth = fn.copy(cls, name, flags).entered.asTerm
+      polyDefDef(fwdMeth, tprefs => prefss => ref(fn).appliedToTypes(tprefs).appliedToArgss(prefss))
     }
     val forwarders = (fns, methNames).zipped.map(forwarder)
     val cdef = ClassDef(cls, DefDef(constr), forwarders)

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

@@ -585,6 +585,10 @@ class Definitions {
 
   lazy val PartialFunctionType: TypeRef         = ctx.requiredClassRef("scala.PartialFunction")
   def PartialFunctionClass(implicit ctx: Context) = PartialFunctionType.symbol.asClass
+
+    lazy val PartialFunction_applyOrElseR = PartialFunctionClass.requiredMethodRef(nme.applyOrElse)
+    def PartialFunction_applyOrElse(implicit ctx: Context) = PartialFunction_applyOrElseR.symbol
+
   lazy val AbstractPartialFunctionType: TypeRef = ctx.requiredClassRef("scala.runtime.AbstractPartialFunction")
   def AbstractPartialFunctionClass(implicit ctx: Context) = AbstractPartialFunctionType.symbol.asClass
   lazy val FunctionXXLType: TypeRef         = ctx.requiredClassRef("scala.FunctionXXL")

compiler/src/dotty/tools/dotc/transform/ExpandSAMs.scala

@@ -13,7 +13,7 @@ import dotty.tools.dotc.util.Positions.Position
 /** Expand SAM closures that cannot be represented by the JVM as lambdas to anonymous classes.
  *  These fall into five categories
  *
- *   1. Partial function closures, we need to generate a isDefinedAt method for these.
+ *   1. Partial function closures, we need to generate isDefinedAt and applyOrElse methods for these.
  *   2. Closures implementing non-trait classes.
  *   3. Closures implementing classes that inherit from a class other than Object
  *      (a lambda cannot not be a run-time subtype of such a class)
@@ -54,38 +54,70 @@ class ExpandSAMs extends MiniPhase {
     val Block(
           (applyDef @ DefDef(nme.ANON_FUN, Nil, List(List(param)), _, _)) :: Nil,
           Closure(_, _, tpt)) = tree
-    val applyRhs: Tree = applyDef.rhs
+
+    def translateMatch(tree: Match, selector: Tree, cases: List[CaseDef], defaultValue: Tree) = {
+      assert(tree.selector.symbol == param.symbol)
+      val selectorTpe = selector.tpe.widen
+      val defaultSym = ctx.newSymbol(selector.symbol.owner, nme.WILDCARD, Synthetic, selectorTpe)
+      val defaultCase =
+        CaseDef(
+          Bind(defaultSym, Underscore(selectorTpe)),
+          EmptyTree,
+          defaultValue)
+      val unchecked = Annotated(selector, New(ref(defn.UncheckedAnnotType)))
+      cpy.Match(tree)(unchecked, cases :+ defaultCase)
+        .subst(param.symbol :: Nil, selector.symbol :: Nil)
+          // Needed because  a partial function can be written as:
+          // param => param match { case "foo" if foo(param) => param }
+          // And we need to update all references to 'param'
+    }
+
+    val applyRhs = applyDef.rhs
     val applyFn = applyDef.symbol.asTerm
 
     val MethodTpe(paramNames, paramTypes, _) = applyFn.info
     val isDefinedAtFn = applyFn.copy(
         name  = nme.isDefinedAt,
         flags = Synthetic | Method,
         info = MethodType(paramNames, paramTypes, defn.BooleanType)).asTerm
-    val tru = Literal(Constant(true))
-    def isDefinedAtRhs(paramRefss: List[List[Tree]]) = applyRhs match {
-      case Match(selector, cases) =>
-        assert(selector.symbol == param.symbol)
-        val paramRef = paramRefss.head.head
-        // Again, the alternative
-        //     val List(List(paramRef)) = paramRefs
-        // fails with a similar self instantiation error
-        def translateCase(cdef: CaseDef): CaseDef =
-          cpy.CaseDef(cdef)(body = tru).changeOwner(applyFn, isDefinedAtFn)
-        val defaultSym = ctx.newSymbol(isDefinedAtFn, nme.WILDCARD, Synthetic, selector.tpe.widen)
-        val defaultCase =
-          CaseDef(
-            Bind(defaultSym, Underscore(selector.tpe.widen)),
-            EmptyTree,
-            Literal(Constant(false)))
-        val annotated = Annotated(paramRef, New(ref(defn.UncheckedAnnotType)))
-        cpy.Match(applyRhs)(annotated, cases.map(translateCase) :+ defaultCase)
-      case _ =>
-        tru
+
+    val applyOrElseFn = applyFn.copy(
+        name  = nme.applyOrElse,
+        flags = Synthetic | Method,
+        info = tpt.tpe.memberInfo(defn.PartialFunction_applyOrElse)).asTerm
+
+    def isDefinedAtRhs(paramRefss: List[List[Tree]]) = {
+      val tru = Literal(Constant(true))
+      applyRhs match {
+        case tree @ Match(_, cases) =>
+          def translateCase(cdef: CaseDef)=
+            cpy.CaseDef(cdef)(body = tru).changeOwner(applyFn, isDefinedAtFn)
+          val paramRef = paramRefss.head.head
+          val defaultValue = Literal(Constant(false))
+          translateMatch(tree, paramRef, cases.map(translateCase), defaultValue)
+        case _ =>
+          tru
+      }
+    }
+
+    def applyOrElseRhs(paramRefss: List[List[Tree]]) = {
+      val List(paramRef, defaultRef) = paramRefss.head
+      applyRhs match {
+        case tree @ Match(_, cases) =>
+          def translateCase(cdef: CaseDef) =
+            cdef.changeOwner(applyFn, applyOrElseFn)
+          val defaultValue = defaultRef.select(nme.apply).appliedTo(paramRef)
+          translateMatch(tree, paramRef, cases.map(translateCase), defaultValue)
+        case _ =>
+          ref(applyFn).appliedTo(paramRef)
+        }
     }
+
     val isDefinedAtDef = transformFollowingDeep(DefDef(isDefinedAtFn, isDefinedAtRhs(_)))
-    val anonCls = AnonClass(tpt.tpe :: Nil, List(applyFn, isDefinedAtFn), List(nme.apply, nme.isDefinedAt))
-    cpy.Block(tree)(List(applyDef, isDefinedAtDef), anonCls)
+    val applyOrElseDef = transformFollowingDeep(DefDef(applyOrElseFn, applyOrElseRhs(_)))
+
+    val anonCls = AnonClass(tpt.tpe :: Nil, List(applyFn, isDefinedAtFn, applyOrElseFn), List(nme.apply, nme.isDefinedAt, nme.applyOrElse))
+    cpy.Block(tree)(List(applyDef, isDefinedAtDef, applyOrElseDef), anonCls)
   }
 
   private def checkRefinements(tpe: Type, pos: Position)(implicit ctx: Context): Type = tpe match {

tests/pos/i4177.scala

@@ -0,0 +1,15 @@
+class Test {
+
+  object Foo { def unapply(x: Int) = if (x == 2) Some(x.toString) else None }
+
+  def test: Unit = {
+    val a: PartialFunction[Int, String] = { case Foo(x) => x }
+    val b: PartialFunction[Int, String] = { case x => x.toString }
+    val c: PartialFunction[Int, String] = { x => x.toString }
+    val d: PartialFunction[Int, String] = x => x.toString
+
+    val e: PartialFunction[String, String] = { case x @ "abc" => x }
+    val f: PartialFunction[String, String] = x => x match { case "abc" => x }
+    val g: PartialFunction[String, String] = x => x match { case "abc" if x.isEmpty => x }
+  }
+}

tests/run/i4177.scala

@@ -0,0 +1,18 @@
+object Test {
+  private[this] var count = 0
+
+  def test(x: Int) = { count += 1; true }
+
+  object Foo {
+    def unapply(x: Int): Option[Int] = { count += 1; Some(x) }
+  }
+
+  def main(args: Array[String]): Unit = {
+    val res = List(1, 2).collect { case x if test(x) => x }
+    assert(count == 2)
+
+    count = 0
+    val res2 = List(1, 2).collect { case Foo(x) => x }
+    assert(count == 2)
+  }
+}