Address reviewer comments on #416

3 files changed, 22 insertions, 11 deletions

diff --git a/src/dotty/tools/dotc/CompilationUnit.scala b/src/dotty/tools/dotc/CompilationUnit.scala
index 2f5f1202d..de51a84cf 100644
--- a/src/dotty/tools/dotc/CompilationUnit.scala
+++ b/src/dotty/tools/dotc/CompilationUnit.scala
@@ -17,10 +17,25 @@ class CompilationUnit(val source: SourceFile) {
   var tpdTree: tpd.Tree = tpd.EmptyTree
 
   def isJava = source.file.name.endsWith(".java")
-  
-  lazy val pickled: TastyPickler = new TastyPickler()
 
+  /**
+   * 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)
 }
\ No newline at end of file
diff --git a/src/dotty/tools/dotc/core/pickling/TreePickler.scala b/src/dotty/tools/dotc/core/pickling/TreePickler.scala
index e23efd860..c5e07ffa9 100644
--- a/src/dotty/tools/dotc/core/pickling/TreePickler.scala
+++ b/src/dotty/tools/dotc/core/pickling/TreePickler.scala
@@ -535,12 +535,8 @@ class TreePickler(pickler: TastyPickler) {
       withLength { pickleType(ann.symbol.typeRef); pickleTree(ann.tree) }
     }
 
-    def updateMapWithDeltas[T](mp: collection.mutable.Map[T, Addr]) = {
-      mp.map{
-        case (key, addr) => (key, adjusted(addr))
-      }.foreach(mp += _)
-    }
-
+    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}%, %")
diff --git a/src/dotty/tools/dotc/transform/Pickler.scala b/src/dotty/tools/dotc/transform/Pickler.scala
index c7a308f93..e238f21ef 100644
--- a/src/dotty/tools/dotc/transform/Pickler.scala
+++ b/src/dotty/tools/dotc/transform/Pickler.scala
@@ -32,7 +32,7 @@ class Pickler extends Phase {
       pickling.println(i"unpickling in run ${ctx.runId}")
       if (ctx.settings.YtestPickler.value) beforePickling(unit) = tree.show
 
-      val pickler = unit.pickled
+      val pickler = unit.pickler
       val treePkl = new TreePickler(pickler)
       treePkl.pickle(tree :: Nil)
       unit.addrOfTree = treePkl.buf.addrOfTree
@@ -41,7 +41,7 @@ class Pickler extends Phase {
         new PositionPickler(pickler, treePkl.buf.addrOfTree).picklePositions(tree :: Nil, tree.pos)
 
       def rawBytes = // not needed right now, but useful to print raw format.
-        unit.pickled.assembleParts().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
@@ -61,7 +61,7 @@ class Pickler extends Phase {
     ctx.definitions.init
     val unpicklers = 
       for (unit <- units) yield {
-        val unpickler = new DottyUnpickler(unit.pickled.assembleParts())
+        val unpickler = new DottyUnpickler(unit.pickler.assembleParts())
         unpickler.enter(roots = Set())
         unpickler
       }