Populate addresses of symbols, types and trees after pickler

src/dotty/tools/dotc/CompilationUnit.scala

@@ -1,8 +1,12 @@
 package dotty.tools
 package dotc
 
+import dotty.tools.dotc.core.Types.Type
+import dotty.tools.dotc.core.pickling.{TastyBuffer, TastyPickler}
 import util.SourceFile
 import ast.{tpd, untpd}
+import TastyBuffer._
+import dotty.tools.dotc.core.Symbols._
 
 class CompilationUnit(val source: SourceFile) {
 
@@ -13,6 +17,25 @@ class CompilationUnit(val source: SourceFile) {
   var tpdTree: tpd.Tree = tpd.EmptyTree
 
   def isJava = source.file.name.endsWith(".java")
-
-  var pickled: Array[Byte] = Array()
+
+  /**
+   * Pickler used to create TASTY sections.
+   * Sections: Header, ASTs and Positions are populated by `pickler` phase.
+   * Subsequent phases can add new sections.
+   */
+  lazy val pickler: TastyPickler = new TastyPickler()
+
+  /**
+   * 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`
+   */
+  var addrOfTree: tpd.Tree => Option[Addr] = (_ => None)
+
+  /**
+   * 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`
+   */
+  var addrOfSym: Symbol => Option[Addr] = (_ => None)
 }

src/dotty/tools/dotc/core/pickling/TreePickler.scala

@@ -26,7 +26,11 @@ class TreePickler(pickler: TastyPickler) {
     op
     fillRef(lengthAddr, currentAddr, relative = true)
   }
-
+
+  def addrOfSym(sym: Symbol): Option[Addr] = {
+    symRefs.get(sym)
+  }
+
   private var makeSymbolicRefsTo: Symbol = NoSymbol
 
   /** All references to members of class `sym` are pickled
@@ -531,8 +535,12 @@ class TreePickler(pickler: TastyPickler) {
       withLength { pickleType(ann.symbol.typeRef); pickleTree(ann.tree) }
     }
 
+    def updateMapWithDeltas[T](mp: collection.mutable.Map[T, Addr]) =
+      for (key <- mp.keysIterator.toBuffer[T]) mp(key) = adjusted(mp(key))
+
     trees.foreach(tree => if (!tree.isEmpty) pickleTree(tree))
     assert(forwardSymRefs.isEmpty, i"unresolved symbols: ${forwardSymRefs.keySet.toList}%, %")
     compactify()
-  }  
+    updateMapWithDeltas(symRefs)
+  }
 }

src/dotty/tools/dotc/transform/NormalizeFlags.scala

@@ -15,7 +15,7 @@ import Flags._, Symbols._
  *     or alias type definition or a deferred val or def.
  */
 class NormalizeFlags extends MiniPhaseTransform with SymTransformer { thisTransformer =>
-  override def phaseName = "elimLocals"
+  override def phaseName = "normalizeFlags"
 
   def transformSym(ref: SymDenotation)(implicit ctx: Context) = {
     var newFlags = ref.flags &~ Local

src/dotty/tools/dotc/transform/Pickler.scala

@@ -11,7 +11,7 @@ import Periods._
 import Phases._
 import collection.mutable
 
-/** This miniphase pickles trees */
+/** This phase pickles trees */
 class Pickler extends Phase {
   import ast.tpd._
 
@@ -32,19 +32,20 @@ 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.pickler
       val treePkl = new TreePickler(pickler)
       treePkl.pickle(tree :: Nil)
+      unit.addrOfTree = treePkl.buf.addrOfTree
+      unit.addrOfSym = treePkl.addrOfSym
       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.pickler.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 +61,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.pickler.assembleParts())
         unpickler.enter(roots = Set())
         unpickler
       }