Keep package structure in rootTreeOrProvider

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -39,6 +39,7 @@ class Compiler {
     List(new sbt.ExtractDependencies) :: // Sends information on classes' dependencies to sbt via callbacks
     List(new PostTyper) ::          // Additional checks and cleanups after type checking
     List(new sbt.ExtractAPI) ::     // Sends a representation of the API of classes to sbt via callbacks
+    List(new SetRootTree) ::        // Set the `rootTreeOrProvider` on class symbols
     Nil
 
   /** Phases dealing with TASTY tree pickling and unpickling */

compiler/src/dotty/tools/dotc/Run.scala

@@ -200,7 +200,8 @@ class Run(comp: Compiler, ictx: Context) extends ImplicitRunInfo with Constraint
 
   /** Enter top-level definitions of classes and objects contain in Scala source file `file`.
    *  The newly added symbols replace any previously entered symbols.
-   *  If `typeCheck = true`, also run typer on the compilation unit.
+   *  If `typeCheck = true`, also run typer on the compilation unit, and set
+   *  `rootTreeOrProvider`.
    */
   def lateCompile(file: AbstractFile, typeCheck: Boolean)(implicit ctx: Context): Unit =
     if (!files.contains(file) && !lateFiles.contains(file)) {
@@ -211,9 +212,13 @@ class Run(comp: Compiler, ictx: Context) extends ImplicitRunInfo with Constraint
           if (unit.isJava) new JavaParser(unit.source).parse()
           else new Parser(unit.source).parse()
         ctx.typer.lateEnter(unit.untpdTree)
-        def typeCheckUnit() = unit.tpdTree = ctx.typer.typedExpr(unit.untpdTree)
+        def processUnit() = {
+          unit.tpdTree = ctx.typer.typedExpr(unit.untpdTree)
+          val phase = new transform.SetRootTree()
+          phase.run
+        }
         if (typeCheck)
-          if (compiling) finalizeActions += (() => typeCheckUnit()) else typeCheckUnit()
+          if (compiling) finalizeActions += (() => processUnit()) else processUnit()
       }
       process()(runContext.fresh.setCompilationUnit(unit))
     }

compiler/src/dotty/tools/dotc/interactive/InteractiveCompiler.scala

@@ -13,6 +13,7 @@ class InteractiveCompiler extends Compiler {
   // after each phase group instead of waiting for the pipeline to finish.
   override def phases: List[List[Phase]] = List(
     List(new FrontEnd),
+    List(new transform.SetRootTree),
     List(new transform.CookComments)
   )
 }

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

@@ -0,0 +1,44 @@
+package dotty.tools.dotc.transform
+
+import dotty.tools.dotc.CompilationUnit
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.core.Phases.Phase
+
+/** Set the `rootTreeOrProvider` property of class symbols. */
+class SetRootTree extends Phase {
+
+  override val phaseName: String = SetRootTree.name
+  override def isRunnable(implicit ctx: Context) =
+    super.isRunnable && ctx.settings.YretainTrees.value
+
+  override def run(implicit ctx: Context): Unit = {
+    val tree = ctx.compilationUnit.tpdTree
+    traverser.traverse(tree)
+  }
+
+  private def traverser = new tpd.TreeTraverser {
+    override def traverse(tree: tpd.Tree)(implicit ctx: Context): Unit = {
+      tree match {
+        case pkg: tpd.PackageDef =>
+          traverseChildren(pkg)
+        case td: tpd.TypeDef =>
+          if (td.symbol.isClass) {
+            val sym = td.symbol.asClass
+            tpd.sliceTopLevel(ctx.compilationUnit.tpdTree, sym) match {
+              case (pkg: tpd.PackageDef) :: Nil =>
+                sym.rootTreeOrProvider = pkg
+              case _ =>
+                sym.rootTreeOrProvider = td
+            }
+          }
+        case _ =>
+          ()
+      }
+    }
+  }
+}
+
+object SetRootTree {
+  val name: String = "SetRootTree"
+}

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

@@ -1642,7 +1642,11 @@ class Typer extends Namer
       // check value class constraints
       checkDerivedValueClass(cls, body1)
 
-      if (ctx.settings.YretainTrees.value) cls.rootTreeOrProvider = cdef1
+
+      // Temporarily set the typed class def as root tree so that we have at least some
+      // information in the IDE in case we never reach `SetRootTree`.
+      if (ctx.mode.is(Mode.Interactive) && ctx.settings.YretainTrees.value)
+        cls.rootTreeOrProvider = cdef1
 
       cdef1
 

language-server/test/dotty/tools/languageserver/RenameTest.scala

@@ -236,6 +236,19 @@ class RenameTest {
 
   }
 
+  @Test def renameImportFromTasty: Unit = {
+    // Note that everything here is in the empty package; this ensures that we will
+    // use the sourcefile loader to load `class Bar`.
+    def testRename(m: CodeMarker) = {
+      withSources(
+        code"""object O { class ${m1}Foo${m2} }""",
+        tasty"""import O.${m3}Foo${m4}
+                class Bar extends ${m5}Foo${m6}"""
+      ).rename(m, "NewName", Set(m1 to m2, m3 to m4, m5 to m6))
+    }
 
+    testRename(m1)
+    testRename(m2)
+  }
 
 }