Skip to content

Commit d711455

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

File tree

71 files changed

+598
-307
lines changed

Some content is hidden

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

71 files changed

+598
-307
lines changed

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

+142-32
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,23 @@
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
18+
import util.Property
819

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
20+
import scala.collection.mutable
1821

1922
/** Inlines all calls to inline methods that are not in an inline method or a quote */
2023
class Inlining extends MacroTransform, IdentityDenotTransformer {
@@ -56,38 +59,96 @@ class Inlining extends MacroTransform, IdentityDenotTransformer {
5659

5760
def newTransformer(using Context): Transformer = new Transformer {
5861
override def transform(tree: tpd.Tree)(using Context): tpd.Tree =
59-
new InliningTreeMap().transform(tree)
62+
InliningTreeMap().transform(tree)
6063
}
6164

62-
private class InliningTreeMap extends TreeMapWithImplicits {
65+
private class MemoizeStatsTreeMap extends TreeMapWithImplicits {
66+
67+
// It is safe to assume that the companion of a tree is in the same scope
68+
// Therefore, when expanding MacroAnnotations, we will only keep track of
69+
// the trees in the same scope as the current transformed tree
70+
71+
private val TrackedTrees = new Property.Key[mutable.Map[Symbol, tpd.MemberDef]]
72+
73+
private def trackedTreesCtx(trees: mutable.Map[Symbol, tpd.MemberDef])(using Context): Context =
74+
ctx.fresh.setProperty(TrackedTrees, trees)
75+
76+
private def trackedTrees(using Context): mutable.Map[Symbol, MemberDef] =
77+
ctx.property(TrackedTrees).getOrElse(mutable.Map.empty)
78+
79+
protected final def getTracked(sym: Symbol)(using Context): Option[MemberDef] =
80+
for trees <- ctx.property(TrackedTrees); tree <- trees.get(sym) yield tree
81+
82+
protected final def updateTracked(sym: Symbol, tree: MemberDef)(using Context): Unit =
83+
for trees <- ctx.property(TrackedTrees) do trees.update(sym, tree)
84+
85+
override def transform(tree: Tree)(using Context): Tree =
86+
tree match
87+
case PackageDef(_, stats) =>
88+
// Phase I: Collect and memoize all the stats
89+
val treesToTrack = stats.collect { case m: MemberDef => (m.symbol, m) }
90+
val withTrackedTreeCtx = trackedTreesCtx(mutable.Map(treesToTrack*))
91+
// Phase II & III: Transform the tree with this definitions and reconcile them with the tracked trees
92+
super.transform(tree)(using withTrackedTreeCtx) match
93+
case pkg@PackageDef(pid, stats) =>
94+
val trackedTree = trackedTrees(using withTrackedTreeCtx)
95+
val updatedStats = stats.mapConserve:
96+
case tree: MemberDef if trackedTree.contains(tree.symbol) =>
97+
trackedTree(tree.symbol)
98+
case stat => stat
99+
cpy.PackageDef(pkg)(pid = pid, stats = updatedStats)
100+
case tree => tree
101+
case block: Block =>
102+
// Phase I: Fetch all the member definitions in the block
103+
val trees = block.stats.collect { case m: MemberDef => (m.symbol, m) }
104+
val withTrackedTreeCtx = trackedTreesCtx(mutable.Map(trees*))
105+
106+
// Phase II / III: Transform the tree and Reconcile between the symbols in syms and the tree
107+
// TODO: Should we define a substitution method where we change the trees
108+
// and not the symbols (see Tree::subst)
109+
// result.subst(MacroAnnotations.trackedTrees(using withTrackedTreeCtx))
110+
super.transform(tree)(using withTrackedTreeCtx) match
111+
case b@Block(stats, expr) =>
112+
val trackedTree = trackedTrees(using withTrackedTreeCtx)
113+
cpy.Block(b)(
114+
expr = expr,
115+
stats = stats.mapConserve:
116+
case ddef: MemberDef if trackedTree.contains(ddef.symbol) =>
117+
trackedTree(ddef.symbol)
118+
case stat => stat
119+
)
120+
case tree => tree
121+
case TypeDef(_, impl: Template) =>
122+
// Phase I: Collect and memoize all the stats
123+
val trees = impl.body.collect { case m: MemberDef => (m.symbol, m) }
124+
val withTrackedTreeCtx = trackedTreesCtx(mutable.Map(trees*))
125+
// Phase II / III: Transform the tree and Reconcile between the symbols in syms and the tree
126+
super.transform(tree)(using withTrackedTreeCtx) match
127+
case tree@TypeDef(name, impl: Template) =>
128+
val trackedTree = trackedTrees(using withTrackedTreeCtx)
129+
cpy.TypeDef(tree)(
130+
name = name,
131+
rhs = cpy.Template(impl)(
132+
body = impl.body.mapConserve:
133+
case ddef: MemberDef if trackedTree.contains(ddef.symbol) =>
134+
trackedTree(ddef.symbol)
135+
case stat => stat
136+
)
137+
)
138+
case tree => tree
139+
case _ => super.transform(tree)
140+
}
141+
142+
private class InliningTreeMap extends MemoizeStatsTreeMap {
63143

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

69149
override def transform(tree: Tree)(using Context): Tree = {
70150
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)
151+
case tree: MemberDef => transformMemberDef(tree)
91152
case _: Typed | _: Block =>
92153
super.transform(tree)
93154
case _: PackageDef =>
@@ -113,7 +174,56 @@ class Inlining extends MacroTransform, IdentityDenotTransformer {
113174
else Inlines.inlineCall(tree1)
114175
else super.transform(tree)
115176
}
177+
178+
private def transformMemberDef(tree: MemberDef)(using Context) : Tree =
179+
if tree.symbol.is(Inline) then tree
180+
else if tree.symbol.is(Param) then
181+
super.transform(tree)
182+
else if
183+
!tree.symbol.isPrimaryConstructor
184+
&& StagingLevel.level == 0
185+
&& MacroAnnotations.hasMacroAnnotation(tree.symbol)
186+
then
187+
// Fetch the companion's tree
188+
val companionTree =
189+
getTracked(
190+
sym = if tree.symbol.is(ModuleClass) then tree.symbol.companionClass
191+
else if tree.symbol.is(ModuleVal) then NoSymbol
192+
else tree.symbol.companionModule.moduleClass)
193+
194+
// Fetch the latest tracked tree (It might have already been processed by its companion)
195+
val latestTree = getTracked(tree.symbol).getOrElse(tree)
196+
197+
// Expand and process MacroAnnotations
198+
val (trees, companion) =
199+
MacroAnnotations(self).expandAnnotations(latestTree, companionTree)
200+
201+
// Update the tracked trees
202+
for case tree : MemberDef <- trees do updateTracked(tree.symbol, tree)
203+
for tree <- companion do updateTracked(tree.symbol, tree)
204+
205+
// Perform inlining on the expansion of the annotations
206+
val trees1 = trees.map(super.transform)
207+
208+
for case tree: MemberDef <- trees1 do updateTracked(tree.symbol, tree)
209+
210+
// Find classes added to the top level from a package object
211+
val (topClasses, trees2) =
212+
if ctx.owner.isPackageObject then trees1.partition(_.symbol.owner == ctx.owner.owner)
213+
else (Nil, trees1)
214+
if topClasses.nonEmpty then
215+
newTopClasses.getOrElseUpdate(ctx.owner.owner, new mutable.ListBuffer) ++= topClasses
216+
flatTree(trees2)
217+
else
218+
super.transform(tree) match
219+
case tree: MemberDef =>
220+
updateTracked(tree.symbol, tree)
221+
tree
222+
case tree => tree
223+
end transformMemberDef
224+
116225
}
226+
117227
}
118228

119229
object Inlining:

0 commit comments

Comments
 (0)