Skip to content

Commit 1c26cda

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

File tree

75 files changed

+658
-350
lines changed

Some content is hidden

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

75 files changed

+658
-350
lines changed

‎compiler/src/dotty/tools/dotc/CompilationUnit.scala

+2-2
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import config.{SourceVersion, Feature}
1717
import StdNames.nme
1818
import scala.annotation.internal.sharable
1919
import scala.util.control.NoStackTrace
20-
import transform.MacroAnnotations
20+
import transform.MacroAnnotations.isMacroAnnotation
2121

2222
class CompilationUnit protected (val source: SourceFile, val info: CompilationUnitInfo | Null) {
2323

@@ -193,7 +193,7 @@ object CompilationUnit {
193193
case _ =>
194194
case _ =>
195195
for annot <- tree.symbol.annotations do
196-
if MacroAnnotations.isMacroAnnotation(annot) then
196+
if annot.isMacroAnnotation then
197197
ctx.compilationUnit.hasMacroAnnotations = true
198198
traverseChildren(tree)
199199
}

‎compiler/src/dotty/tools/dotc/ast/TreeMapWithTrackedStats.scala

+80
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package dotty.tools.dotc
2+
package ast
3+
4+
import tpd.*
5+
import core.Contexts.*
6+
import core.Symbols.*
7+
import util.Property
8+
9+
import scala.collection.mutable
10+
11+
abstract class TreeMapWithTrackedStats extends TreeMapWithImplicits:
12+
13+
import TreeMapWithTrackedStats.*
14+
15+
// It is safe to assume that the companion of a tree is in the same scope
16+
// Therefore, when expanding MacroAnnotations, we will only keep track of
17+
// the trees in the same scope as the current transformed tree
18+
19+
protected final def getTracked(sym: Symbol)(using Context): Option[MemberDef] =
20+
for trees <- ctx.property(TrackedTrees)
21+
tree <- trees.get(sym)
22+
yield tree
23+
24+
protected final def updateTracked(tree: Tree)(using Context): Tree =
25+
tree match
26+
case tree: MemberDef =>
27+
trackedTrees.update(tree.symbol, tree)
28+
tree
29+
case _ => tree
30+
end updateTracked
31+
32+
private final def withUpdatedTree(stats: List[Tree])(using Context) =
33+
val trackedTrees = TreeMapWithTrackedStats.trackedTrees
34+
stats.mapConserve:
35+
case tree: MemberDef if trackedTrees.contains(tree.symbol) =>
36+
trackedTrees(tree.symbol)
37+
case stat => stat
38+
39+
override def transform(tree: Tree)(using Context): Tree =
40+
tree match
41+
case PackageDef(_, stats) =>
42+
// Step I: Collect and memoize all the stats
43+
inContext(trackedTreesCtx(stats)):
44+
// Step II: Transform the tree
45+
val result = super.transform(tree)
46+
// Step III: Reconcile between the symbols in syms and the tree
47+
(result: @unchecked) match
48+
case pkg@PackageDef(pid, stats) =>
49+
cpy.PackageDef(pkg)(pid = pid, stats = withUpdatedTree(stats))
50+
case block: Block =>
51+
// Step I: Collect all the member definitions in the block
52+
inContext(trackedTreesCtx(block.stats)):
53+
// Step II: Transform the tree
54+
val result = super.transform(tree)
55+
// Step III: Reconcile between the symbols in syms and the tree
56+
(result: @unchecked) match
57+
case b@Block(stats, expr) =>
58+
cpy.Block(b)(expr = expr, stats = withUpdatedTree(stats))
59+
case TypeDef(_, impl: Template) =>
60+
// Step I: Collect and memoize all the stats
61+
inContext(trackedTreesCtx(impl.body)):
62+
// Step II: Transform the tree
63+
val result = super.transform(tree)
64+
// Step III: Reconcile between the symbols in syms and the tree
65+
(result : @unchecked) match
66+
case tree@TypeDef(name, impl: Template) =>
67+
cpy.TypeDef(tree)(rhs = cpy.Template(impl)(body = withUpdatedTree(impl.body)))
68+
case _ => super.transform(tree)
69+
70+
end TreeMapWithTrackedStats
71+
72+
object TreeMapWithTrackedStats:
73+
private val TrackedTrees = new Property.Key[mutable.Map[Symbol, tpd.MemberDef]]
74+
75+
private def trackedTrees(using Context): mutable.Map[Symbol, MemberDef] =
76+
ctx.property(TrackedTrees).get
77+
78+
private def trackedTreesCtx(stats: List[Tree])(using Context): Context =
79+
val treesToTrack = stats.collect { case m: MemberDef => (m.symbol, m) }
80+
ctx.fresh.setProperty(TrackedTrees, mutable.Map(treesToTrack*))

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

+58-31
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,24 @@
11
package dotty.tools.dotc
22
package transform
33

4+
import ast.tpd
5+
import ast.Trees.*
6+
import ast.TreeMapWithTrackedStats
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 MacroAnnotations.hasMacroAnnotation
16+
import inlines.Inlines
17+
import quoted.*
18+
import staging.StagingLevel
19+
import util.Property
820

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
21+
import scala.collection.mutable
1822

1923
/** Inlines all calls to inline methods that are not in an inline method or a quote */
2024
class Inlining extends MacroTransform, IdentityDenotTransformer {
@@ -56,38 +60,21 @@ class Inlining extends MacroTransform, IdentityDenotTransformer {
5660

5761
def newTransformer(using Context): Transformer = new Transformer {
5862
override def transform(tree: tpd.Tree)(using Context): tpd.Tree =
59-
new InliningTreeMap().transform(tree)
63+
InliningTreeMap().transform(tree)
6064
}
6165

62-
private class InliningTreeMap extends TreeMapWithImplicits {
66+
private class InliningTreeMap extends TreeMapWithTrackedStats {
6367

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

6973
override def transform(tree: Tree)(using Context): Tree = {
7074
tree match
7175
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)
76+
// Fetch the latest tracked tree (It might have already been processed by its companion)
77+
transformMemberDef(getTracked(tree.symbol).getOrElse(tree))
9178
case _: Typed | _: Block =>
9279
super.transform(tree)
9380
case _: PackageDef =>
@@ -113,7 +100,47 @@ class Inlining extends MacroTransform, IdentityDenotTransformer {
113100
else Inlines.inlineCall(tree1)
114101
else super.transform(tree)
115102
}
103+
104+
private def transformMemberDef(tree: MemberDef)(using Context) : Tree =
105+
if tree.symbol.is(Inline) then tree
106+
else if tree.symbol.is(Param) then
107+
super.transform(tree)
108+
else if
109+
!tree.symbol.isPrimaryConstructor
110+
&& StagingLevel.level == 0
111+
&& tree.symbol.hasMacroAnnotation
112+
then
113+
// Fetch the companion's tree
114+
val companionSym =
115+
if tree.symbol.is(ModuleClass) then tree.symbol.companionClass
116+
else if tree.symbol.is(ModuleVal) then NoSymbol
117+
else tree.symbol.companionModule.moduleClass
118+
119+
// Expand and process MacroAnnotations
120+
val (trees, companion) = MacroAnnotations.expandAnnotations(tree, getTracked(companionSym))
121+
122+
// Enter the new symbols & Update the tracked trees
123+
(companion.toList ::: trees).foreach: tree =>
124+
MacroAnnotations.enterMissingSymbols(tree, self)
125+
updateTracked(tree)
126+
127+
// Perform inlining on the expansion of the annotations
128+
val trees1 = trees.map(super.transform)
129+
trees1.foreach(updateTracked)
130+
131+
// Find classes added to the top level from a package object
132+
val (topClasses, trees2) =
133+
if ctx.owner.isPackageObject then trees1.partition(_.symbol.owner == ctx.owner.owner)
134+
else (Nil, trees1)
135+
if topClasses.nonEmpty then
136+
newTopClasses.getOrElseUpdate(ctx.owner.owner, new mutable.ListBuffer) ++= topClasses
137+
flatTree(trees2)
138+
else
139+
updateTracked(super.transform(tree))
140+
end transformMemberDef
141+
116142
}
143+
117144
}
118145

119146
object Inlining:

0 commit comments

Comments
 (0)