Address review

Author: liufengyun

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

@@ -28,7 +28,7 @@ class SpecializeFunctions extends MiniPhase with InfoTransformer {
       val apply = tp.decls.lookup(nme.apply)
       if (!apply.exists) return tp
 
-      val newApplys = new mutable.ListBuffer[Symbol]
+      var newApplys: List[Symbol] = Nil
 
       var arity = 0
       while (arity < 3) {
@@ -39,14 +39,15 @@ class SpecializeFunctions extends MiniPhase with InfoTransformer {
           val argTypes   = paramTypes.init
           val retType    = paramTypes.last
 
+          val isSpecializable = defn.isSpecializableFunction(sym.asClass, argTypes, retType)
+          if (!isSpecializable) return tp
+
           val apply = baseType.member(nme.apply)
           val overrideApply = tp.decls.find {  sym =>
             sym.is(Flags.Method, butNot = Flags.Deferred) && sym.name == nme.apply && sym.info.matches(apply.info)
           }
 
-          def isSpecializable = defn.isSpecializableFunction(sym.asClass, argTypes, retType)
-
-          if (overrideApply.exists && isSpecializable) {
+          if (overrideApply.exists) {
             val specializedMethodName: Name = nme.apply.specializedFunction(retType, argTypes)
             val applySpecialized = newSymbol(
               sym,
@@ -55,7 +56,7 @@ class SpecializeFunctions extends MiniPhase with InfoTransformer {
               overrideApply.info
             )
 
-            newApplys += applySpecialized
+            newApplys ::= applySpecialized
           }
         }
         arity += 1
@@ -71,49 +72,40 @@ class SpecializeFunctions extends MiniPhase with InfoTransformer {
     case _ => tp
   }
 
-  /** Transforms the `Template` of the classes to contain forwarders from the
-   *  generic applys to the specialized ones. Also inserts the specialized applys
-   *  in the template body.
+  /** Create forwarders from the generic applys to the specialized ones.
    */
-  override def transformTemplate(tree: Template)(using Context) = {
-    val cls = tree.symbol.owner.asClass
-
-    val apply = cls.info.decls.lookup(nme.apply)
-    if (!apply.exists || !derivesFromFn012(cls)) return tree
-
-    val applyBuf = new mutable.ListBuffer[Tree]
-    val newBody = tree.body.mapConserve {
-      case ddef: DefDef
-        if ddef.name == nme.apply &&
-          ddef.vparamss.length == 1 &&
-          ddef.vparamss.head.length < 3
-      =>
-        val paramTypes = ddef.vparamss.head.map(_.symbol.info)
-        val retType = ddef.tpe.widen.finalResultType
-
-        val specName = nme.apply.specializedFunction(retType, paramTypes)
-        val specializedApply = cls.info.decls.lookup(specName)
-
-        if (specializedApply.exists) {
-          val specializedDecl =
-            DefDef(specializedApply.asTerm, vparamss => {
-              ddef.rhs
-                .changeOwner(ddef.symbol, specializedApply)
-                .subst(ddef.vparamss.head.map(_.symbol), vparamss.head.map(_.symbol))
-            })
-          applyBuf += specializedDecl
-
-          // create a forwarding to the specialized apply
-          val args = ddef.vparamss.head.map(vparam => ref(vparam.symbol))
-          val rhs = This(cls).select(specializedApply).appliedToArgs(args)
-          cpy.DefDef(ddef)(rhs = rhs)
-        } else ddef
-
-      case x => x
-    }
-
-    if applyBuf.isEmpty then tree
-    else cpy.Template(tree)(body = applyBuf.toList ::: newBody)
+  override def transformDefDef(ddef: DefDef)(using Context) = {
+    val sym = ddef.symbol
+
+    if ddef.name != nme.apply
+       || ddef.vparamss.length != 1
+       || ddef.vparamss.head.length > 2
+       || !sym.owner.isClass
+       || !derivesFromFn012(sym.owner.asClass)
+    then
+      return ddef
+
+    val cls = sym.owner.asClass
+    val paramTypes = ddef.vparamss.head.map(_.symbol.info)
+    val retType = ddef.tpe.widen.finalResultType
+
+    val specName = nme.apply.specializedFunction(retType, paramTypes)
+    val specializedApply = cls.asClass.info.decls.lookup(specName)
+
+    if (!specializedApply.exists) return ddef
+
+    val specializedDecl =
+      DefDef(specializedApply.asTerm, vparamss => {
+        ddef.rhs
+          .changeOwner(ddef.symbol, specializedApply)
+          .subst(ddef.vparamss.head.map(_.symbol), vparamss.head.map(_.symbol))
+      })
+
+    // create a forwarding to the specialized apply
+    val args = ddef.vparamss.head.map(vparam => ref(vparam.symbol))
+    val rhs = This(cls).select(specializedApply).appliedToArgs(args)
+    val ddef1 = cpy.DefDef(ddef)(rhs = rhs)
+    Thicket(ddef1, specializedDecl)
   }
 
   /** Dispatch to specialized `apply`s in user code when available */