Skip to content

Commit b562ff4

Browse files
committed
Try use transformTemplate
1 parent 989aaf9 commit b562ff4

File tree

1 file changed

+49
-38
lines changed

1 file changed

+49
-38
lines changed

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

Lines changed: 49 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -21,44 +21,55 @@ class SpecializeFunctions extends MiniPhase {
2121

2222
/** Create forwarders from the generic applys to the specialized ones.
2323
*/
24-
override def transformDefDef(ddef: DefDef)(using Context) = {
25-
val sym = ddef.symbol
26-
27-
if ddef.name != nme.apply
28-
|| sym.is(Flags.Deferred)
29-
|| ddef.vparamss.length != 1
30-
|| ddef.vparamss.head.length > 2
31-
|| !sym.owner.isClass
32-
then
33-
return ddef
34-
35-
val cls = sym.owner.asClass
36-
val paramTypes = ddef.vparamss.head.map(_.symbol.info)
37-
val retType = sym.info.finalResultType
38-
val specName = nme.apply.specializedFunction(retType, paramTypes)
39-
40-
val isSpecializable = defn.isSpecializableFunction(cls, paramTypes, retType)
41-
if (!isSpecializable || cls.info.decls.lookup(specName).exists) return ddef
42-
43-
val specializedApply = newSymbol(
44-
cls,
45-
specName,
46-
sym.flags | Flags.Synthetic,
47-
sym.info
48-
).entered
49-
50-
val specializedDecl =
51-
DefDef(specializedApply.asTerm, vparamss => {
52-
ddef.rhs
53-
.changeOwner(ddef.symbol, specializedApply)
54-
.subst(ddef.vparamss.head.map(_.symbol), vparamss.head.map(_.symbol))
55-
})
56-
57-
// create a forwarding to the specialized apply
58-
val args = ddef.vparamss.head.map(vparam => ref(vparam.symbol))
59-
val rhs = This(cls).select(specializedApply).appliedToArgs(args)
60-
val ddef1 = cpy.DefDef(ddef)(rhs = rhs)
61-
Thicket(ddef1, specializedDecl)
24+
override def transformTemplate(tree: Template)(using Context) = {
25+
val cls = tree.symbol.owner.asClass
26+
27+
val apply = cls.info.decls.lookup(nme.apply)
28+
if (!apply.exists || !derivesFromFn012(cls)) return tree
29+
30+
var applyBuf: List[Tree] = Nil
31+
val newBody = tree.body.mapConserve {
32+
case ddef: DefDef
33+
if ddef.name == nme.apply &&
34+
ddef.vparamss.length == 1 &&
35+
ddef.vparamss.head.length < 3 &&
36+
!ddef.symbol.is(Flags.Deferred)
37+
=>
38+
val sym = ddef.symbol
39+
val paramTypes = ddef.vparamss.head.map(_.symbol.info)
40+
val retType = sym.info.finalResultType
41+
42+
val isSpecializable = defn.isSpecializableFunction(cls, paramTypes, retType)
43+
if (isSpecializable) {
44+
val specName = nme.apply.specializedFunction(retType, paramTypes)
45+
val specializedApply = newSymbol(
46+
cls,
47+
specName,
48+
sym.flags | Flags.Synthetic,
49+
sym.info
50+
).entered
51+
52+
val specializedDecl =
53+
DefDef(specializedApply.asTerm, vparamss => {
54+
ddef.rhs
55+
.changeOwner(ddef.symbol, specializedApply)
56+
.subst(ddef.vparamss.head.map(_.symbol), vparamss.head.map(_.symbol))
57+
})
58+
59+
applyBuf ::= specializedDecl
60+
61+
// create a forwarding to the specialized apply
62+
val args = ddef.vparamss.head.map(vparam => ref(vparam.symbol))
63+
val rhs = This(cls).select(specializedApply).appliedToArgs(args)
64+
cpy.DefDef(ddef)(rhs = rhs)
65+
}
66+
else ddef
67+
68+
case x => x
69+
}
70+
71+
if applyBuf.isEmpty then tree
72+
else cpy.Template(tree)(body = applyBuf ::: newBody)
6273
}
6374

6475
/** Dispatch to specialized `apply`s in user code when available */

0 commit comments

Comments
 (0)