Add unpickled source file path as annotation to root symbols

1 files changed, 8 insertions, 4 deletions

diff --git a/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala b/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala
index 6d0e73efb..edaeacd9a 100644
--- a/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala
+++ b/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala
@@ -3,13 +3,14 @@ package dotc
 package core
 package tasty
 
-import Contexts._, SymDenotations._
+import Contexts._, SymDenotations._, Symbols._
 import dotty.tools.dotc.ast.tpd
 import TastyUnpickler._, TastyBuffer._
 import dotty.tools.dotc.core.tasty.DottyUnpickler.{SourceFileUnpickler, TreeSectionUnpickler, PositionsSectionUnpickler}
 import util.Positions._
 import util.{SourceFile, NoSource}
 import PositionUnpickler._
+import Annotations.Annotation
 import classfile.ClassfileParser
 
 object DottyUnpickler {
@@ -43,21 +44,24 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded {
 
   val unpickler = new TastyUnpickler(bytes)
   private val treeUnpickler = unpickler.unpickle(new TreeSectionUnpickler).get
+  private val source = unpickler.unpickle(new SourceFileUnpickler).getOrElse(NoSource)
 
   /** Enter all toplevel classes and objects into their scopes
    *  @param roots          a set of SymDenotations that should be overwritten by unpickling
    */
-  def enter(roots: Set[SymDenotation])(implicit ctx: Context): Unit =
+  def enter(roots: Set[SymDenotation])(implicit ctx: Context): Unit = {
     treeUnpickler.enterTopLevel(roots)
+    if (source.exists)
+      for (root <- roots) root.addAnnotation(Annotation.makeSourceFile(source.path))
+  }
 
   /** The unpickled trees, and the source file they come from
    *  @param readPositions if true, trees get decorated with position information.
    */
   def body(readPositions: Boolean = false)(implicit ctx: Context): (List[Tree], SourceFile) = {
-    val source = unpickler.unpickle(new SourceFileUnpickler).getOrElse(NoSource)
     if (readPositions)
       for ((totalRange, positions) <- unpickler.unpickle(new PositionsSectionUnpickler))
         treeUnpickler.usePositions(totalRange, positions)
-    (treeUnpickler.unpickle(), source)
+    (treeUnpickler.unpickle(source), source)
   }
 }