Keep language imports around longer

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

@@ -1151,12 +1151,57 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
               case _ => tree1 :: rest1
           case nil => nil
       recur(trees, 0)
+
+    /** 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
+     */
+    inline def mapStatements(exprOwner: Symbol, inline op: Tree => Context ?=> Tree)(using Context): List[Tree] =
+      @tailrec
+      def loop(mapped: mutable.ListBuffer[Tree] | Null, unchanged: List[Tree], pending: List[Tree])(using Context): List[Tree] =
+        pending match
+          case stat :: rest =>
+            val statCtx = stat match
+              case _: DefTree | _: ImportOrExport => ctx
+              case _ => ctx.exprContext(stat, exprOwner)
+            val stat1 = op(stat)(using statCtx)
+            val restCtx = stat match
+              case stat: Import => ctx.importContext(stat, stat.symbol)
+              case _ => ctx
+            if stat1 eq stat then
+              loop(mapped, unchanged, rest)(using restCtx)
+            else
+              val buf = if mapped == null then new mutable.ListBuffer[Tree] else mapped
+              var xc = unchanged
+              while xc ne pending do
+                buf += xc.head
+                xc = xc.tail
+              stat1 match
+                case Thicket(stats1) => buf ++= stats1
+                case _ => buf += stat1
+              loop(buf, rest, rest)(using restCtx)
+          case nil =>
+            if mapped == null then unchanged
+            else mapped.prependToList(unchanged)
+
+      loop(null, trees, trees)
+    end mapStatements
   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):
+    override def transformStats(trees: List[Tree], exprOwner: Symbol)(using Context): List[Tree] =
+      trees.mapStatements(exprOwner, transform(_))
+
   /** 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/FirstTransform.scala

@@ -25,7 +25,8 @@ object FirstTransform {
 }
 
 /** The first tree transform
- *   - eliminates some kinds of trees: Imports, NamedArgs
+ *   - eliminates some kinds of trees: Imports other than language imports,
+ *     Exports, NamedArgs, type trees other than TypeTree
  *   - stubs out native methods
  *   - eliminates self tree in Template and self symbol in ClassInfo
  *   - collapses all type trees to trees of class TypeTree
@@ -58,7 +59,7 @@ class FirstTransform extends MiniPhase with InfoTransformer { thisPhase =>
             tree.symbol.is(JavaStatic) && qualTpe.derivesFrom(tree.symbol.enclosingClass),
           i"non member selection of ${tree.symbol.showLocated} from ${qualTpe} in $tree")
       case _: TypeTree =>
-      case _: Import | _: NamedArg | _: TypTree =>
+      case _: Export | _: NamedArg | _: TypTree =>
         assert(false, i"illegal tree: $tree")
       case _ =>
     }
@@ -136,7 +137,8 @@ class FirstTransform extends MiniPhase with InfoTransformer { thisPhase =>
   }
 
   override def transformOther(tree: Tree)(using Context): Tree = tree match {
-    case tree: ImportOrExport => EmptyTree
+    case tree: Import if untpd.languageImport(tree.expr).isEmpty => EmptyTree
+    case tree: Export => EmptyTree
     case tree: NamedArg => transformAllDeep(tree.arg)
     case tree => if (tree.isType) toTypeTree(tree) else tree
   }

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
 }

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

@@ -432,16 +432,10 @@ class MegaPhase(val miniPhases: Array[MiniPhase]) extends Phase {
   def transformSpecificTree[T <: Tree](tree: T, start: Int)(using Context): T =
     transformTree(tree, start).asInstanceOf[T]
 
-  def transformStats(trees: List[Tree], exprOwner: Symbol, start: Int)(using Context): List[Tree] = {
-    def transformStat(stat: Tree)(using Context): Tree = stat match {
-      case _: Import | _: DefTree => transformTree(stat, start)
-      case Thicket(stats) => cpy.Thicket(stat)(stats.mapConserve(transformStat))
-      case _ => transformTree(stat, start)(using ctx.exprContext(stat, exprOwner))
-    }
+  def transformStats(trees: List[Tree], exprOwner: Symbol, start: Int)(using Context): List[Tree] =
     val nestedCtx = prepStats(trees, start)
-    val trees1 = trees.mapInline(transformStat(_)(using nestedCtx))
+    val trees1 = trees.mapStatements(exprOwner, transformTree(_, start))(using nestedCtx)
     goStats(trees1, start)(using nestedCtx)
-  }
 
   def transformUnit(tree: Tree)(using Context): Tree = {
     val nestedCtx = prepUnit(tree, 0)

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

@@ -14,13 +14,15 @@ import StdNames.nme
 import ast.tpd
 import SymUtils._
 import config.Feature
+import Decorators.*
 
 /** This phase makes all erased term members of classes private so that they cannot
  *  conflict with non-erased members. This is needed so that subsequent phases like
  *  ResolveSuper that inspect class members work correctly.
  *  The phase also replaces all expressions that appear in an erased context by
  *  default values. This is necessary so that subsequent checking phases such
  *  as IsInstanceOfChecker don't give false negatives.
+ *  Finally, the phase drops (language-) imports.
  */
 class PruneErasedDefs extends MiniPhase with SymTransformer { thisTransform =>
   import tpd._
@@ -54,10 +56,18 @@ class PruneErasedDefs extends MiniPhase with SymTransformer { thisTransform =>
     checkErasedInExperimental(tree.symbol)
     tree
 
+  override def transformOther(tree: Tree)(using Context): Tree = tree match
+    case tree: Import => EmptyTree
+    case _ => tree
+
   def checkErasedInExperimental(sym: Symbol)(using Context): Unit =
     // Make an exception for Scala 2 experimental macros to allow dual Scala 2/3 macros under non experimental mode
     if sym.is(Erased, butNot = Macro) && sym != defn.Compiletime_erasedValue && !sym.isInExperimentalScope then
       Feature.checkExperimentalFeature("erased", sym.sourcePos)
+
+  override def checkPostCondition(tree: Tree)(using Context): Unit = tree match
+    case _: tpd.Import => assert(false, i"illegal tree: $tree")
+    case _ =>
 }
 
 object PruneErasedDefs {

compiler/src/dotty/tools/dotc/typer/Typer.scala

@@ -2556,7 +2556,7 @@ class Typer(@constructorOnly nestingLevel: Int = 0) extends Namer
                 |The selector is not a member of an object or package.""")
     else typd(imp.expr, AnySelectionProto)
 
-  def typedImport(imp: untpd.Import, sym: Symbol)(using Context): Import =
+  def typedImport(imp: untpd.Import, sym: Symbol)(using Context): Tree =
     val expr1 = typedImportQualifier(imp, typedExpr(_, _)(using ctx.withOwner(sym)))
     checkLegalImportPath(expr1)
     val selectors1 = typedSelectors(imp.selectors)