Generalize lift behavior between Flatten and LambdaLift

To be combinable with follow-up mini-phases the lift operation needs to handle Thickets specially.
This commit factors out the behavior from LambdaLift, so that Flatten can do the same thing.

3 files changed, 9 insertions, 5 deletions

diff --git a/src/dotty/tools/dotc/ast/Trees.scala b/src/dotty/tools/dotc/ast/Trees.scala
index 71026a449..648f19a4f 100644
--- a/src/dotty/tools/dotc/ast/Trees.scala
+++ b/src/dotty/tools/dotc/ast/Trees.scala
@@ -347,6 +347,11 @@ object Trees {
       s
     }
 
+    /** If this is a thicket, gerform `op` on each of its trees
+     *  otherwise, perform `op` ion tree itself.
+     */
+    def foreachInThicket(op: Tree[T] => Unit): Unit = op(this)
+
     override def toText(printer: Printer) = printer.toText(this)
 
     override def hashCode(): Int = System.identityHashCode(this)
@@ -809,6 +814,8 @@ object Trees {
       val newTrees = trees.map(_.withPos(pos))
       new Thicket[T](newTrees).asInstanceOf[this.type]
     }
+    override def foreachInThicket(op: Tree[T] => Unit): Unit =
+      trees foreach (_.foreachInThicket(op))
   }
 
   class EmptyValDef[T >: Untyped] extends ValDef[T](
diff --git a/src/dotty/tools/dotc/transform/Flatten.scala b/src/dotty/tools/dotc/transform/Flatten.scala
index 769503f9d..dceefc0bc 100644
--- a/src/dotty/tools/dotc/transform/Flatten.scala
+++ b/src/dotty/tools/dotc/transform/Flatten.scala
@@ -32,7 +32,7 @@ class Flatten extends MiniPhaseTransform with SymTransformer { thisTransform =>
   private def liftIfNested(tree: Tree)(implicit ctx: Context, info: TransformerInfo) =
     if (ctx.owner is Package) tree
     else {
-      liftedDefs += transformFollowing(tree)
+      transformFollowing(tree).foreachInThicket(liftedDefs += _)
       EmptyTree
     }
 
diff --git a/src/dotty/tools/dotc/transform/LambdaLift.scala b/src/dotty/tools/dotc/transform/LambdaLift.scala
index 051c6065e..f36ff6247 100644
--- a/src/dotty/tools/dotc/transform/LambdaLift.scala
+++ b/src/dotty/tools/dotc/transform/LambdaLift.scala
@@ -357,10 +357,7 @@ class LambdaLift extends MiniPhaseTransform with IdentityDenotTransformer { this
 
   private def liftDef(tree: MemberDef)(implicit ctx: Context, info: TransformerInfo): Tree = {
     val buf = liftedDefs(tree.symbol.owner)
-    transformFollowing(rename(tree, tree.symbol.name)) match {
-      case Thicket(trees) => buf ++= trees
-      case tree => buf += tree
-    }
+    transformFollowing(rename(tree, tree.symbol.name)).foreachInThicket(buf += _)
     EmptyTree
   }