Record pickled trees in a hashmap

The idea is that we want to use the mapping from tree to Addr
in other sections, most immeditaely for positions, but it could
be others as well.

3 files changed, 26 insertions, 4 deletions

diff --git a/src/dotty/tools/dotc/core/pickling/TreeBuffer.scala b/src/dotty/tools/dotc/core/pickling/TreeBuffer.scala
index 287e2f334..c1eae5014 100644
--- a/src/dotty/tools/dotc/core/pickling/TreeBuffer.scala
+++ b/src/dotty/tools/dotc/core/pickling/TreeBuffer.scala
@@ -6,6 +6,7 @@ package pickling
 import util.Util.{bestFit, dble}
 import TastyBuffer.{Addr, AddrWidth}
 import config.Printers.pickling
+import ast.tpd.Tree
 
 class TreeBuffer extends TastyBuffer(1000000) {
 
@@ -16,6 +17,13 @@ class TreeBuffer extends TastyBuffer(1000000) {
   private var isRelative = new Array[Boolean](initialOffsetSize)
   private var delta: Array[Int] = _
   private var numOffsets = 0
+
+  private[pickling] val pickledTrees = new java.util.IdentityHashMap[Tree, Any] // Value type is really Addr, but that's not compatible with null  
+  
+  def addrOfTree(tree: Tree): Option[Addr] = pickledTrees.get(tree) match {
+    case null => None
+    case n => Some(n.asInstanceOf[Addr])
+  }
             
   private def offset(i: Int): Addr = Addr(offsets(i))
 
@@ -74,7 +82,7 @@ class TreeBuffer extends TastyBuffer(1000000) {
     }
   }
   
-  /** The absoluate or relative adjusted address at index `i` of `offsets` array*/
+  /** The absolute or relative adjusted address at index `i` of `offsets` array*/
   private def adjustedOffset(i: Int): Addr = {
     val at = offset(i)
     val original = getAddr(at)
@@ -140,13 +148,21 @@ class TreeBuffer extends TastyBuffer(1000000) {
     wasted
   }
   
+  def adjustPickledTrees(): Unit = {
+    val it = pickledTrees.keySet.iterator
+    while (it.hasNext) {
+      val tree = it.next
+      pickledTrees.put(tree, adjusted(pickledTrees.get(tree).asInstanceOf[Addr]))
+    }
+  }
+  
   /** Final assembly, involving the following steps:
    *   - compute deltas
    *   - adjust deltas until additional savings are < 1% of total
    *   - adjust offsets according to the adjusted deltas
    *   - shrink buffer, skipping zeroes.
    */
-  override def assemble(): Unit = {
+  def compactify(): Unit = {
     val origLength = length
     computeDeltas()
     //println(s"offsets: ${offsets.take(numOffsets).deep}")
@@ -157,6 +173,7 @@ class TreeBuffer extends TastyBuffer(1000000) {
       pickling.println(s"adjusting deltas, saved = $saved")
     } while (saved > 0 && length / saved < 100)
     adjustOffsets()
+    adjustPickledTrees()
     val wasted = compress()
     pickling.println(s"original length: $origLength, compressed to: $length, wasted: $wasted") // DEBUG, for now.
   }
diff --git a/src/dotty/tools/dotc/core/pickling/TreePickler.scala b/src/dotty/tools/dotc/core/pickling/TreePickler.scala
index 32eb5f47f..0e6c45d30 100644
--- a/src/dotty/tools/dotc/core/pickling/TreePickler.scala
+++ b/src/dotty/tools/dotc/core/pickling/TreePickler.scala
@@ -10,7 +10,7 @@ import Contexts._, Symbols._, Types._, Names._, Constants._, Decorators._, Annot
 import collection.mutable
 import TastyBuffer._
 
-class TreePickler(pickler: TastyPickler, picklePositions: Boolean) {
+class TreePickler(pickler: TastyPickler) {
   val buf = new TreeBuffer
   pickler.newSection("ASTs", buf)
   import buf._
@@ -244,6 +244,7 @@ class TreePickler(pickler: TastyPickler, picklePositions: Boolean) {
       if (!tree.isEmpty) pickleTree(tree)
 
     def pickleTree(tree: Tree): Unit = try {
+      pickledTrees.put(tree, currentAddr)
       tree match {
       case Ident(name) =>
         tree.tpe match {
@@ -468,5 +469,6 @@ class TreePickler(pickler: TastyPickler, picklePositions: Boolean) {
 
     pickleTree(tree)
     assert(forwardSymRefs.isEmpty, i"unresolved symbols: ${forwardSymRefs.keySet.toList}%, %")
+    compactify()
   }  
 }
diff --git a/src/dotty/tools/dotc/transform/Pickler.scala b/src/dotty/tools/dotc/transform/Pickler.scala
index c21910e65..f080aa7ad 100644
--- a/src/dotty/tools/dotc/transform/Pickler.scala
+++ b/src/dotty/tools/dotc/transform/Pickler.scala
@@ -17,12 +17,15 @@ class Pickler extends MiniPhaseTransform { thisTransform =>
   override def transformUnit(tree: Tree)(implicit ctx: Context, info: TransformerInfo): Tree = {
     if (!ctx.compilationUnit.isJava) {
       val pickler = new TastyPickler
-      new TreePickler(pickler, picklePositions = false).pickle(tree)
+      
+      val treePkl = new TreePickler(pickler)
+      treePkl.pickle(tree)
       val bytes = pickler.assembleParts()
       def rawBytes = // not needed right now, but useful to print raw format.
         bytes.iterator.grouped(10).toList.zipWithIndex.map {
           case (row, i) => s"${i}0: ${row.mkString(" ")}"
         }
+      // println(i"rawBytes = \n$rawBytes%\n%") // DEBUG
       if (Printers.pickling ne Printers.noPrinter) new TastyPrinter(bytes).printContents()
     }
     tree