Facturing out EnclosingMethodTraverser from CapturedVars

The idea to traverse with currently enclosing methid is also used in
LambdaLift and could be used elsewhere.

2 files changed, 26 insertions, 15 deletions

diff --git a/src/dotty/tools/dotc/ast/tpd.scala b/src/dotty/tools/dotc/ast/tpd.scala
index 216b1c0c5..45e0aff54 100644
--- a/src/dotty/tools/dotc/ast/tpd.scala
+++ b/src/dotty/tools/dotc/ast/tpd.scala
@@ -653,6 +653,22 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
     Ident(defn.ScalaRuntimeModule.requiredMethod(name).termRef).appliedToArgs(args)
   }
 
+  /** A traverser that passes the enlcosing class or method as an argumenr
+   *  to the traverse method.
+   */
+  abstract class EnclosingMethodTraverser(implicit ctx: Context) extends TreeAccumulator[Symbol] {
+    def traverse(enclMeth: Symbol, tree: Tree): Unit
+    def apply(enclMeth: Symbol, tree: Tree) = {
+      tree match {
+        case _: DefTree if tree.symbol.exists =>
+          traverse(tree.symbol.enclosingMethod, tree)
+        case _ =>
+          traverse(enclMeth, tree)
+      }
+      enclMeth
+    }
+  }
+
   // ensure that constructors are fully applied?
   // ensure that normal methods are fully applied?
 
diff --git a/src/dotty/tools/dotc/transform/CapturedVars.scala b/src/dotty/tools/dotc/transform/CapturedVars.scala
index d5bd56bc3..f13cd5690 100644
--- a/src/dotty/tools/dotc/transform/CapturedVars.scala
+++ b/src/dotty/tools/dotc/transform/CapturedVars.scala
@@ -27,21 +27,16 @@ class CapturedVars extends MiniPhaseTransform with SymTransformer { thisTransfor
 
   private var captured: mutable.HashSet[Symbol] = _
 
-  private class CollectCaptured(implicit ctx: Context) extends TreeAccumulator[Symbol] {
-    def apply(enclMeth: Symbol, tree: Tree) = {
-      tree match {
-        case id: Ident =>
-          val sym = id.symbol
-          if (sym.is(Mutable, butNot = Method) && sym.owner.isTerm && sym.enclosingMethod != enclMeth) {
-            ctx.log(i"capturing $sym in ${sym.enclosingMethod}, referenced from $enclMeth")
-            captured += sym
-          }
-        case tree: DefTree if tree.symbol.exists =>
-          foldOver(tree.symbol.enclosingMethod, tree)
-        case _ =>
-          foldOver(enclMeth, tree)
-      }
-      enclMeth
+  private class CollectCaptured(implicit ctx: Context) extends EnclosingMethodTraverser {
+    def traverse(enclMeth: Symbol, tree: Tree) = tree match {
+      case id: Ident =>
+        val sym = id.symbol
+        if (sym.is(Mutable, butNot = Method) && sym.owner.isTerm && sym.enclosingMethod != enclMeth) {
+          ctx.log(i"capturing $sym in ${sym.enclosingMethod}, referenced from $enclMeth")
+          captured += sym
+        }
+      case _ =>
+        foldOver(enclMeth, tree)
     }
     def runOver(tree: Tree) = {
       captured = mutable.HashSet()