Set the positions of inlined trees wehn read form Tasty

This required a major change in the way positions are handled, as the
previous scheme did not allow to read the positions of arbitrary subtrees
selectively.

Fortunately, it's altogether a major simplification.

Also, this fixed a bug in the previous scheme, where positions were
generated before compactification, resulting in addresses being wrong.

Author: odersky

src/dotty/tools/dotc/FromTasty.scala

@@ -86,7 +86,7 @@ object FromTasty extends Driver {
               case info: ClassfileLoader =>
                 info.load(clsd) match {
                   case Some(unpickler: DottyUnpickler) =>
-                    val List(unpickled) = unpickler.body(readPositions = true)
+                    val List(unpickled) = unpickler.body(ctx.addMode(Mode.ReadPositions))
                     val unit1 = new CompilationUnit(new SourceFile(clsd.symbol.sourceFile, Seq()))
                     unit1.tpdTree = unpickled
                     unit1.unpicklers += (clsd.classSymbol -> unpickler.unpickler)

src/dotty/tools/dotc/core/Mode.scala

@@ -89,5 +89,8 @@ object Mode {
    */
   val AllowLambdaWildcardApply = newMode(15, "AllowHKApplyToWildcards")
 
+  /** Read original positions when unpickling from TASTY */
+  val ReadPositions = newMode(16, "ReadPositions")
+
   val PatternOrType = Pattern | Type
 }

src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala

@@ -8,23 +8,24 @@ import dotty.tools.dotc.ast.tpd
 import TastyUnpickler._, TastyBuffer._
 import util.Positions._
 import util.{SourceFile, NoSource}
-import PositionUnpickler._
 import Annotations.Annotation
+import core.Mode
 import classfile.ClassfileParser
 
 object DottyUnpickler {
 
   /** Exception thrown if classfile is corrupted */
   class BadSignature(msg: String) extends RuntimeException(msg)
 
-  class TreeSectionUnpickler extends SectionUnpickler[TreeUnpickler]("ASTs") {
+  class TreeSectionUnpickler(posUnpickler: Option[PositionUnpickler])
+  extends SectionUnpickler[TreeUnpickler]("ASTs") {
     def unpickle(reader: TastyReader, tastyName: TastyName.Table) =
-      new TreeUnpickler(reader, tastyName)
+      new TreeUnpickler(reader, tastyName, posUnpickler)
   }
 
-  class PositionsSectionUnpickler extends SectionUnpickler[(Position, AddrToPosition)]("Positions") {
+  class PositionsSectionUnpickler extends SectionUnpickler[PositionUnpickler]("Positions") {
     def unpickle(reader: TastyReader, tastyName: TastyName.Table) =
-      new PositionUnpickler(reader).unpickle()
+      new PositionUnpickler(reader)
   }
 }
 
@@ -36,21 +37,17 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded {
   import DottyUnpickler._
 
   val unpickler = new TastyUnpickler(bytes)
-  private val treeUnpickler = unpickler.unpickle(new TreeSectionUnpickler).get
+  private val posUnpicklerOpt = unpickler.unpickle(new PositionsSectionUnpickler)
+  private val treeUnpickler = unpickler.unpickle(new TreeSectionUnpickler(posUnpicklerOpt)).get
 
   /** Enter all toplevel classes and objects into their scopes
    *  @param roots          a set of SymDenotations that should be overwritten by unpickling
    */
   def enter(roots: Set[SymDenotation])(implicit ctx: Context): Unit =
     treeUnpickler.enterTopLevel(roots)
 
-  /** The unpickled trees, and the source file they come from
-   *  @param readPositions if true, trees get decorated with position information.
-   */
-  def body(readPositions: Boolean = false)(implicit ctx: Context): List[Tree] = {
-    if (readPositions)
-      for ((totalRange, positions) <- unpickler.unpickle(new PositionsSectionUnpickler))
-        treeUnpickler.usePositions(totalRange, positions)
+  /** The unpickled trees, and the source file they come from. */
+  def body(implicit ctx: Context): List[Tree] = {
     treeUnpickler.unpickle()
   }
 }

src/dotty/tools/dotc/core/tasty/PositionPickler.scala

@@ -3,7 +3,8 @@ package dotc
 package core
 package tasty
 
-import ast.tpd._
+import ast._
+import ast.Trees._
 import ast.Trees.WithLazyField
 import TastyFormat._
 import core._
@@ -12,64 +13,47 @@ import collection.mutable
 import TastyBuffer._
 import util.Positions._
 
-object PositionPickler {
-
-  trait DeferredPosition {
-    var parentPos: Position = NoPosition
-  }
-
-  def traverse(x: Any, parentPos: Position, op: (Tree, Position) => Unit)(implicit ctx: Context): Unit =
-    if (parentPos.exists)
-      x match {
-        case x: Tree @unchecked =>
-          op(x, parentPos)
-          x match {
-            case x: MemberDef @unchecked => traverse(x.symbol.annotations, x.pos, op)
-            case _ =>
-          }
-          traverse(x.productIterator, x.pos, op)
-        case x: DeferredPosition =>
-          x.parentPos = parentPos
-        case xs: TraversableOnce[_] =>
-          xs.foreach(traverse(_, parentPos, op))
-        case _ =>
-      }
-}
-import PositionPickler._
-
-class PositionPickler(pickler: TastyPickler, addrOfTree: Tree => Option[Addr]) {
+class PositionPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr]) {
   val buf = new TastyBuffer(5000)
   pickler.newSection("Positions", buf)
   import buf._
+  import ast.tpd._
+
+  def header(addrDelta: Int, hasStartDelta: Boolean, hasEndDelta: Boolean) = {
+    def toInt(b: Boolean) = if (b) 1 else 0
+    (addrDelta << 2) | (toInt(hasStartDelta) << 1) | toInt(hasEndDelta)
+  }
 
-  def picklePositions(roots: List[Tree], totalRange: Position)(implicit ctx: Context) = {
+  def picklePositions(roots: List[Tree])(implicit ctx: Context) = {
     var lastIndex = 0
-    def record(tree: Tree, parentPos: Position): Unit =
-      if (tree.pos.exists) {
-        def msg = s"failure to pickle $tree at ${tree.pos}, parent = $parentPos"
-        val endPos = tree.pos.end min parentPos.end
-          // end positions can be larger than their parents
-          // e.g. in the case of synthetic empty ranges, which are placed at the next token after
-          // the current construct.
-        val endDelta = endPos - parentPos.end
-        val startPos =
-          if (endDelta == 0) tree.pos.start max parentPos.start else tree.pos.start min endPos
-          // Since end positions are corrected above, start positions have to follow suit.
-        val startDelta = startPos - parentPos.start
-        if (startDelta != 0 || endDelta != 0)
-          for (addr <- addrOfTree(tree)) {
-            buf.writeInt(addr.index - lastIndex)
-            lastIndex = addr.index
-            if (startDelta != 0) buf.writeInt(startDelta)
-            if (endDelta != 0) {
-              assert(endDelta < 0, msg)
-              buf.writeInt(endDelta)
-            } else
-              assert(startDelta >= 0, msg)
+    var lastPos = Position(0, 0)
+    def pickleDeltas(index: Int, pos: Position) = {
+      val addrDelta = index - lastIndex
+      val startDelta = pos.start - lastPos.start
+      val endDelta = pos.end - lastPos.end
+      buf.writeInt(header(addrDelta, startDelta != 0, endDelta != 0))
+      if (startDelta != 0) buf.writeInt(startDelta)
+      if (endDelta != 0) buf.writeInt(endDelta)
+      lastIndex = index
+      lastPos = pos
+    }
+    def traverse(x: Any, parentPos: Position): Unit = x match {
+      case x: Tree @unchecked =>
+        if (x.pos.exists && x.pos.toSynthetic != parentPos.toSynthetic) {
+          addrOfTree(x) match {
+            case Some(addr) => pickleDeltas(addr.index, x.pos)
+            case _ =>
           }
-      }
-
-    buf.writeNat(totalRange.end)
-    traverse(roots, totalRange, record)
+        }
+        x match {
+          case x: MemberDef @unchecked => traverse(x.symbol.annotations, x.pos)
+          case _ =>
+        }
+        traverse(x.productIterator, x.pos)
+      case xs: TraversableOnce[_] =>
+        xs.foreach(traverse(_, parentPos))
+      case _ =>
+    }
+    traverse(roots, NoPosition)
   }
 }

src/dotty/tools/dotc/core/tasty/PositionUnpickler.scala

@@ -6,33 +6,31 @@ package tasty
 
 import util.Positions._
 import collection.mutable
-import TastyBuffer.Addr
-
-object PositionUnpickler {
-  type AddrToPosition = mutable.HashMap[Addr, Position]
-}
+import TastyBuffer.{Addr, NoAddr}
 
 /** Unpickler for tree positions */
 class PositionUnpickler(reader: TastyReader) {
-  import PositionUnpickler._
   import reader._
 
-  def unpickle(): (Position, AddrToPosition) = {
-    val positions = new mutable.HashMap[Addr, Position] // Dotty deviation: Can't use new AddrToPosition here. TODO: fix this!
-    val sourceLength = readNat()
-    def readDelta() = if (isAtEnd) 0 else readInt()
-    var curIndex: Addr = Addr(readDelta())
+  private[tasty] lazy val positions = {
+    val positions = new mutable.HashMap[Addr, Position]
+    var curIndex = 0
+    var curStart = 0
+    var curEnd = 0
     while (!isAtEnd) {
-      val delta1 = readDelta()
-      val delta2 = readDelta()
-      val (startDelta, endDelta, indexDelta) =
-        if (delta2 <= 0) (delta1, -delta2, readDelta())
-        else if (delta1 < 0) (0, -delta1, delta2)
-        else (delta1, 0, delta2)
-      positions(curIndex) = Position(startDelta, endDelta, startDelta)
-        // make non-synthetic position; will be made synthetic by normalization.
-      curIndex += indexDelta
+      val header = readInt()
+      val addrDelta = header >> 2
+      val hasStart = (header & 2) != 0
+      val hasEnd = (header & 1) != 0
+      curIndex += addrDelta
+      assert(curIndex >= 0)
+      if (hasStart) curStart += readInt()
+      if (hasEnd) curEnd += readInt()
+      positions(Addr(curIndex)) = Position(curStart, curEnd)
     }
-    (Position(0, sourceLength), positions)
+    positions
   }
+
+  def posAt(addr: Addr) = positions.getOrElse(addr, NoPosition)
 }
+

src/dotty/tools/dotc/core/tasty/TastyFormat.scala

@@ -185,21 +185,16 @@ Note: Tree tags are grouped into 5 categories that determine what follows, and t
   Category 4 (tags 112-127):  tag Nat AST
   Category 5 (tags 128-255):  tag Length <payload>
 
-Standard Section: "Positions" sourceLength_Nat Assoc*
-
-  Assoc         = addr_Delta offset_Delta offset_Delta?
-                                            // addr_Delta      :
-                                            //    Difference of address to last recorded node.
-                                            //    All but the first addr_Deltas are > 0, the first is >= 0.
-                                            // 2nd offset_Delta:
-                                            //    Difference of end offset of addressed node vs parent node. Always <= 0
-                                            // 1st offset Delta, if delta >= 0 or 2nd offset delta exists
-                                            //    Difference of start offset of addressed node vs parent node.
-                                            // 1st offset Delta, if delta < 0 and 2nd offset delta does not exist:
-                                            //    Difference of end offset of addressed node vs parent node.
-                                            // Offsets and addresses are difference encoded.
+Standard Section: "Positions" Assoc*
+
+  Assoc         = Header offset_Delta? offset_Delta?
+  Header        = addr_Delta +              // in one Nat: difference of address to last recorded node << 2 +
+                  hasStartDiff +            // one bit indicating whether there follows a start address delta << 1
+                  hasEndDiff                // one bit indicating whether there follows an end address delta
                                             // Nodes which have the same positions as their parents are omitted.
-  Delta         = Int                       // Difference between consecutive offsets / tree addresses,
+                                            // offset_Deltas give difference of start/end offset wrt to the
+                                            // same offset in the previously recorded node (or 0 for the first recorded node)
+  Delta         = Int                       // Difference between consecutive offsets,
 
 **************************************************************************************/
 

src/dotty/tools/dotc/core/tasty/TastyPickler.scala

@@ -31,7 +31,6 @@ class TastyPickler {
     sections += ((nameBuffer.nameIndex(name), buf))
 
   def assembleParts(): Array[Byte] = {
-    treePkl.compactify()
     def lengthWithLength(buf: TastyBuffer) = {
       buf.assemble()
       buf.length + natSize(buf.length)

src/dotty/tools/dotc/core/tasty/TastyPrinter.scala

@@ -113,8 +113,8 @@ class TastyPrinter(bytes: Array[Byte])(implicit ctx: Context) {
   class PositionSectionUnpickler extends SectionUnpickler[Unit]("Positions") {
     def unpickle(reader: TastyReader, tastyName: TastyName.Table): Unit = {
       print(s"${reader.endAddr.index - reader.currentAddr.index}")
-      val (totalRange, positions) = new PositionUnpickler(reader).unpickle()
-      println(s" position bytes in $totalRange:")
+      val positions = new PositionUnpickler(reader).positions
+      println(s" position bytes:")
       val sorted = positions.toSeq.sortBy(_._1.index)
       for ((addr, pos) <- sorted) println(s"${addr.index}: ${offsetToInt(pos.start)} .. ${pos.end}")
     }