Fixes for TypePositions scheme

odersky · odersky · commit 18fddc16cbd3 · 2016-10-27T10:22:56.000+02:00
Notable, deal with the fact that type trees can be shared, so
the map from trees to positioned types must be a multi-map.
diff --git a/src/dotty/tools/dotc/core/tasty/PositionPickler.scala b/src/dotty/tools/dotc/core/tasty/PositionPickler.scala
@@ -11,30 +11,18 @@ import core._
 import Contexts._, Symbols._, Types._, Names._, Constants._, Decorators._, Annotations._
 import collection.mutable
 import TastyBuffer._
+import TreeBuffer._
+import ast.TypePositions.PosAnnot
 import util.Positions._
 
-class PositionPickler(pickler: TastyPickler, addrsOfTree: tpd.Tree => List[Addr]) {
+class PositionPickler(pickler: TastyPickler, addrsInfo: AddrsInfo) {
   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
-  }
+  
+  private val addrOfIterator = addrsInfo.addrs.iterator
+  private val positionedOfIterator = addrsInfo.positioned.iterator
 
   def header(addrDelta: Int, hasStartDelta: Boolean, hasEndDelta: Boolean, hasPoint: Boolean) = {
     def toInt(b: Boolean) = if (b) 1 else 0
@@ -65,24 +53,34 @@ class PositionPickler(pickler: TastyPickler, addrsOfTree: tpd.Tree => List[Addr]
       case _ => false
     }
 
+    def picklePos(ad: Addressed, pos: Position) = 
+      if (addrOfIterator.hasNext(ad)) {
+        val addr = addrOfIterator.next(ad)
+        //println(i"pickling $ad with $pos at $addr")
+        pickleDeltas(addr.index, pos)
+      }
+
     def traverse(x: Any): Unit = x match {
       case x: Tree @unchecked =>
-        val pos = if (x.isInstanceOf[MemberDef]) x.pos else x.pos.toSynthetic
-        if (pos.exists && (pos != x.initialPos.toSynthetic || alwaysNeedsPos(x))) {
-          nextTreeAddr(x) match {
-            case Some(addr) =>
-              //println(i"pickling $x with $pos at $addr")
-              pickleDeltas(addr.index, pos)
+        if (positionedOfIterator.hasNext(x)) {
+          val posType = positionedOfIterator.next(x)
+          //println(s"pickle type pos $posType")
+          posType.foreachPart {
+            case t @ AnnotatedType(_, PosAnnot(pos)) => picklePos(t, pos.toSynthetic)
             case _ =>
-              //println(i"no address for $x")
-          }
+          }          
         }
-        //else if (x.pos.exists) println(i"skipping $x")
-        x match {
-          case x: MemberDef @unchecked => traverse(x.symbol.annotations.map(_.tree))
-          case _ =>
+        else {
+          val pos = if (x.isInstanceOf[MemberDef]) x.pos else x.pos.toSynthetic
+          if (pos.exists && (pos != x.initialPos.toSynthetic || alwaysNeedsPos(x)))
+            picklePos(x, pos)
+          //else if (x.pos.exists) println(i"skipping $x")
+          x match {
+            case x: MemberDef @unchecked => traverse(x.symbol.annotations.map(_.tree))
+            case _ =>
+          }
+          traverse(x.productIterator)
         }
-        traverse(x.productIterator)
       case xs: TraversableOnce[_] =>
         xs.foreach(traverse)
       case _ =>
diff --git a/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala b/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala
@@ -10,41 +10,102 @@ import ast.untpd.Tree
 import ast.TypePositions.PosAnnot
 import Types.Type
 
-class TreeBuffer extends TastyBuffer(50000) {
+object TreeBuffer {
+
+  /** A set of iterators indexed by a key */
+  abstract class MultiIterator[Key, Value] {
+    def hasNext(key: Key): Boolean
+    def next(key: Key): Value 
+  }
+
+  /** A map from `Key`s to lists of `Value`s. Optimized for the case
+   *  where only a single value is associated with a key.
+   */
+  class MultiMap[Key <: AnyRef, Value] {
+    private type Values = Any // really: Value | List[Values]
+    private val m = new java.util.IdentityHashMap[Key, Values]
+
+    /** Register `key -> value` binding as last value for `key`. */
+    def register(key: Key, value: Value) =
+      m.put(key,
+        m.get(key) match {
+          case null => value
+          case vs: List[_] => vs :+ value
+          case v => v :: value :: Nil
+        })
+
+    /** The list of values associated with to `key` */
+    def apply(key: Key): List[Value] = m.get(key) match {
+       case null => Nil
+       case vs: List[Value] => vs
+       case v: Value => v :: Nil
+    }
+
+    /** Transform all values stored in the map using `op` */
+    def updateValues(op: Value => Value): Unit = {
+      val it = m.keySet.iterator
+      while (it.hasNext) {
+        val x = it.next()
+        m.put(x, 
+          m.get(x) match {
+            case vs: List[Value] => vs.map(op)
+            case v: Value => op(v)
+          })
+      }
+    }
+
+    /** A new multi-iterator over this map */
+    def iterator = new MultiIterator[Key, Value] {
+      private val remaining = new java.util.IdentityHashMap[Key, Iterator[Value]]
+      def hasNext(key: Key) = remaining.get(key) match {
+        case it: Iterator[_] => it.hasNext
+        case _ => m.get(key) != null
+      }
+      def next(key: Key) = remaining.get(key) match {
+        case it: Iterator[_] => it.next
+        case _ => 
+          m.get(key) match {
+            case null => 
+              Iterator.empty.next
+            case (v: Value) :: (vs: List[Value]) => 
+              remaining.put(key, if (vs.isEmpty) Iterator.empty else vs.iterator)
+              v
+            case v: Value =>
+              remaining.put(key, Iterator.empty)
+              v
+          }
+      }
+    }
+  }
+
+  /** The type of things that have addresses of type `Addr`.
+   *  This is really `Tree | Type`, but that can't be expressed in current Scala
+   */
+  type Addressed = AnyRef
+
+  /** Address info associated with a TreeBuffer, consisting of two multi-maps,
+   *  one from trees and types to addresses, and
+   *  one from trees to positioned types.
+   */
+  trait AddrsInfo {
+    def addrs: MultiMap[Addressed, Addr]
+    def positioned: MultiMap[Tree, Type]
+  }
+}
+
+class TreeBuffer extends TastyBuffer(50000) with TreeBuffer.AddrsInfo {
+  import TreeBuffer._
 
   private final val ItemsOverOffsets = 2
   private val initialOffsetSize = bytes.length / (AddrWidth * ItemsOverOffsets)
   private var offsets = new Array[Int](initialOffsetSize)
   private var isRelative = new Array[Boolean](initialOffsetSize)
   private var delta: Array[Int] = _
   private var numOffsets = 0
-
-  private type TreeAddrs = Any // really: Addr | List[Addr]
-  
-  private type Addressed = Any // really: Tree | Positioned  
-  
-  /** 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[Addressed, 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
-  }
-
-  val posTypeOfTree = new java.util.IdentityHashMap[Tree, Type]
   
+  val addrs = new MultiMap[Addressed, Addr]
+  val positioned = new MultiMap[Tree, Type]
+    
   private def offset(i: Int): Addr = Addr(offsets(i))
 
   private def keepOffset(relative: Boolean): Unit = {
@@ -168,16 +229,7 @@ class TreeBuffer extends TastyBuffer(50000) {
     wasted
   }
 
-  def adjustTreeAddrs(): Unit = {
-    val it = treeAddrs.keySet.iterator
-    while (it.hasNext) {
-      val tree = it.next
-      treeAddrs.get(tree) match {
-        case addr: Addr => treeAddrs.put(tree, adjusted(addr))
-        case addrs: List[Addr] => treeAddrs.put(tree, addrs.map(adjusted))
-      }
-    }
-  }
+  def adjustTreeAddrs(): Unit = addrs.updateValues(adjusted)
 
   /** Final assembly, involving the following steps:
    *   - compute deltas
diff --git a/src/dotty/tools/dotc/core/tasty/TreePickler.scala b/src/dotty/tools/dotc/core/tasty/TreePickler.scala
@@ -143,6 +143,7 @@ class TreePickler(pickler: TastyPickler) {
 
   def pickleType(tpe0: Type, richTypes: Boolean = false)(implicit ctx: Context): Unit = tpe0 match {
     case AnnotatedType(tpe1, PosAnnot(pos)) =>
+      addrs.register(tpe0, currentAddr)
       pickleType(tpe1, richTypes)
     case _ =>
       try {
@@ -164,7 +165,7 @@ class TreePickler(pickler: TastyPickler) {
   }
 
   def pickleTypeWithPos(tpe: Type, tree: Tree)(implicit ctx: Context): Unit = {
-    registerTreeAddr(tree)
+    addrs.register(tree, currentAddr)
     pickleType(tpe)
   }
 
@@ -315,34 +316,33 @@ class TreePickler(pickler: TastyPickler) {
     pickled
   }
   
-  def pickleTpt(tpt: Tree)(implicit ctx: Context): Unit = tpt match {
-    case _: TypeTree => pickleTypeWithPos(tpt.tpe, tpt)
-    case _ => 
-      val posType = TypePositions.toType(tpt)
-      
-      def checkTypePositions() = {
-        val tree = TypePositions.toTree(posType)
-        val origRefs = TypePositions.refPositions(tpt)
-        val newRefs = TypePositions.refPositions(tree)
-        assert(TypePositions.refPositions(tpt).toSet subsetOf TypePositions.refPositions(tree).toSet, 
-            i"""$tpt 
-               |--> 
-               |$posType 
-               |--> 
-               |$tree
-               |/ 
-               |${posType.toString}
-               |/
-               |${tree.toString}""")
+  /** Pickle type tree. This will pickle its type but will try to keep position
+   *  info with it.
+   */
+  def pickleTpt(tpt: Tree)(implicit ctx: Context): Unit = 
+    picklePosTypeOfTree(TypePositions.toType(tpt), tpt)
+
+  /** Pickle positioned type associated with type tree `tpt`. */
+  private def picklePosTypeOfTree(posType: Type, tpt: Tree)(implicit ctx: Context): Unit = {
+    def checkTypePositions() = {
+      val tree = TypePositions.toTree(posType)
+      val origRefs = TypePositions.refPositions(tpt)
+      val newRefs = TypePositions.refPositions(tree)
+      assert(TypePositions.refPositions(tpt).toSet subsetOf TypePositions.refPositions(tree).toSet,
+        i"""$tpt 
+           |--> 
+           |$posType 
+           |--> 
+           |$tree
+           |/ 
+           |${posType.toString}
+           |/
+           |${tree.toString}""")
 
-      }
-      if (false) checkTypePositions
-      
-      posType match {
-        case pd @ AnnotatedType(_, _: PosAnnot) => posTypeOfTree.put(tpt, pd)
-        case _ => 
-      }
-      pickleTypeWithPos(posType, tpt)
+    }
+    if (false) checkTypePositions
+    positioned.register(tpt, posType)
+    pickleType(posType)
   }
 
   def pickleTreeUnlessEmpty(tree: Tree)(implicit ctx: Context): Unit =
@@ -365,7 +365,7 @@ class TreePickler(pickler: TastyPickler) {
   }
 
   def pickleParam(tree: Tree)(implicit ctx: Context): Unit = {
-    registerTreeAddr(tree)
+    addrs.register(tree, currentAddr)
     tree match {
       case tree: ValDef => pickleDef(PARAM, tree.symbol, tree.tpt)
       case tree: DefDef => pickleDef(PARAM, tree.symbol, tree.tpt, tree.rhs)
@@ -384,7 +384,7 @@ class TreePickler(pickler: TastyPickler) {
   }
 
   def pickleTree(tree: Tree)(implicit ctx: Context): Unit = try {
-    registerTreeAddr(tree)
+    addrs.register(tree, currentAddr)
     tree match {
       case tree if tree.isType =>
         pickleTpt(tree)
@@ -528,13 +528,12 @@ class TreePickler(pickler: TastyPickler) {
           if ((selfInfo ne NoType) || !tree.self.isEmpty) {
             writeByte(SELFDEF)
             pickleName(tree.self.name)
-            if (!tree.self.isEmpty) registerTreeAddr(tree.self.tpt)
-            pickleType {
-              cinfo.selfInfo match {
-                case sym: Symbol => sym.info
-                case tp: Type => tp
-              }
+            val selfTp = cinfo.selfInfo match {
+              case sym: Symbol => sym.info
+              case tp: Type => tp
             }
+            if (!tree.self.tpt.isEmpty) pickleTpt(tree.self.tpt)
+            else pickleType(selfTp)
           }
           pickleStats(tree.constr :: rest)
         }
@@ -561,7 +560,7 @@ class TreePickler(pickler: TastyPickler) {
   }
 
   def pickleSelector(tag: Int, id: untpd.Ident)(implicit ctx: Context): Unit = {
-    registerTreeAddr(id)
+    addrs.register(id, currentAddr)
     writeByte(tag)
     pickleName(id.name)
   }
diff --git a/src/dotty/tools/dotc/printing/RefinedPrinter.scala b/src/dotty/tools/dotc/printing/RefinedPrinter.scala
@@ -15,6 +15,7 @@ import Decorators._
 import config.Config
 import scala.annotation.switch
 import language.implicitConversions
+import dotty.tools.dotc.ast.TypePositions
 
 class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
 
@@ -529,7 +530,12 @@ class RefinedPrinter(_ctx: Context) extends PlainPrinter(_ctx) {
       txt = toText(tree.typeOpt)
     if (ctx.settings.Yprintpos.value && !tree.isInstanceOf[WithoutTypeOrPos[_]]) {
       val pos = 
-        if (homogenizedView && !tree.isInstanceOf[MemberDef]) tree.pos.toSynthetic 
+        if (homogenizedView) 
+          tree match {
+            case _: MemberDef => tree.pos // point is important for MemberDefs since it points to defined id
+            case SingletonTypeTree(ref) if ref.pos.exists => ref.pos.toSynthetic // we only serialize pos of underlying ref
+            case _ => tree.pos.toSynthetic // deserialized position is synthetic
+          }
         else tree.pos
       val clsStr = "" // DEBUG: if (tree.isType) tree.getClass.toString else ""
       txt = (txt ~ "@" ~ pos.toString ~ clsStr).close
diff --git a/src/dotty/tools/dotc/transform/Pickler.scala b/src/dotty/tools/dotc/transform/Pickler.scala
@@ -46,10 +46,10 @@ class Pickler extends Phase {
       val treePkl = pickler.treePkl
       treePkl.pickle(tree :: Nil)
       treePkl.compactify()
-      pickler.addrsOfTree = treePkl.buf.addrsOfTree
+      pickler.addrsOfTree = treePkl.buf.addrs.apply
       pickler.addrOfSym = treePkl.addrOfSym
       if (tree.pos.exists)
-        new PositionPickler(pickler, treePkl.buf.addrsOfTree).picklePositions(tree :: Nil)
+        new PositionPickler(pickler, treePkl.buf).picklePositions(tree :: Nil)
 
       def rawBytes = // not needed right now, but useful to print raw format.
         pickler.assembleParts().iterator.grouped(10).toList.zipWithIndex.map {
