Skip to content

Commit f70157f

Browse files
committed
Refactor LambdaLift
Break out the dependency collection into a separate class. That makes it more re-usable and makes LambdaLift itself easier to understand.
1 parent 06df92b commit f70157f

File tree

2 files changed

+281
-255
lines changed

2 files changed

+281
-255
lines changed
Lines changed: 266 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,266 @@
1+
package dotty.tools.dotc
2+
package transform
3+
4+
import MegaPhase._
5+
import core.Denotations.NonSymSingleDenotation
6+
import core.DenotTransformers._
7+
import core.Symbols._
8+
import core.Contexts._
9+
import core.Types._
10+
import core.Flags._
11+
import core.Decorators._
12+
import core.StdNames.nme
13+
import core.Names._
14+
import core.NameOps._
15+
import core.NameKinds.ExpandPrefixName
16+
import ast.Trees._
17+
import SymUtils._
18+
import ExplicitOuter.outer
19+
import util.Store
20+
import collection.mutable
21+
import collection.mutable.{ HashMap, HashSet, LinkedHashMap, TreeSet }
22+
23+
abstract class DependencyCollector:
24+
import ast.tpd._
25+
26+
def enclosure(using Context): Symbol
27+
def isExpr(sym: Symbol)(using Context): Boolean
28+
29+
protected type SymSet = TreeSet[Symbol]
30+
31+
/** A map storing free variables of functions and classes */
32+
val free: mutable.LinkedHashMap[Symbol, SymSet] = new LinkedHashMap
33+
34+
/** A hashtable storing calls between functions */
35+
val called = new LinkedHashMap[Symbol, SymSet]
36+
37+
/** A map from local methods and classes to the owners to which they will be lifted as members.
38+
* For methods and classes that do not have any dependencies this will be the enclosing package.
39+
* symbols with packages as lifted owners will subsequently represented as static
40+
* members of their toplevel class, unless their enclosing class was already static.
41+
* Note: During tree transform (which runs at phase LambdaLift + 1), liftedOwner
42+
* is also used to decide whether a method had a term owner before.
43+
*/
44+
val liftedOwner = new LinkedHashMap[Symbol, Symbol]
45+
46+
/** A flag to indicate whether new free variables have been found */
47+
private var changedFreeVars: Boolean = _
48+
49+
/** A flag to indicate whether lifted owners have changed */
50+
private var changedLiftedOwner: Boolean = _
51+
52+
private val ord: Ordering[Symbol] = Ordering.by(_.id)
53+
private def newSymSet = TreeSet.empty[Symbol](ord)
54+
55+
private def symSet(f: LinkedHashMap[Symbol, SymSet], sym: Symbol): SymSet =
56+
f.getOrElseUpdate(sym, newSymSet)
57+
58+
def freeVars(sym: Symbol): List[Symbol] = free.get(sym) match
59+
case Some(set) => set.toList
60+
case None => Nil
61+
62+
/** A symbol is local if it is owned by a term or a local trait,
63+
* or if it is a constructor of a local symbol.
64+
* Note: we count members of local traits as local since their free variables
65+
* have to be passed on from their callers. By contrast, class members get their
66+
* free variable proxies from their enclosing class.
67+
*/
68+
def isLocal(sym: Symbol)(using Context): Boolean =
69+
val owner = sym.maybeOwner
70+
owner.isTerm
71+
|| owner.is(Trait) && isLocal(owner)
72+
|| sym.isConstructor && isLocal(owner)
73+
74+
/** Set `liftedOwner(sym)` to `owner` if `owner` is more deeply nested
75+
* than the previous value of `liftedowner(sym)`.
76+
*/
77+
def narrowLiftedOwner(sym: Symbol, owner: Symbol)(using Context): Unit =
78+
if sym.maybeOwner.isTerm
79+
&& owner.isProperlyContainedIn(liftedOwner(sym))
80+
&& owner != sym
81+
then
82+
report.log(i"narrow lifted $sym to $owner")
83+
changedLiftedOwner = true
84+
liftedOwner(sym) = owner
85+
86+
private class NoPath extends Exception
87+
88+
/** Mark symbol `sym` as being free in `enclosure`, unless `sym` is defined
89+
* in `enclosure` or there is an intermediate class properly containing `enclosure`
90+
* in which `sym` is also free. Also, update `liftedOwner` of `enclosure` so
91+
* that `enclosure` can access `sym`, or its proxy in an intermediate class.
92+
* This means:
93+
*
94+
* 1. If there is an intermediate class in which `sym` is free, `enclosure`
95+
* must be contained in that class (in order to access the `sym proxy stored
96+
* in the class).
97+
*
98+
* 2. If there is no intermediate class, `enclosure` must be contained
99+
* in the class enclosing `sym`.
100+
*
101+
* @return If there is a non-trait class between `enclosure` and
102+
* the owner of `sym`, the largest such class.
103+
* Otherwise, if there is a trait between `enclosure` and
104+
* the owner of `sym`, the largest such trait.
105+
* Otherwise, NoSymbol.
106+
*
107+
* @pre sym.owner.isTerm, (enclosure.isMethod || enclosure.isClass)
108+
*
109+
* The idea of `markFree` is illustrated with an example:
110+
*
111+
* def f(x: int) = {
112+
* class C {
113+
* class D {
114+
* val y = x
115+
* }
116+
* }
117+
* }
118+
*
119+
* In this case `x` is free in the primary constructor of class `C`.
120+
* but it is not free in `D`, because after lambda lift the code would be transformed
121+
* as follows:
122+
*
123+
* def f(x$0: int) {
124+
* class C(x$0: int) {
125+
* val x$1 = x$0
126+
* class D {
127+
* val y = outer.x$1
128+
* }
129+
* }
130+
* }
131+
*/
132+
private def markFree(sym: Symbol, enclosure: Symbol)(using Context): Symbol = try {
133+
if (!enclosure.exists) throw new NoPath
134+
if (enclosure == sym.enclosure) NoSymbol
135+
else {
136+
def nestedInConstructor(sym: Symbol): Boolean =
137+
sym.isConstructor ||
138+
sym.isTerm && nestedInConstructor(sym.enclosure)
139+
report.debuglog(i"mark free: ${sym.showLocated} with owner ${sym.maybeOwner} marked free in $enclosure")
140+
val intermediate =
141+
if (enclosure.is(PackageClass)) enclosure
142+
else if (enclosure.isConstructor) markFree(sym, enclosure.owner.enclosure)
143+
else markFree(sym, enclosure.enclosure)
144+
if (intermediate.exists) narrowLiftedOwner(enclosure, intermediate)
145+
if !intermediate.isRealClass || nestedInConstructor(enclosure) then
146+
// Constructors and methods nested inside traits get the free variables
147+
// of the enclosing trait or class.
148+
// Conversely, local traits do not get free variables.
149+
// Methods inside constructors also don't have intermediates,
150+
// need to get all their free variables passed directly.
151+
if (!enclosure.is(Trait))
152+
if (symSet(free, enclosure).add(sym)) {
153+
changedFreeVars = true
154+
report.log(i"$sym is free in $enclosure")
155+
}
156+
if (intermediate.isRealClass) intermediate
157+
else if (enclosure.isRealClass) enclosure
158+
else if (intermediate.isClass) intermediate
159+
else if (enclosure.isClass) enclosure
160+
else NoSymbol
161+
}
162+
}
163+
catch {
164+
case ex: NoPath =>
165+
println(i"error lambda lifting ${ctx.compilationUnit}: $sym is not visible from $enclosure")
166+
throw ex
167+
}
168+
169+
private def markCalled(callee: Symbol, caller: Symbol)(using Context): Unit = {
170+
report.debuglog(i"mark called: $callee of ${callee.owner} is called by $caller in ${caller.owner}")
171+
assert(isLocal(callee))
172+
symSet(called, caller) += callee
173+
}
174+
175+
def process(tree: Tree)(using Context) =
176+
val sym = tree.symbol
177+
178+
def narrowTo(thisClass: ClassSymbol) =
179+
val enclMethod = enclosure
180+
val enclClass = enclMethod.enclosingClass
181+
narrowLiftedOwner(enclMethod,
182+
if enclClass.isContainedIn(thisClass) then thisClass
183+
else enclClass) // unknown this reference, play it safe and assume the narrowest possible owner
184+
185+
tree match
186+
case tree: Ident =>
187+
if isLocal(sym) then
188+
if isExpr(sym) then markCalled(sym, enclosure)
189+
else if sym.isTerm then markFree(sym, enclosure)
190+
def captureImplicitThis(x: Type): Unit = x match
191+
case tr@TermRef(x, _) if !tr.termSymbol.isStatic => captureImplicitThis(x)
192+
case x: ThisType if !x.tref.typeSymbol.isStaticOwner => narrowTo(x.tref.typeSymbol.asClass)
193+
case _ =>
194+
captureImplicitThis(tree.tpe)
195+
case tree: Select =>
196+
if isExpr(sym) && isLocal(sym) then markCalled(sym, enclosure)
197+
case tree: This =>
198+
narrowTo(tree.symbol.asClass)
199+
case tree: DefDef =>
200+
if sym.owner.isTerm then
201+
liftedOwner(sym) = sym.enclosingPackageClass
202+
// this will make methods in supercall constructors of top-level classes owned
203+
// by the enclosing package, which means they will be static.
204+
// On the other hand, all other methods will be indirectly owned by their
205+
// top-level class. This avoids possible deadlocks when a static method
206+
// has to access its enclosing object from the outside.
207+
else if sym.isConstructor then
208+
if sym.isPrimaryConstructor && isLocal(sym.owner) && !sym.owner.is(Trait) then
209+
// add a call edge from the constructor of a local non-trait class to
210+
// the class itself. This is done so that the constructor inherits
211+
// the free variables of the class.
212+
symSet(called, sym) += sym.owner
213+
case tree: TypeDef =>
214+
if sym.owner.isTerm then liftedOwner(sym) = sym.topLevelClass.owner
215+
case _ =>
216+
end process
217+
218+
private class CollectDependencies extends TreeTraverser:
219+
def traverse(tree: Tree)(using Context) =
220+
try
221+
process(tree)
222+
traverseChildren(tree)
223+
catch case ex: Exception =>
224+
println(i"$ex while traversing $tree")
225+
throw ex
226+
227+
/** Compute final free variables map `fvs by closing over caller dependencies. */
228+
private def computeFreeVars()(using Context): Unit =
229+
while
230+
changedFreeVars = false
231+
for
232+
caller <- called.keys
233+
callee <- called(caller)
234+
fvs <- free get callee
235+
fv <- fvs
236+
do
237+
markFree(fv, caller)
238+
changedFreeVars
239+
do ()
240+
241+
/** Compute final liftedOwner map by closing over caller dependencies */
242+
private def computeLiftedOwners()(using Context): Unit =
243+
while
244+
changedLiftedOwner = false
245+
for
246+
caller <- called.keys
247+
callee <- called(caller)
248+
do
249+
val normalizedCallee = callee.skipConstructor
250+
val calleeOwner = normalizedCallee.owner
251+
if calleeOwner.isTerm then narrowLiftedOwner(caller, liftedOwner(normalizedCallee))
252+
else
253+
assert(calleeOwner.is(Trait))
254+
// methods nested inside local trait methods cannot be lifted out
255+
// beyond the trait. Note that we can also call a trait method through
256+
// a qualifier; in that case no restriction to lifted owner arises.
257+
if caller.isContainedIn(calleeOwner) then
258+
narrowLiftedOwner(caller, calleeOwner)
259+
changedLiftedOwner
260+
do ()
261+
262+
def collectDependencies()(using Context): Unit =
263+
CollectDependencies().traverse(ctx.compilationUnit.tpdTree)
264+
computeFreeVars()
265+
computeLiftedOwners()
266+
end DependencyCollector

0 commit comments

Comments
 (0)