Export pickler and pickled addresses so that other phases can pickle their info

2 files changed, 10 insertions, 6 deletions

diff --git a/src/dotty/tools/dotc/CompilationUnit.scala b/src/dotty/tools/dotc/CompilationUnit.scala
index 44f9b7e5a..c2cd9361c 100644
--- a/src/dotty/tools/dotc/CompilationUnit.scala
+++ b/src/dotty/tools/dotc/CompilationUnit.scala
@@ -1,8 +1,10 @@
 package dotty.tools
 package dotc
 
+import dotty.tools.dotc.core.pickling.{TastyBuffer, TastyPickler}
 import util.SourceFile
 import ast.{tpd, untpd}
+import TastyBuffer._
 
 class CompilationUnit(val source: SourceFile) {
 
@@ -14,5 +16,7 @@ class CompilationUnit(val source: SourceFile) {
 
   def isJava = source.file.name.endsWith(".java")
   
-  var pickled: Array[Byte] = Array()
+  lazy val pickled: TastyPickler = new TastyPickler()
+
+  var addrOfTree: tpd.Tree => Option[Addr] = (_ => None)
 }
\ No newline at end of file
diff --git a/src/dotty/tools/dotc/transform/Pickler.scala b/src/dotty/tools/dotc/transform/Pickler.scala
index fb6b58e36..0d8acf38a 100644
--- a/src/dotty/tools/dotc/transform/Pickler.scala
+++ b/src/dotty/tools/dotc/transform/Pickler.scala
@@ -32,19 +32,19 @@ class Pickler extends Phase {
       pickling.println(i"unpickling in run ${ctx.runId}")
       if (ctx.settings.YtestPickler.value) beforePickling(unit) = tree.show
 
-      val pickler = new TastyPickler            
+      val pickler = unit.pickled
       val treePkl = new TreePickler(pickler)
       treePkl.pickle(tree :: Nil)
+      unit.addrOfTree = treePkl.buf.addrOfTree
       if (tree.pos.exists)
         new PositionPickler(pickler, treePkl.buf.addrOfTree).picklePositions(tree :: Nil, tree.pos)
 
-      unit.pickled = pickler.assembleParts()
       def rawBytes = // not needed right now, but useful to print raw format.
-        unit.pickled.iterator.grouped(10).toList.zipWithIndex.map {
+        unit.pickled.assembleParts().iterator.grouped(10).toList.zipWithIndex.map {
           case (row, i) => s"${i}0: ${row.mkString(" ")}"
         }
       // println(i"rawBytes = \n$rawBytes%\n%") // DEBUG
-      if (pickling ne noPrinter) new TastyPrinter(unit.pickled).printContents()
+      if (pickling ne noPrinter) new TastyPrinter(pickler.assembleParts()).printContents()
     } 
   }
   
@@ -60,7 +60,7 @@ class Pickler extends Phase {
     ctx.definitions.init
     val unpicklers = 
       for (unit <- units) yield {
-        val unpickler = new DottyUnpickler(unit.pickled)
+        val unpickler = new DottyUnpickler(unit.pickled.assembleParts())
         unpickler.enter(roots = Set())
         unpickler
       }