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.

8 files changed, 103 insertions, 159 deletions

diff --git a/src/dotty/tools/dotc/core/Mode.scala b/src/dotty/tools/dotc/core/Mode.scala
index 3e9b7effe..7a9bb0572 100644
--- a/src/dotty/tools/dotc/core/Mode.scala
+++ b/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
 }
diff --git a/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala b/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala
index 0ad5d6966..2c93819d5 100644
--- a/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala
+++ b/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala
@@ -8,8 +8,8 @@ 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 {
@@ -17,14 +17,15 @@ 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,7 +37,8 @@ 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
@@ -44,13 +46,8 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded {
   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()
   }
 }
diff --git a/src/dotty/tools/dotc/core/tasty/PositionPickler.scala b/src/dotty/tools/dotc/core/tasty/PositionPickler.scala
index b0550b70a..63bb00a71 100644
--- a/src/dotty/tools/dotc/core/tasty/PositionPickler.scala
+++ b/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)
   }
 }
diff --git a/src/dotty/tools/dotc/core/tasty/PositionUnpickler.scala b/src/dotty/tools/dotc/core/tasty/PositionUnpickler.scala
index fa80a2769..c29aeba70 100644
--- a/src/dotty/tools/dotc/core/tasty/PositionUnpickler.scala
+++ b/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)
 }
+
diff --git a/src/dotty/tools/dotc/core/tasty/TastyFormat.scala b/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
index e9de68e7f..8e8d58b47 100644
--- a/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
+++ b/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,
 
 **************************************************************************************/
 
diff --git a/src/dotty/tools/dotc/core/tasty/TastyPickler.scala b/src/dotty/tools/dotc/core/tasty/TastyPickler.scala
index 83e6020d5..98b0dc7c6 100644
--- a/src/dotty/tools/dotc/core/tasty/TastyPickler.scala
+++ b/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)
diff --git a/src/dotty/tools/dotc/core/tasty/TastyPrinter.scala b/src/dotty/tools/dotc/core/tasty/TastyPrinter.scala
index 915ae3f21..7fcd7c29e 100644
--- a/src/dotty/tools/dotc/core/tasty/TastyPrinter.scala
+++ b/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}")
     }
diff --git a/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala b/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
index 11f2eddac..60cad0445 100644
--- a/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
+++ b/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
@@ -10,38 +10,23 @@ import ast.{tpd, Trees, untpd}
 import typer.Inliner
 import Trees._
 import Decorators._
-import TastyUnpickler._, TastyBuffer._, PositionPickler._
+import TastyUnpickler._, TastyBuffer._
 import scala.annotation.{tailrec, switch}
 import scala.collection.mutable.ListBuffer
 import scala.collection.{ mutable, immutable }
 import config.Printers.pickling
 
 /** Unpickler for typed trees
- *  @param reader         the reader from which to unpickle
- *  @param tastyName      the nametable
+ *  @param reader          the reader from which to unpickle
+ *  @param tastyName       the nametable
+ *  @param posUNpicklerOpt the unpickler for positions, if it exists
  */
-class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table) {
+class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpicklerOpt: Option[PositionUnpickler]) {
   import TastyFormat._
   import TastyName._
   import TreeUnpickler._
   import tpd._
 
-  private var readPositions = false
-  private var totalRange = NoPosition
-  private var positions: collection.Map[Addr, Position] = _
-
-  /** Make a subsequent call to `unpickle` return trees with positions
-   *  @param totalRange     the range position enclosing all returned trees,
-   *                        or NoPosition if positions should not be unpickled
-   *  @param positions      a map from tree addresses to their positions relative
-   *                        to positions of parent nodes.
-   */
-  def usePositions(totalRange: Position, positions: collection.Map[Addr, Position]): Unit = {
-    readPositions = true
-    this.totalRange = totalRange
-    this.positions = positions
-  }
-
   /** A map from addresses of definition entries to the symbols they define */
   private val symAtAddr  = new mutable.HashMap[Addr, Symbol]
 
@@ -86,10 +71,7 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table) {
   /** The unpickled trees */
   def unpickle()(implicit ctx: Context): List[Tree] = {
     assert(roots != null, "unpickle without previous enterTopLevel")
-    val stats = new TreeReader(reader)
-      .readTopLevel()(ctx.addMode(Mode.AllowDependentFunctions))
-    normalizePos(stats, totalRange)
-    stats
+    new TreeReader(reader).readTopLevel()(ctx.addMode(Mode.AllowDependentFunctions))
   }
 
   def toTermName(tname: TastyName): TermName = tname match {
@@ -508,7 +490,8 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table) {
       }
       else (annots.find(_.symbol == defn.InlineAnnot)) match {
         case Some(inlineAnnot) =>
-          Inliner.attachBody(inlineAnnot, forkAt(rhsStart).readTerm()(localContext(sym)))
+          Inliner.attachBody(inlineAnnot,
+            forkAt(rhsStart).readTerm()(localContext(sym).addMode(Mode.ReadPositions)))
         case none =>
       }
       goto(start)
@@ -1017,44 +1000,29 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table) {
       new LazyReader(localReader, op)
     }
 
-// ------ Hooks for positions ------------------------------------------------
+// ------ Setting positions ------------------------------------------------
 
-    /** Record address from which tree was created as a temporary position in the tree.
-     *  The temporary position contains deltas relative to the position of the (as yet unknown)
-     *  parent node. It is marked as a non-synthetic source position.
-     */
-    def setPos[T <: Tree](addr: Addr, tree: T): T = {
-      if (readPositions)
-        tree.setPosUnchecked(positions.getOrElse(addr, Position(0, 0, 0)))
-      tree
-    }
+    /** Set position of `tree` at given `addr`. */
+    def setPos[T <: Tree](addr: Addr, tree: T)(implicit ctx: Context): tree.type =
+      if (ctx.mode.is(Mode.ReadPositions)) {
+        posUnpicklerOpt match {
+          case Some(posUnpickler) => tree.withPos(posUnpickler.posAt(addr))
+          case _  => tree
+        }
+      }
+      else tree
   }
 
-  private def setNormalized(tree: Tree, parentPos: Position): Unit =
-    tree.setPosUnchecked(
-      if (tree.pos.exists)
-        Position(parentPos.start + offsetToInt(tree.pos.start), parentPos.end - tree.pos.end)
-      else
-        parentPos)
-
-  def normalizePos(x: Any, parentPos: Position)(implicit ctx: Context): Unit =
-    traverse(x, parentPos, setNormalized)
-
-  class LazyReader[T <: AnyRef](reader: TreeReader, op: TreeReader => Context => T) extends Trees.Lazy[T] with DeferredPosition {
+  class LazyReader[T <: AnyRef](reader: TreeReader, op: TreeReader => Context => T) extends Trees.Lazy[T] {
     def complete(implicit ctx: Context): T = {
       pickling.println(i"starting to read at ${reader.reader.currentAddr}")
-      val res = op(reader)(ctx.addMode(Mode.AllowDependentFunctions).withPhaseNoLater(ctx.picklerPhase))
-      normalizePos(res, parentPos)
-      res
+      op(reader)(ctx.addMode(Mode.AllowDependentFunctions).withPhaseNoLater(ctx.picklerPhase))
     }
   }
 
-  class LazyAnnotationReader(sym: Symbol, reader: TreeReader)
-      extends LazyAnnotation(sym) with DeferredPosition {
+  class LazyAnnotationReader(sym: Symbol, reader: TreeReader) extends LazyAnnotation(sym) {
     def complete(implicit ctx: Context) = {
-      val res = reader.readTerm()(ctx.withPhaseNoLater(ctx.picklerPhase))
-      normalizePos(res, parentPos)
-      res
+      reader.readTerm()(ctx.withPhaseNoLater(ctx.picklerPhase))
     }
   }