3 files changed, 18 insertions, 16 deletions
diff --git a/src/dotty/tools/dotc/CompilationUnit.scala b/src/dotty/tools/dotc/CompilationUnit.scala
index f0ce19320..f5899f0f2 100644
--- a/src/dotty/tools/dotc/CompilationUnit.scala
+++ b/src/dotty/tools/dotc/CompilationUnit.scala
@@ -24,18 +24,4 @@ class CompilationUnit(val source: SourceFile) {
    * Subsequent phases can add new sections.
    */
   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 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`
-   */
-  var addrOfSym: Symbol => Option[Addr] = (_ => None)
 }
diff --git a/src/dotty/tools/dotc/core/pickling/TastyPickler.scala b/src/dotty/tools/dotc/core/pickling/TastyPickler.scala
index 6bd6f1c44..74cc72b8c 100644
--- a/src/dotty/tools/dotc/core/pickling/TastyPickler.scala
+++ b/src/dotty/tools/dotc/core/pickling/TastyPickler.scala
@@ -7,6 +7,8 @@ import PickleFormat._
 import collection.mutable
 import TastyBuffer._
 import java.util.UUID
+import core.Symbols.Symbol
+import ast.tpd
 
 class TastyPickler {
 
@@ -51,4 +53,18 @@ class TastyPickler {
     assert(all.length == totalSize && all.bytes.length == totalSize, s"totalSize = $totalSize, all.length = ${all.length}, all.bytes.length = ${all.bytes.length}")
     all.bytes
   }
+
+  /**
+   * 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 directly 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)
 }
diff --git a/src/dotty/tools/dotc/transform/Pickler.scala b/src/dotty/tools/dotc/transform/Pickler.scala
index 5da92c079..50a419589 100644
--- a/src/dotty/tools/dotc/transform/Pickler.scala
+++ b/src/dotty/tools/dotc/transform/Pickler.scala
@@ -46,8 +46,8 @@ class Pickler extends Phase {
       unit.picklers += (cls -> pickler)
       val treePkl = new TreePickler(pickler)
       treePkl.pickle(tree :: Nil)
-      unit.addrOfTree = treePkl.buf.addrOfTree
-      unit.addrOfSym = treePkl.addrOfSym
+      pickler.addrOfTree = treePkl.buf.addrOfTree
+      pickler.addrOfSym = treePkl.addrOfSym
       if (unit.source.exists)
         pickleSourcefile(pickler, unit.source)
       if (tree.pos.exists)