SpecializeFunction1: don't roll over parents, use mapConserve

Author: felixmulder

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

@@ -256,7 +256,7 @@ object NameOps {
       case nme.clone_ => nme.clone_
     }
 
-    def specializedFor(classTargs: List[Types.Type], classTargsNames: List[Name], methodTargs: List[Types.Type] = scala.Nil, methodTarsNames: List[Name] = scala.Nil)(implicit ctx: Context): name.ThisName = {
+    def specializedFor(classTargs: List[Types.Type], classTargsNames: List[Name], methodTargs: List[Types.Type], methodTarsNames: List[Name])(implicit ctx: Context): name.ThisName = {
 
       def typeToTag(tp: Types.Type): Name = {
         tp.classSymbol match {

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

@@ -13,27 +13,119 @@ class SpecializeFunction1 extends MiniPhaseTransform with DenotTransformer {
 
   val phaseName = "specializeFunction1"
 
+  // Setup ---------------------------------------------------------------------
+  private[this] val functionName = "JFunction1".toTermName
+  private[this] val functionPkg  = "scala.compat.java8.".toTermName
+  private[this] var argTypes: Set[Symbol] = _
+  private[this] var retTypes: Set[Symbol] = _
+
+  override def prepareForUnit(tree: Tree)(implicit ctx: Context) = {
+    argTypes = Set(defn.DoubleClass,
+                   defn.FloatClass,
+                   defn.IntClass,
+                   defn.LongClass,
+                   defn.UnitClass,
+                   defn.BooleanClass)
+
+    retTypes = Set(defn.DoubleClass,
+                   defn.FloatClass,
+                   defn.IntClass,
+                   defn.LongClass)
+    this
+  }
+
+  // Transformations -----------------------------------------------------------
+
   /** Transforms all classes extending `Function1[-T1, +R]` so that
     * they instead extend the specialized version `JFunction$mp...`
     */
   def transform(ref: SingleDenotation)(implicit ctx: Context) = ref match {
-    case ShouldTransformDenot(cref, t1, r, func1) =>
-      transformDenot(cref, t1, r, func1)
+    case ShouldTransformDenot(cref, t1, r, func1) => {
+      val specializedFunction: Symbol =
+        ctx.getClassIfDefined(functionPkg ++ specializedName(functionName, t1, r))
+
+      def replaceFunction1(in: List[TypeRef]): List[TypeRef] =
+        in.mapConserve { tp =>
+          if (tp.isRef(defn.FunctionClass(1)) && (specializedFunction ne NoSymbol))
+            specializedFunction.typeRef
+          else tp
+        }
+
+      def specializeApply(scope: Scope): Scope =
+        if (specializedFunction ne NoSymbol) {
+          def specializedApply: Symbol = {
+            val specializedMethodName = specializedName(nme.apply, t1, r)
+            ctx.newSymbol(
+              cref.symbol,
+              specializedMethodName,
+              Flags.Override | Flags.Method,
+              specializedFunction.info.decls.lookup(specializedMethodName).info
+            )
+          }
+
+          val alteredScope = scope.cloneScope
+          alteredScope.enter(specializedApply)
+          alteredScope
+        }
+        else scope
+
+      val ClassInfo(prefix, cls, parents, decls, info) = cref.classInfo
+      val newInfo = ClassInfo(prefix, cls, replaceFunction1(in = parents), specializeApply(decls), info)
+      cref.copySymDenotation(info = newInfo)
+    }
     case _ => ref
   }
 
+  /** Transform the class definition's `Template`:
+    *
+    * - change the tree to have the correct parent
+    * - add the specialized apply method to the template body
+    * - forward the old `apply` to the specialized version
+    */
   override def transformTemplate(tree: Template)(implicit ctx: Context, info: TransformerInfo) =
     tree match {
-      case ShouldTransformTree(func1, t1, r) => transformTree(tree, func1, t1, r)
+      case tmpl @ ShouldTransformTree(func1, t1, r) => {
+        val specializedFunc1 =
+          TypeTree(ctx.requiredClassRef(functionPkg ++ specializedName(functionName, t1, r)))
+
+        val parents = tmpl.parents.foldRight(List.empty[Tree]) { (t, trees) =>
+          (if (func1 eq t) specializedFunc1 else t) :: trees
+        }
+
+        val body = tmpl.body.foldRight(List.empty[Tree]) {
+          case (tree: DefDef, acc) if tree.name == nme.apply => {
+            val specializedMethodName = specializedName(nme.apply, t1, r)
+            val specializedApply = ctx.owner.info.decls.lookup(specializedMethodName).asTerm
+
+            val forwardingBody =
+              tpd.ref(specializedApply)
+              .appliedToArgs(tree.vparamss.head.map(vparam => ref(vparam.symbol)))
+
+            val applyWithForwarding = cpy.DefDef(tree)(rhs = forwardingBody)
+
+            val specializedApplyDefDef = polyDefDef(specializedApply, trefs => vrefss => {
+              tree.rhs
+                .changeOwner(tree.symbol, specializedApply)
+                .subst(tree.vparamss.flatten.map(_.symbol), vrefss.flatten.map(_.symbol))
+            })
+
+            applyWithForwarding :: specializedApplyDefDef :: acc
+          }
+          case (tree, acc) => tree :: acc
+        }
+
+        cpy.Template(tmpl)(parents = parents, body = body)
+      }
       case _ => tree
     }
 
+  /** Dispatch to specialized `apply`s in user code */
   override def transformApply(tree: Apply)(implicit ctx: Context, info: TransformerInfo) = {
     import ast.Trees._
     tree match {
       case Apply(select @ Select(id, nme.apply), arg :: Nil) =>
         val params = List(arg.tpe, tree.tpe)
-        val specializedApply = nme.apply.specializedFor(params, params.map(_.typeSymbol.name))
+        val specializedApply = nme.apply.specializedFor(params, params.map(_.typeSymbol.name), Nil, Nil)
         val hasOverridenSpecializedApply = id.tpe.decls.iterator.exists { sym =>
           sym.is(Flags.Override) && (sym.name eq specializedApply)
         }
@@ -44,137 +136,32 @@ class SpecializeFunction1 extends MiniPhaseTransform with DenotTransformer {
     }
   }
 
-  private[this] val functionName = "JFunction1".toTermName
-  private[this] val functionPkg  = "scala.compat.java8.".toTermName
-  private[this] var Function1: Symbol = _
-  private[this] var argTypes: Set[Symbol] = _
-  private[this] var returnTypes: Set[Symbol] = _
-  private[this] var blacklisted: Set[Symbol] = _
-
-  /** Do setup of `argTypes` and `returnTypes` */
-  override def prepareForUnit(tree: Tree)(implicit ctx: Context) = {
-    argTypes = Set(defn.DoubleClass,
-                   defn.FloatClass,
-                   defn.IntClass,
-                   defn.LongClass,
-                   defn.UnitClass,
-                   defn.BooleanClass)
-
-    returnTypes = Set(defn.DoubleClass,
-                      defn.FloatClass,
-                      defn.IntClass,
-                      defn.LongClass)
-
-    Function1 = ctx.requiredClass("scala.Function1")
-
-    blacklisted = Set(
-      "scala.compat.java8.JFunction1",
-      "scala.runtime.AbstractFunction1",
-      "scala.PartialFunction",
-      "scala.runtime.AbstractPartialFunction"
-    ).map(ctx.requiredClass(_))
-
-    this
-  }
+  private def specializedName(name: Name, t1: Type, r: Type)(implicit ctx: Context): Name =
+    name.specializedFor(List(t1, r), List(t1, r).map(_.typeSymbol.name), Nil, Nil)
 
+  // Extractors ----------------------------------------------------------------
   private object ShouldTransformDenot {
-    def unapply(cref: ClassDenotation)(implicit ctx: Context): Option[(ClassDenotation, Type, Type, Type)] = {
-      def collectFunc1(xs: List[Type])(implicit ctx: Context): Option[(Type, Type, Type)] =
-        xs.collect {
-          case func1 @ RefinedType(RefinedType(parent, _, t1), _, r)
-            if func1.isRef(Function1) => (t1, r, func1)
-        }.headOption
-
-      collectFunc1(cref.info.parentsWithArgs).flatMap { case (t1, r, func1) =>
-        if (
-          !argTypes.contains(t1.typeSymbol) ||
-          !returnTypes.contains(r.typeSymbol) ||
-          blacklisted.exists(sym => cref.symbol.derivesFrom(sym))
-        ) None
-        else Some((cref, t1, r, func1))
+    def unapply(cref: ClassDenotation)(implicit ctx: Context): Option[(ClassDenotation, Type, Type, Type)] =
+      getFunc1(cref.typeRef).map {
+        case (t1, r, func1) => (cref, t1, r, func1)
       }
-    }
   }
 
   private object ShouldTransformTree {
     def unapply(tree: Template)(implicit ctx: Context): Option[(Tree, Type, Type)] =
       tree.parents
-      .map { t => (t.tpe, t) }
-      .collect {
-        case (tp @ RefinedType(RefinedType(parent, _, t1), _, r), func1)
-        if tp.isRef(Function1) &&
-           argTypes.contains(t1.typeSymbol) &&
-           returnTypes.contains(r.typeSymbol) => (func1, t1, r)
-      }
-      .headOption
-  }
-
-  private def specializedName(name: Name, t1: Type, r: Type)(implicit ctx: Context): Name =
-    name.specializedFor(List(t1, r), List(t1, r).map(_.typeSymbol.name))
-
-  def transformDenot(cref: ClassDenotation, t1: Type, r: Type, func1: Type)(implicit ctx: Context): SingleDenotation = {
-    val specializedFunction: TypeRef =
-      ctx.requiredClassRef(functionPkg ++ specializedName(functionName, t1, r))
-
-    def replaceFunction1(in: List[TypeRef]): List[TypeRef] =
-      in.foldRight(List.empty[TypeRef]) { (tp, acc) =>
-        val newTp =
-          if (tp.isRef(Function1)) specializedFunction
-          else tp
-        newTp :: acc
-      }
-
-    def specializeApply(scope: Scope): Scope = {
-      def specializedApply: Symbol = {
-        val specializedMethodName = specializedName(nme.apply, t1, r)
-        ctx.newSymbol(
-          cref.symbol,
-          specializedMethodName,
-          Flags.Override | Flags.Method,
-          specializedFunction.info.decls.lookup(specializedMethodName).info
-        )
-      }
-
-      val alteredScope = scope.cloneScope
-      alteredScope.enter(specializedApply)
-      alteredScope
-    }
-
-    val ClassInfo(prefix, cls, parents, decls, info) = cref.classInfo
-    val newInfo = ClassInfo(prefix, cls, replaceFunction1(in = parents), specializeApply(decls), info)
-    cref.copySymDenotation(info = newInfo)
+      .map(t => getFunc1(t.tpe).map { case (t1, r, _) => (t, t1, r) })
+      .flatten.headOption
   }
 
-  private def transformTree(tmpl: Template, func1: Tree, t1: Type, r: Type)(implicit ctx: Context) = {
-    val specializedFunc1 =
-      TypeTree(ctx.requiredClassRef(functionPkg ++ specializedName(functionName, t1, r)))
-
-    val parents = tmpl.parents.foldRight(List.empty[Tree]) { (t, trees) =>
-      (if (func1 eq t) specializedFunc1 else t) :: trees
-    }
-
-    val body = tmpl.body.foldRight(List.empty[Tree]) {
-      case (tree: DefDef, acc) if tree.name == nme.apply => {
-        val specializedMethodName = specializedName(nme.apply, t1, r)
-        val specializedApply = ctx.owner.info.decls.lookup(specializedMethodName).asTerm
-
-        val forwardingBody =
-          tpd.ref(specializedApply)
-          .appliedToArgs(tree.vparamss.head.map(vparam => ref(vparam.symbol)))
-
-        val applyWithForwarding = cpy.DefDef(tree)(rhs = forwardingBody)
-
-        val specializedApplyDefDef = polyDefDef(specializedApply, trefs => vrefss => {
-          tree.rhs
-            .changeOwner(tree.symbol, specializedApply)
-            .subst(tree.vparamss.flatten.map(_.symbol), vrefss.flatten.map(_.symbol))
-        })
-
-        applyWithForwarding :: specializedApplyDefDef :: acc
+  private def getFunc1(tpe: Type)(implicit ctx: Context): Option[(Type, Type, Type)] =
+    if (!tpe.derivesFrom(defn.FunctionClass(1)))
+      None
+    else
+      tpe.baseTypeWithArgs(defn.FunctionClass(1)) match {
+        case func1 @ RefinedType(RefinedType(parent, _, t1), _, r) if (
+          argTypes.contains(t1.typeSymbol) && retTypes.contains(r.typeSymbol)
+        ) => Some((t1, r, func1))
+        case _ => None
       }
-      case (tree, acc) => tree :: acc
-    }
-
-    cpy.Template(tmpl)(parents = parents, body = body)
-  }
 }