Maintain source files in pickled info

So far: Only one source file is recorded. Should evaluate
whether more are needed. Will programs composed from several
source files be pickled? They will certainly be generated after
inlining, but maybe all that happens after pickling?

5 files changed, 31 insertions, 17 deletions

diff --git a/src/dotty/tools/dotc/CompilationUnit.scala b/src/dotty/tools/dotc/CompilationUnit.scala
index 4f8c30aab..f0ce19320 100644
--- a/src/dotty/tools/dotc/CompilationUnit.scala
+++ b/src/dotty/tools/dotc/CompilationUnit.scala
@@ -25,14 +25,14 @@ class CompilationUnit(val source: SourceFile) {
    */
   var picklers: Map[ClassSymbol, TastyPickler] = Map()
 
-  /**
+  /** TODO: I'd prefer we do not put this in CompilationUnit
    * Addresses in TASTY file of trees, stored by pickling.
    * Note that trees are checked for reference equality,
-   * so one can reliably use this function only dirrectly after `pickler`
+   * so one can reliably use this function only directly after `pickler`
    */
   var addrOfTree: tpd.Tree => Option[Addr] = (_ => None)
 
-  /**
+  /** TODO: I'd prefer we do not put this in CompilationUnit
    * Addresses in TASTY file of symbols, stored by pickling.
    * Note that trees are checked for reference equality,
    * so one can reliably use this function only dirrectly after `pickler`
diff --git a/src/dotty/tools/dotc/FromTasty.scala b/src/dotty/tools/dotc/FromTasty.scala
index 844f5fe76..2c704efc2 100644
--- a/src/dotty/tools/dotc/FromTasty.scala
+++ b/src/dotty/tools/dotc/FromTasty.scala
@@ -66,13 +66,9 @@ object FromTasty extends Driver {
               case info: ClassfileLoader =>
                 info.load(clsd) match {
                   case Some(unpickler: DottyUnpickler) =>
-                    val (List(unpickled), sources) = unpickler.body(readPositions = false)
-                    val unit1 = sources.get(unpickled) match {
-                      case Some(source) => new CompilationUnit(source)
-                      case _ => unit
-                    }
+                    val (List(unpickled), source) = unpickler.body(readPositions = false)
+                    val unit1 = new CompilationUnit(source)
                     unit1.tpdTree = unpickled
-                    unit1.embeddedSources = sources - unpickled
                     unit1
                   case _ =>
                     cannotUnpickle(s"it does not have a TASTY attribute")
diff --git a/src/dotty/tools/dotc/core/pickling/DottyUnpickler.scala b/src/dotty/tools/dotc/core/pickling/DottyUnpickler.scala
index 6e7bd9210..2d8f571ec 100644
--- a/src/dotty/tools/dotc/core/pickling/DottyUnpickler.scala
+++ b/src/dotty/tools/dotc/core/pickling/DottyUnpickler.scala
@@ -7,6 +7,7 @@ import Contexts._, SymDenotations._
 import dotty.tools.dotc.ast.tpd
 import TastyUnpickler._, TastyBuffer._
 import util.Positions._
+import util.{SourceFile, NoSource}
 import PositionUnpickler._
 
 object DottyUnpickler {
@@ -30,14 +31,20 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded {
   def enter(roots: Set[SymDenotation])(implicit ctx: Context): Unit =
     treeUnpickler.enterTopLevel(roots)
 
-  /** The unpickled trees
+  /** 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] = {
+  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()))
+      for ((totalRange, positions) <- unpickler.unpickle(new PositionsSectionUnpickler))
         treeUnpickler.usePositions(totalRange, positions)
-    treeUnpickler.unpickle()
+    (treeUnpickler.unpickle(), source)
+  }
+
+  private class SourceFileUnpickler extends SectionUnpickler[SourceFile]("Sourcefile") {
+    def unpickle(reader: TastyReader, tastyName: TastyName.Table) =
+      new SourceFile(tastyName(reader.readNameRef()).toString, Seq())
   }
 
   private class TreeSectionUnpickler extends SectionUnpickler[TreeUnpickler]("ASTs") {
diff --git a/src/dotty/tools/dotc/core/pickling/PickleFormat.scala b/src/dotty/tools/dotc/core/pickling/PickleFormat.scala
index d12a879ba..4f3841212 100644
--- a/src/dotty/tools/dotc/core/pickling/PickleFormat.scala
+++ b/src/dotty/tools/dotc/core/pickling/PickleFormat.scala
@@ -185,6 +185,8 @@ Note: Tree tags are grouped into 5 categories that determine what follows, and t
   Category 4 (tags 112-127):  tag Nat AST
   Category 5 (tags 128-255):  tag Length <payload>
 
+Standard Section: "Sourcefile" sourcefile_NameRef
+
 Standard Section: "Positions" sourceLength_Nat Assoc*
 
   Assoc         = addr_Delta offset_Delta offset_Delta?
diff --git a/src/dotty/tools/dotc/transform/Pickler.scala b/src/dotty/tools/dotc/transform/Pickler.scala
index c09ba6889..d30c876bd 100644
--- a/src/dotty/tools/dotc/transform/Pickler.scala
+++ b/src/dotty/tools/dotc/transform/Pickler.scala
@@ -11,6 +11,7 @@ import Periods._
 import Phases._
 import Symbols._
 import Flags.Module
+import util.SourceFile
 import collection.mutable
 
 /** This phase pickles trees */
@@ -47,6 +48,8 @@ class Pickler extends Phase {
       treePkl.pickle(tree :: Nil)
       unit.addrOfTree = treePkl.buf.addrOfTree
       unit.addrOfSym = treePkl.addrOfSym
+      if (unit.source.exists)
+        pickleSourcefile(pickler, unit.source)
       if (tree.pos.exists)
         new PositionPickler(pickler, treePkl.buf.addrOfTree).picklePositions(tree :: Nil, tree.pos)
 
@@ -62,6 +65,12 @@ class Pickler extends Phase {
     }
   }
 
+  private def pickleSourcefile(pickler: TastyPickler, source: SourceFile): Unit = {
+    val buf = new TastyBuffer(10)
+    pickler.newSection("Sourcefile", buf)
+    buf.writeNat(pickler.nameBuffer.nameIndex(source.file.path).index)
+  }
+
   override def runOn(units: List[CompilationUnit])(implicit ctx: Context): List[CompilationUnit] = {
     val result = super.runOn(units)
     if (ctx.settings.YtestPickler.value)
@@ -80,16 +89,16 @@ class Pickler extends Phase {
       }
     pickling.println("************* entered toplevel ***********")
     for ((cls, unpickler) <- unpicklers) {
-      val unpickled = unpickler.body(readPositions = false)
-      testSame(i"$unpickled%\n%", beforePickling(cls), cls)
+      val (unpickled, source) = unpickler.body(readPositions = false)
+      testSame(i"$unpickled%\n%", beforePickling(cls), cls, source)
     }
   }
 
-  private def testSame(unpickled: String, previous: String, cls: ClassSymbol)(implicit ctx: Context) =
+  private def testSame(unpickled: String, previous: String, cls: ClassSymbol, source: SourceFile)(implicit ctx: Context) =
     if (previous != unpickled) {
       output("before-pickling.txt", previous)
       output("after-pickling.txt", unpickled)
-      ctx.error(s"""pickling difference for ${cls.fullName}, for details:
+      ctx.error(s"""pickling difference for ${cls.fullName} in $source, for details:
                    |
                    |  diff before-pickling.txt after-pickling.txt""".stripMargin)
     }