Make import contexts visible in macro transforms

Todo: We have very similar and involved operations now sor handling statements
in tpd.TreeMapWithPreciseStatContexts and MegaPhase. Can we factor out the common logic?
But we should not create any closures doing so, to keep things fast.

Author: odersky

compiler/src/dotty/tools/dotc/ast/TreeMapWithImplicits.scala

@@ -15,51 +15,12 @@ import scala.annotation.tailrec
  *
  *  This incudes implicits defined in scope as well as imported implicits.
  */
-class TreeMapWithImplicits extends tpd.TreeMap {
+class TreeMapWithImplicits extends tpd.TreeMapWithPreciseStatContexts {
   import tpd._
 
   def transformSelf(vd: ValDef)(using Context): ValDef =
     cpy.ValDef(vd)(tpt = transform(vd.tpt))
 
-  /** Transform statements, while maintaining import contexts and expression contexts
-   *  in the same way as Typer does. The code addresses additional concerns:
-   *   - be tail-recursive where possible
-   *   - don't re-allocate trees where nothing has changed
-   */
-  override def transformStats(stats: List[Tree], exprOwner: Symbol)(using Context): List[Tree] = {
-
-    @tailrec def traverse(curStats: List[Tree])(using Context): List[Tree] = {
-
-      def recur(stats: List[Tree], changed: Tree, rest: List[Tree])(using Context): List[Tree] =
-        if (stats eq curStats) {
-          val rest1 = transformStats(rest, exprOwner)
-          changed match {
-            case Thicket(trees) => trees ::: rest1
-            case tree => tree :: rest1
-          }
-        }
-        else stats.head :: recur(stats.tail, changed, rest)
-
-      curStats match {
-        case stat :: rest =>
-          val statCtx = stat match {
-            case stat: DefTree => ctx
-            case _ => ctx.exprContext(stat, exprOwner)
-          }
-          val restCtx = stat match {
-            case stat: Import => ctx.importContext(stat, stat.symbol)
-            case _ => ctx
-          }
-          val stat1 = transform(stat)(using statCtx)
-          if (stat1 ne stat) recur(stats, stat1, rest)(using restCtx)
-          else traverse(rest)(using restCtx)
-        case nil =>
-          stats
-      }
-    }
-    traverse(stats)
-  }
-
   private def nestedScopeCtx(defs: List[Tree])(using Context): Context = {
     val nestedCtx = ctx.fresh.setNewScope
     defs foreach {

compiler/src/dotty/tools/dotc/ast/tpd.scala

@@ -1153,10 +1153,54 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
       recur(trees, 0)
   end extension
 
+  /** A treemap that generates the same contexts as the original typer for statements.
+   *  This means:
+   *    - statements that are not definitions get the exprOwner as owner
+   *    - imports are reflected in the contexts of subsequent statements
+   */
+  class TreeMapWithPreciseStatContexts(cpy: TreeCopier = tpd.cpy) extends TreeMap(cpy):
+
+    /** Transform statements, while maintaining import contexts and expression contexts
+     *  in the same way as Typer does. The code addresses additional concerns:
+     *   - be tail-recursive where possible
+     *   - don't re-allocate trees where nothing has changed
+     */
+    override def transformStats(stats: List[Tree], exprOwner: Symbol)(using Context): List[Tree] =
+
+      @tailrec def traverse(curStats: List[Tree])(using Context): List[Tree] =
+
+        def recur(stats: List[Tree], changed: Tree, rest: List[Tree])(using Context): List[Tree] =
+          if stats eq curStats then
+            val rest1 = transformStats(rest, exprOwner)
+            changed match
+              case Thicket(trees) => trees ::: rest1
+              case tree => tree :: rest1
+          else
+            stats.head :: recur(stats.tail, changed, rest)
+
+        curStats match
+          case stat :: rest =>
+            val statCtx = stat match
+              case _: DefTree | _: ImportOrExport => ctx
+              case _ => ctx.exprContext(stat, exprOwner)
+            val restCtx = stat match
+              case stat: Import => ctx.importContext(stat, stat.symbol)
+              case _ => ctx
+            val stat1 = transform(stat)(using statCtx)
+            if stat1 ne stat then recur(stats, stat1, rest)(using restCtx)
+            else traverse(rest)(using restCtx)
+          case nil =>
+            stats
+
+      traverse(stats)
+    end transformStats
+
+  end TreeMapWithPreciseStatContexts
+
   /** Map Inlined nodes, NamedArgs, Blocks with no statements and local references to underlying arguments.
    *  Also drops Inline and Block with no statements.
    */
-  class MapToUnderlying extends TreeMap {
+  private class MapToUnderlying extends TreeMap {
     override def transform(tree: Tree)(using Context): Tree = tree match {
       case tree: Ident if isBinding(tree.symbol) && skipLocal(tree.symbol) =>
         tree.symbol.defTree match {

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

@@ -29,19 +29,11 @@ abstract class MacroTransform extends Phase {
    */
   protected def transformPhase(using Context): Phase = this
 
-  class Transformer extends TreeMap(cpy = cpyBetweenPhases) {
+  class Transformer extends TreeMapWithPreciseStatContexts(cpy = cpyBetweenPhases):
 
-    protected def localCtx(tree: Tree)(using Context): FreshContext = 
+    protected def localCtx(tree: Tree)(using Context): FreshContext =
       ctx.fresh.setTree(tree).setOwner(localOwner(tree))
 
-    override def transformStats(trees: List[Tree], exprOwner: Symbol)(using Context): List[Tree] = {
-      def transformStat(stat: Tree): Tree = stat match {
-        case _: Import | _: DefTree => transform(stat)
-        case _ => transform(stat)(using ctx.exprContext(stat, exprOwner))
-      }
-      flatten(trees.mapconserve(transformStat(_)))
-    }
-
     override def transform(tree: Tree)(using Context): Tree =
       try
         tree match {
@@ -67,5 +59,5 @@ abstract class MacroTransform extends Phase {
 
     def transformSelf(vd: ValDef)(using Context): ValDef =
       cpy.ValDef(vd)(tpt = transform(vd.tpt))
-  }
+  end Transformer
 }