Handle shared trees

Shared trees are pickled under multiple addresses. Previously, only the
last address was stored, which led to trees with unknown positions.

Now, all addresses are stored.

3 files changed, 48 insertions, 13 deletions

diff --git a/src/dotty/tools/dotc/core/tasty/PositionPickler.scala b/src/dotty/tools/dotc/core/tasty/PositionPickler.scala
index 65e9d12be..b528b8aaf 100644
--- a/src/dotty/tools/dotc/core/tasty/PositionPickler.scala
+++ b/src/dotty/tools/dotc/core/tasty/PositionPickler.scala
@@ -13,12 +13,29 @@ import collection.mutable
 import TastyBuffer._
 import util.Positions._
 
-class PositionPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr]) {
+class PositionPickler(pickler: TastyPickler, addrsOfTree: tpd.Tree => List[Addr]) {
   val buf = new TastyBuffer(5000)
   pickler.newSection("Positions", buf)
   import buf._
   import ast.tpd._
 
+  private val remainingAddrs = new java.util.IdentityHashMap[Tree, Iterator[Addr]]
+
+  def nextTreeAddr(tree: Tree): Option[Addr] = remainingAddrs.get(tree) match {
+    case null =>
+      addrsOfTree(tree) match {
+        case Nil =>
+          None
+        case addr :: Nil =>
+          Some(addr)
+        case addrs =>
+          remainingAddrs.put(tree, addrs.iterator)
+          nextTreeAddr(tree)
+      }
+    case it: Iterator[_] =>
+      if (it.hasNext) Some(it.next) else None
+  }
+
   def header(addrDelta: Int, hasStartDelta: Boolean, hasEndDelta: Boolean) = {
     def toInt(b: Boolean) = if (b) 1 else 0
     (addrDelta << 2) | (toInt(hasStartDelta) << 1) | toInt(hasEndDelta)
@@ -40,9 +57,9 @@ class PositionPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr
     def traverse(x: Any): Unit = x match {
       case x: Tree @unchecked =>
         if (x.pos.exists /*&& x.pos.toSynthetic != x.initialPos.toSynthetic*/) {
-          addrOfTree(x) match {
+          nextTreeAddr(x) match {
             case Some(addr) =>
-              //println(i"pickling $x")
+              //println(i"pickling $x ar $addr")
               pickleDeltas(addr.index, x.pos)
             case _ =>
               //println(i"no address for $x")
diff --git a/src/dotty/tools/dotc/core/tasty/TastyPickler.scala b/src/dotty/tools/dotc/core/tasty/TastyPickler.scala
index 98b0dc7c6..f847dda68 100644
--- a/src/dotty/tools/dotc/core/tasty/TastyPickler.scala
+++ b/src/dotty/tools/dotc/core/tasty/TastyPickler.scala
@@ -57,9 +57,11 @@ class 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 directly after `pickler`
+   * so one can reliably use this function only directly after `pickler`.
+   * Note that a tree can have several addresses, if it is shared,
+   * i.e. accessible from different paths. Any such sharing is undone by pickling.
    */
-  var addrOfTree: tpd.Tree => Option[Addr] = (_ => None)
+  var addrsOfTree: tpd.Tree => List[Addr] = (_ => Nil)
 
   /**
    * Addresses in TASTY file of symbols, stored by pickling.
diff --git a/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala b/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala
index b9e1d2b45..f2681ecde 100644
--- a/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala
+++ b/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala
@@ -17,13 +17,26 @@ class TreeBuffer extends TastyBuffer(50000) {
   private var delta: Array[Int] = _
   private var numOffsets = 0
 
-  private val treeAddr = new java.util.IdentityHashMap[Tree, Any] // Value type is really Addr, but that's not compatible with null
+  private type TreeAddrs = Any // really: Addr | List[Addr]
 
-  def registerTreeAddr(tree: Tree) = treeAddr.put(tree, currentAddr)
-
-  def addrOfTree(tree: Tree): Option[Addr] = treeAddr.get(tree) match {
-    case null => None
-    case n => Some(n.asInstanceOf[Addr])
+  /** A map from trees to the address(es) at which a tree is pickled. There may be several
+   *  such addresses if the tree is shared. To keep the map compact, the value type is a
+   *  disjunction of a single address (which is the common case) and a list of addresses.
+   */
+  private val treeAddrs = new java.util.IdentityHashMap[Tree, TreeAddrs]
+
+  def registerTreeAddr(tree: Tree) =
+    treeAddrs.put(tree,
+      treeAddrs.get(tree) match {
+        case null => currentAddr
+        case x: Addr => x :: currentAddr :: Nil
+        case xs: List[_] => xs :+ currentAddr
+      })
+
+  def addrsOfTree(tree: Tree): List[Addr] = treeAddrs.get(tree) match {
+    case null => Nil
+    case addr: Addr => addr :: Nil
+    case addrs: List[Addr] => addrs
   }
 
   private def offset(i: Int): Addr = Addr(offsets(i))
@@ -150,10 +163,13 @@ class TreeBuffer extends TastyBuffer(50000) {
   }
 
   def adjustTreeAddrs(): Unit = {
-    val it = treeAddr.keySet.iterator
+    val it = treeAddrs.keySet.iterator
     while (it.hasNext) {
       val tree = it.next
-      treeAddr.put(tree, adjusted(treeAddr.get(tree).asInstanceOf[Addr]))
+      treeAddrs.get(tree) match {
+        case addr: Addr => treeAddrs.put(tree, adjusted(addr))
+        case addrs: List[Addr] => treeAddrs.put(tree, addrs.map(adjusted))
+      }
     }
   }