Code refactoring

Author: liufengyun

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

@@ -1436,6 +1436,31 @@ class Definitions {
         false
     })
 
+  @tu lazy val Function0SpecializedApplyNames: PerRun[collection.Set[TermName]] =
+    new PerRun({
+      for r <- defn.Function0SpecializedReturnTypes
+      yield nme.apply.specializedFunction(r, Nil).asTermName
+    })
+
+  @tu lazy val Function1SpecializedApplyNames: PerRun[collection.Set[TermName]] =
+    new PerRun({
+      for
+        r  <- Function1SpecializedReturnTypes
+        t1 <- Function1SpecializedParamTypes
+      yield
+        nme.apply.specializedFunction(r, List(t1)).asTermName
+    })
+
+  @tu lazy val Function2SpecializedApplyNames: PerRun[collection.Set[TermName]] =
+    new PerRun({
+      for
+        r  <- Function2SpecializedReturnTypes
+        t1 <- Function2SpecializedParamTypes
+        t2 <- Function2SpecializedReturnTypes
+      yield
+        nme.apply.specializedFunction(r, List(t1, t2)).asTermName
+    })
+
   def functionArity(tp: Type)(using Context): Int = tp.dropDependentRefinement.dealias.argInfos.length - 1
 
   /** Return underlying context function type (i.e. instance of an ContextFunctionN class)

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

@@ -6,6 +6,9 @@ import Contexts.Context, Types._, Decorators._, Symbols._, DenotTransformers._
 import SymDenotations._, Scopes._, StdNames._, NameOps._, Names._
 import MegaPhase.MiniPhase
 
+import scala.collection.mutable
+
+
 /** This phase synthesizes specialized methods for FunctionN, this is done
  *  since there are no scala signatures in the bytecode for the specialized
  *  methods.
@@ -19,75 +22,95 @@ class SpecializedApplyMethods extends MiniPhase with InfoTransformer {
 
   val phaseName = "specializedApplyMethods"
 
-  private[this] var func0Applys: collection.Set[Symbol] = _
-  private[this] var func1Applys: collection.Set[Symbol] = _
-  private[this] var func2Applys: collection.Set[Symbol] = _
-  private[this] var func0: Symbol = _
-  private[this] var func1: Symbol = _
-  private[this] var func2: Symbol = _
-
-  private def init()(using Context): Unit = if (func0Applys eq null) {
-    def specApply(sym: Symbol, args: List[Type], ret: Type)(using Context): Symbol = {
-      val name = nme.apply.specializedFunction(ret, args)
-      newSymbol(sym, name, Flags.Method, MethodType(args, ret))
-    }
+  private def specApply(sym: Symbol, args: List[Type], ret: Type)(using Context): Symbol = {
+    val name = nme.apply.specializedFunction(ret, args)
+    newSymbol(sym, name, Flags.Method, MethodType(args, ret))
+  }
 
-    func0 = defn.FunctionClass(0)
-    func0Applys = for (r <- defn.Function0SpecializedReturnTypes) yield specApply(func0, Nil, r)
+  private def specFun0(op: Type => Unit)(using Context): Unit = {
+    for (r <- defn.Function0SpecializedReturnTypes) do
+      op(r)
+  }
 
-    func1 = defn.FunctionClass(1)
-    func1Applys = for {
+  private def specFun1(op: (Type, Type) => Unit)(using Context): Unit = {
+    for
       r  <- defn.Function1SpecializedReturnTypes
       t1 <- defn.Function1SpecializedParamTypes
-    } yield specApply(func1, List(t1), r)
+    do
+      op(t1, r)
+  }
 
-    func2 = defn.FunctionClass(2)
-    func2Applys = for {
+  private def specFun2(op: (Type, Type, Type) => Unit)(using Context): Unit = {
+    for
       r  <- defn.Function2SpecializedReturnTypes
       t1 <- defn.Function2SpecializedParamTypes
       t2 <- defn.Function2SpecializedReturnTypes
-    } yield specApply(func2, List(t1, t2), r)
+    do
+      op(t1, t2, r)
   }
 
   /** Add symbols for specialized methods to FunctionN */
   override def transformInfo(tp: Type, sym: Symbol)(using Context) = tp match {
     case tp: ClassInfo if defn.isPlainFunctionClass(sym) => {
-      init()
-      val newDecls = sym.name.functionArity match {
-        case 0 => func0Applys.foldLeft(tp.decls.cloneScope) {
-          (decls, sym) => decls.enter(sym); decls
-        }
-        case 1 => func1Applys.foldLeft(tp.decls.cloneScope) {
-          (decls, sym) => decls.enter(sym); decls
-        }
-        case 2 => func2Applys.foldLeft(tp.decls.cloneScope) {
-          (decls, sym) => decls.enter(sym); decls
-        }
-        case _ => tp.decls
-      }
+      sym.name.functionArity match {
+        case 0 =>
+          val scope = tp.decls.cloneScope
+          specFun0 { r => scope.enter(specApply(sym, Nil, r)) }
+          tp.derivedClassInfo(decls = scope)
+
+        case 1 =>
+          val scope = tp.decls.cloneScope
+          specFun1 { (t1, r) => scope.enter(specApply(sym, List(t1), r)) }
+          tp.derivedClassInfo(decls = scope)
 
-      tp.derivedClassInfo(decls = newDecls)
+        case 2 =>
+          val scope = tp.decls.cloneScope
+          specFun2 { (t1, t2, r) => scope.enter(specApply(sym, List(t1, t2), r)) }
+          tp.derivedClassInfo(decls = scope)
+
+        case _ =>
+          tp
+      }
     }
     case _ => tp
   }
 
   /** Create bridge methods for FunctionN with specialized applys */
   override def transformTemplate(tree: Template)(using Context) = {
-    val owner = tree.symbol.owner
-    val additionalSymbols =
-      if (owner eq func0) func0Applys
-      else if (owner eq func1) func1Applys
-      else if (owner eq func2) func2Applys
-      else Nil
-
-    if (additionalSymbols eq Nil) tree
-    else cpy.Template(tree)(body = tree.body ++ additionalSymbols.map { apply =>
-      DefDef(apply.asTerm, { vparamss =>
-        This(owner.asClass)
-          .select(nme.apply)
-          .appliedToArgss(vparamss)
-          .ensureConforms(apply.info.finalResultType)
-      })
-    })
+    val cls = tree.symbol.owner.asInstanceOf[ClassSymbol]
+
+    if (!defn.isPlainFunctionClass(cls)) return tree
+
+    def synthesizeApply(names: collection.Set[TermName]): Tree = {
+      val applyBuf = new mutable.ListBuffer[DefDef]
+      names.foreach { name =>
+        val applySym = cls.info.decls.lookup(name)
+        val ddef = DefDef(
+          applySym.asTerm,
+          { vparamss =>
+              This(cls)
+                .select(nme.apply)
+                .appliedToArgss(vparamss)
+                .ensureConforms(applySym.info.finalResultType)
+          }
+        )
+        applyBuf += ddef
+      }
+      cpy.Template(tree)(body = tree.body ++ applyBuf)
+    }
+
+    cls.name.functionArity match {
+      case 0 =>
+        synthesizeApply(defn.Function0SpecializedApplyNames())
+
+      case 1 =>
+        synthesizeApply(defn.Function1SpecializedApplyNames())
+
+      case 2 =>
+        synthesizeApply(defn.Function2SpecializedApplyNames())
+
+      case _ =>
+        tree
+    }
   }
 }