Skip to content

Commit ef74f4d

Browse files
hamzaremmalhamzaremmal
committed
Add support for companion in MacroAnnotations
1 parent 54d67e0 commit ef74f4d

File tree

68 files changed

+561
-302
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

68 files changed

+561
-302
lines changed

‎compiler/src/dotty/tools/dotc/transform/Inlining.scala

Lines changed: 132 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,22 @@
11
package dotty.tools.dotc
22
package transform
33

4+
import ast.tpd
5+
import ast.Trees.*
6+
import ast.TreeMapWithImplicits
47
import core.*
58
import Flags.*
9+
import Decorators.*
610
import Contexts.*
711
import Symbols.*
12+
import Decorators.*
13+
import config.Printers.inlining
14+
import DenotTransformers.IdentityDenotTransformer
15+
import inlines.Inlines
16+
import quoted.*
17+
import staging.StagingLevel
818

9-
import dotty.tools.dotc.ast.tpd
10-
import dotty.tools.dotc.ast.Trees.*
11-
import dotty.tools.dotc.quoted.*
12-
import dotty.tools.dotc.inlines.Inlines
13-
import dotty.tools.dotc.ast.TreeMapWithImplicits
14-
import dotty.tools.dotc.core.DenotTransformers.IdentityDenotTransformer
15-
import dotty.tools.dotc.staging.StagingLevel
16-
17-
import scala.collection.mutable.ListBuffer
19+
import scala.collection.mutable
1820

1921
/** Inlines all calls to inline methods that are not in an inline method or a quote */
2022
class Inlining extends MacroTransform, IdentityDenotTransformer {
@@ -56,38 +58,82 @@ class Inlining extends MacroTransform, IdentityDenotTransformer {
5658

5759
def newTransformer(using Context): Transformer = new Transformer {
5860
override def transform(tree: tpd.Tree)(using Context): tpd.Tree =
59-
new InliningTreeMap().transform(tree)
61+
InliningTreeMap().transform(tree)
62+
}
63+
64+
private class MemoizeStatsTreeMap extends TreeMapWithImplicits {
65+
66+
// It is safe to assume that the companion of a tree is in the same scope
67+
// Therefore, when expanding MacroAnnotations, we will only keep track of
68+
// the trees in the same scope as the current transformed tree
69+
70+
override def transform(tree: Tree)(using Context): Tree =
71+
tree match
72+
case PackageDef(_, stats) =>
73+
// Phase I: Collect and memoize all the stats
74+
val treesToTrack = stats.collect { case m: MemberDef => (m.symbol, m) }
75+
val withTrackedTreeCtx = MacroAnnotations.trackedTreesCtx(mutable.Map(treesToTrack*))
76+
// Phase II & III: Transform the tree with this definitions and reconcile them with the tracked trees
77+
super.transform(tree)(using withTrackedTreeCtx) match
78+
case pkg@PackageDef(pid, stats) =>
79+
val trackedTree = MacroAnnotations.trackedTrees(using withTrackedTreeCtx)
80+
val updatedStats = stats.mapConserve:
81+
case tree: MemberDef if trackedTree.contains(tree.symbol) =>
82+
trackedTree(tree.symbol)
83+
case stat => stat
84+
cpy.PackageDef(pkg)(pid = pid, stats = updatedStats)
85+
case tree => tree
86+
case block: Block =>
87+
// Phase I: Fetch all the member definitions in the block
88+
val trackedTrees = block.stats.collect { case m: MemberDef => (m.symbol, m) }
89+
val withTrackedTreeCtx = MacroAnnotations.trackedTreesCtx(mutable.Map(trackedTrees*))
90+
91+
// Phase II / III: Transform the tree and Reconcile between the symbols in syms and the tree
92+
// TODO: Should we define a substitution method where we change the trees
93+
// and not the symbols (see Tree::subst)
94+
// result.subst(MacroAnnotations.trackedTrees(using withTrackedTreeCtx))
95+
super.transform(tree)(using withTrackedTreeCtx) match
96+
case b@Block(stats, expr) =>
97+
val trackedTree = MacroAnnotations.trackedTrees(using withTrackedTreeCtx)
98+
cpy.Block(b)(
99+
expr = expr,
100+
stats = stats.mapConserve:
101+
case ddef: MemberDef if trackedTree.contains(ddef.symbol) =>
102+
trackedTree(ddef.symbol)
103+
case stat => stat
104+
)
105+
case tree => tree
106+
case TypeDef(_, impl: Template) =>
107+
// Phase I: Collect and memoize all the stats
108+
val treesToTrack = impl.body.collect { case m: MemberDef => (m.symbol, m) }
109+
val withTrackedTreeCtx = MacroAnnotations.trackedTreesCtx(mutable.Map(treesToTrack*))
110+
// Phase II / III: Transform the tree and Reconcile between the symbols in syms and the tree
111+
super.transform(tree)(using withTrackedTreeCtx) match
112+
case tree@TypeDef(name, impl: Template) =>
113+
val trackedTree = MacroAnnotations.trackedTrees(using withTrackedTreeCtx)
114+
cpy.TypeDef(tree)(
115+
name = name,
116+
rhs = cpy.Template(impl)(
117+
body = impl.body.mapConserve:
118+
case ddef: MemberDef if trackedTree.contains(ddef.symbol) =>
119+
trackedTree(ddef.symbol)
120+
case stat => stat
121+
)
122+
)
123+
case tree => tree
124+
case _ => super.transform(tree)
60125
}
61126

62-
private class InliningTreeMap extends TreeMapWithImplicits {
127+
private class InliningTreeMap extends MemoizeStatsTreeMap {
63128

64129
/** List of top level classes added by macro annotation in a package object.
65130
* These are added to the PackageDef that owns this particular package object.
66131
*/
67-
private val newTopClasses = MutableSymbolMap[ListBuffer[Tree]]()
132+
private val newTopClasses = MutableSymbolMap[mutable.ListBuffer[Tree]]()
68133

69134
override def transform(tree: Tree)(using Context): Tree = {
70135
tree match
71-
case tree: MemberDef =>
72-
if tree.symbol.is(Inline) then tree
73-
else if tree.symbol.is(Param) then super.transform(tree)
74-
else if
75-
!tree.symbol.isPrimaryConstructor
76-
&& StagingLevel.level == 0
77-
&& MacroAnnotations.hasMacroAnnotation(tree.symbol)
78-
then
79-
val trees = (new MacroAnnotations(self)).expandAnnotations(tree)
80-
val trees1 = trees.map(super.transform)
81-
82-
// Find classes added to the top level from a package object
83-
val (topClasses, trees2) =
84-
if ctx.owner.isPackageObject then trees1.partition(_.symbol.owner == ctx.owner.owner)
85-
else (Nil, trees1)
86-
if topClasses.nonEmpty then
87-
newTopClasses.getOrElseUpdate(ctx.owner.owner, new ListBuffer) ++= topClasses
88-
89-
flatTree(trees2)
90-
else super.transform(tree)
136+
case tree: MemberDef => transformMemberDef(tree)
91137
case _: Typed | _: Block =>
92138
super.transform(tree)
93139
case _: PackageDef =>
@@ -113,7 +159,61 @@ class Inlining extends MacroTransform, IdentityDenotTransformer {
113159
else Inlines.inlineCall(tree1)
114160
else super.transform(tree)
115161
}
162+
163+
private def transformMemberDef(tree: MemberDef)(using Context) : Tree =
164+
if tree.symbol.is(Inline) then tree
165+
else if tree.symbol.is(Param) then
166+
super.transform(tree)
167+
else if
168+
!tree.symbol.isPrimaryConstructor
169+
&& StagingLevel.level == 0
170+
&& MacroAnnotations.hasMacroAnnotation(tree.symbol)
171+
then
172+
// Fetch the companion's tree
173+
val companionSym =
174+
if tree.symbol.is(ModuleClass) then tree.symbol.companionClass
175+
else if tree.symbol.is(ModuleVal) then NoSymbol
176+
else tree.symbol.companionModule.moduleClass
177+
178+
val companionTree = MacroAnnotations.findTrackedTree(companionSym)
179+
180+
// Fetch the latest tracked tree (It might have already been processed by its companion)
181+
val latestTree = MacroAnnotations.findTrackedTree(tree.symbol)
182+
.getOrElse(tree)
183+
184+
// Expand and process MacroAnnotations
185+
val (trees, companion) =
186+
MacroAnnotations(self).expandAnnotations(latestTree, companionTree)
187+
188+
// Update the tracked trees
189+
for case tree : MemberDef <- trees do
190+
MacroAnnotations.updateTrackedTree(tree.symbol, tree)
191+
for tree <- companion do
192+
MacroAnnotations.updateTrackedTree(tree.symbol, tree)
193+
194+
// Perform inlining on the expansion of the annotations
195+
val trees1 = trees.map(super.transform)
196+
197+
for case tree: MemberDef <- trees1 do
198+
MacroAnnotations.updateTrackedTree(tree.symbol, tree)
199+
200+
// Find classes added to the top level from a package object
201+
val (topClasses, trees2) =
202+
if ctx.owner.isPackageObject then trees1.partition(_.symbol.owner == ctx.owner.owner)
203+
else (Nil, trees1)
204+
if topClasses.nonEmpty then
205+
newTopClasses.getOrElseUpdate(ctx.owner.owner, new mutable.ListBuffer) ++= topClasses
206+
flatTree(trees2)
207+
else
208+
super.transform(tree) match
209+
case tree: MemberDef =>
210+
MacroAnnotations.updateTrackedTree(tree.symbol, tree)
211+
tree
212+
case tree => tree
213+
end transformMemberDef
214+
116215
}
216+
117217
}
118218

119219
object Inlining:

0 commit comments

Comments
 (0)