Use aggregation instead of inheritance for dependency collection

Author: odersky

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

@@ -19,20 +19,24 @@ import ExplicitOuter.outer
 import util.Store
 import collection.mutable
 import collection.mutable.{ HashMap, HashSet, LinkedHashMap, TreeSet }
+import annotation.constructorOnly
 
-abstract class DependencyCollector:
+/** Exposes the dependencies of the typed tree in the current compilation unit
+ *  in sets `freeVars`, `liftedOwner`.
+ */
+abstract class Dependencies(@constructorOnly rootContext: Context):
   import ast.tpd._
 
   def enclosure(using Context): Symbol
   def isExpr(sym: Symbol)(using Context): Boolean
 
-  protected type SymSet = TreeSet[Symbol]
+  type SymSet = TreeSet[Symbol]
 
   /** A map storing free variables of functions and classes */
   val free: mutable.LinkedHashMap[Symbol, SymSet] = new LinkedHashMap
 
   /** A hashtable storing calls between functions */
-  val called = new LinkedHashMap[Symbol, SymSet]
+  private val called = new LinkedHashMap[Symbol, SymSet]
 
   /** A map from local methods and classes to the owners to which they will be lifted as members.
    *  For methods and classes that do not have any dependencies this will be the enclosing package.
@@ -55,17 +59,15 @@ abstract class DependencyCollector:
   private def symSet(f: LinkedHashMap[Symbol, SymSet], sym: Symbol): SymSet =
     f.getOrElseUpdate(sym, newSymSet)
 
-  def freeVars(sym: Symbol): List[Symbol] = free.get(sym) match
-    case Some(set) => set.toList
-    case None => Nil
+  def freeVars(sym: Symbol): collection.Set[Symbol] = free.getOrElse(sym, Set.empty)
 
   /** A symbol is local if it is owned by a term or a local trait,
    *  or if it is a constructor of a local symbol.
    *  Note: we count members of local traits as local since their free variables
    *  have to be passed on from their callers. By contrast, class members get their
    *  free variable proxies from their enclosing class.
    */
-  def isLocal(sym: Symbol)(using Context): Boolean =
+  private def isLocal(sym: Symbol)(using Context): Boolean =
     val owner = sym.maybeOwner
     owner.isTerm
     || owner.is(Trait) && isLocal(owner)
@@ -74,7 +76,7 @@ abstract class DependencyCollector:
   /** Set `liftedOwner(sym)` to `owner` if `owner` is more deeply nested
    *  than the previous value of `liftedowner(sym)`.
    */
-  def narrowLiftedOwner(sym: Symbol, owner: Symbol)(using Context): Unit =
+  private def narrowLiftedOwner(sym: Symbol, owner: Symbol)(using Context): Unit =
     if sym.maybeOwner.isTerm
         && owner.isProperlyContainedIn(liftedOwner(sym))
         && owner != sym
@@ -172,7 +174,7 @@ abstract class DependencyCollector:
     symSet(called, caller) += callee
   }
 
-  def process(tree: Tree)(using Context) =
+  protected def process(tree: Tree)(using Context) =
     val sym = tree.symbol
 
     def narrowTo(thisClass: ClassSymbol) =
@@ -259,8 +261,9 @@ abstract class DependencyCollector:
       changedLiftedOwner
     do ()
 
-  def collectDependencies()(using Context): Unit =
-    CollectDependencies().traverse(ctx.compilationUnit.tpdTree)
+  inContext(rootContext) {
+    CollectDependencies().traverse(rootContext.compilationUnit.tpdTree)
     computeFreeVars()
     computeLiftedOwners()
-end DependencyCollector
+  }
+end Dependencies

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

@@ -27,37 +27,41 @@ object LambdaLift {
   val name: String = "lambdaLift"
 
   /** The core lambda lift functionality. */
-  class Lifter(thisPhase: MiniPhase & DenotTransformer)(using Context) extends DependencyCollector:
-
-    /** A map storing the free variable proxies of functions and classes.
-     *  For every function and class, this is a map from the free variables
-     *  of that function or class to the proxy symbols accessing them.
-     */
-    private val proxyMap = new LinkedHashMap[Symbol, Map[Symbol, Symbol]]
+  class Lifter(thisPhase: MiniPhase & DenotTransformer)(using Context):
 
     /** The outer parameter of a constructor */
     private val outerParam = new HashMap[Symbol, Symbol]
 
     /** Buffers for lifted out classes and methods, indexed by owner */
     val liftedDefs: mutable.HashMap[Symbol, mutable.ListBuffer[Tree]] = new HashMap
 
-    def proxyOf(sym: Symbol, fv: Symbol): Symbol = proxyMap.getOrElse(sym, Map.empty)(fv)
+    val deps = new Dependencies(ctx.withPhase(thisPhase)):
+      def enclosure(using Context) = ctx.owner.enclosingMethod
+      def isExpr(sym: Symbol)(using Context): Boolean = sym.is(Method)
+
+      override def process(tree: Tree)(using Context): Unit =
+        super.process(tree)
+        tree match
+          case tree: DefDef if tree.symbol.isConstructor =>
+            tree.termParamss.head.find(_.name == nme.OUTER) match
+              case Some(vdef) => outerParam(tree.symbol) = vdef.symbol
+              case _ =>
+          case tree: Template =>
+            liftedDefs(tree.symbol.owner) = new mutable.ListBuffer
+          case _ =>
+    end deps
+    export deps.{liftedOwner, free}
 
-    def proxies(sym: Symbol): List[Symbol] =  freeVars(sym).map(proxyOf(sym, _))
+    /** A map storing the free variable proxies of functions and classes.
+     *  For every function and class, this is a map from the free variables
+     *  of that function or class to the proxy symbols accessing them.
+     */
+    private val proxyMap = new LinkedHashMap[Symbol, Map[Symbol, Symbol]]
 
-    def enclosure(using Context) = ctx.owner.enclosingMethod
-    def isExpr(sym: Symbol)(using Context): Boolean = sym.is(Method)
+    def proxyOf(sym: Symbol, fv: Symbol): Symbol = proxyMap.getOrElse(sym, Map.empty)(fv)
 
-    override def process(tree: Tree)(using Context): Unit =
-      super.process(tree)
-      tree match
-        case tree: DefDef if tree.symbol.isConstructor =>
-          tree.termParamss.head.find(_.name == nme.OUTER) match
-            case Some(vdef) => outerParam(tree.symbol) = vdef.symbol
-            case _ =>
-        case tree: Template =>
-          liftedDefs(tree.symbol.owner) = new mutable.ListBuffer
-        case _ =>
+    def proxies(sym: Symbol): List[Symbol] =
+      deps.freeVars(sym).toList.map(proxyOf(sym, _))
 
     private def newName(sym: Symbol)(using Context): Name =
       if (sym.isAnonymousFunction && sym.owner.is(Method))
@@ -133,9 +137,6 @@ object LambdaLift {
     }
 
     // initialization
-    atPhase(thisPhase) {
-      collectDependencies()
-    }
     atPhase(thisPhase.next) {
       generateProxies()
       liftLocals()
@@ -205,7 +206,7 @@ object LambdaLift {
 
         /** Initialize proxy fields from proxy parameters and map `rhs` from fields to parameters */
         def copyParams(rhs: Tree) = {
-          val fvs = freeVars(sym.owner)
+          val fvs = deps.freeVars(sym.owner).toList
           val classProxies = fvs.map(proxyOf(sym.owner, _))
           val constrProxies = fvs.map(proxyOf(sym, _))
           report.debuglog(i"copy params ${constrProxies.map(_.showLocated)}%, % to ${classProxies.map(_.showLocated)}%, %}")