New scheme for recording positions

Single traverser, also handles lazy trees.

10 files changed, 179 insertions, 153 deletions

diff --git a/src/dotty/tools/dotc/core/pickling/DottyUnpickler.scala b/src/dotty/tools/dotc/core/pickling/DottyUnpickler.scala
index 5db07ae2d..9dbbec01e 100644
--- a/src/dotty/tools/dotc/core/pickling/DottyUnpickler.scala
+++ b/src/dotty/tools/dotc/core/pickling/DottyUnpickler.scala
@@ -32,13 +32,11 @@ class DottyUnpickler(bytes: Array[Byte], readPositions: Boolean = false)(implici
 
   class TreeSectionUnpickler()(implicit ctx: Context) extends SectionUnpickler[List[Tree]]("ASTs") {
     def unpickle(reader: TastyReader, tastyName: TastyName.Table): List[Tree] =
-      new TreesUnpickler(reader, tastyName, readPositions).unpickle()
+      new TreeUnpickler(reader, tastyName, readPositions).unpickle()
   }
   
   class PositionsSectionUnpickler(trees: List[Tree])(implicit ctx: Context) extends SectionUnpickler[Unit]("Positions") {
-    def unpickle(reader: TastyReader, tastyName: TastyName.Table): Unit = {
-      new OffsetUnpickler(reader, Edge.left).traverse(trees)
-      new OffsetUnpickler(reader, Edge.right).traverse(trees)
-    }
+    def unpickle(reader: TastyReader, tastyName: TastyName.Table): Unit =
+      new PositionReader(reader).unpickle(trees)
   }
 }
diff --git a/src/dotty/tools/dotc/core/pickling/Edge.scala b/src/dotty/tools/dotc/core/pickling/Edge.scala
deleted file mode 100644
index 59cd85ab2..000000000
--- a/src/dotty/tools/dotc/core/pickling/Edge.scala
+++ /dev/null
@@ -1,54 +0,0 @@
-package dotty.tools.dotc
-package core
-package pickling
-
-import util.Positions._
-import ast.tpd.{Tree, MemberDef}
-import core.Contexts._
-
-abstract class Edge {
-  
-  def offset(pos: Position): Int
-  def seq(op1: () => Unit, op2: () => Unit): Unit
-  def updateOffset(pos: Position, off: Int): Position
-  
-  private val noOp = () => ()
-  
-  def traverseAll(roots: List[Tree])(op: Tree => Unit)(implicit ctx: Context) = {
-    
-    def elemsTraversal(xs: TraversableOnce[Any]): () => Unit = 
-      (noOp /: xs) ((op, x) => () => seq(op, elemTraversal(x)))
-
-    def elemTraversal(x: Any): () => Unit = () => x match {        
-      case x: Tree @ unchecked =>
-        op(x)
-        
-        val annotTraversal = x match {
-          case x: MemberDef => elemsTraversal(x.symbol.annotations)
-          case _ => noOp
-        }
-        val childrenTraversal = elemsTraversal(x.productIterator)
-        seq(annotTraversal, childrenTraversal)
-      case xs: List[_] =>
-        elemsTraversal(xs)()
-      case _ =>
-        ()
-    }
- 
-    elemsTraversal(roots)()
-  }
-}
-
-object Edge {
-
-  object left extends Edge {
-    def offset(pos: Position): Int = pos.start
-    def updateOffset(pos: Position, off: Int) = Position(off, pos.end, 0)
-    def seq(op1: () => Unit, op2: () => Unit) = { op1(); op2() }
-  }
-  object right extends Edge {
-    def offset(pos: Position): Int = pos.end
-    def updateOffset(pos: Position, off: Int) = Position(pos.start, off, 0)
-    def seq(op1: () => Unit, op2: () => Unit) = { op2(); op1() }
-  }
-}
\ No newline at end of file
diff --git a/src/dotty/tools/dotc/core/pickling/OffsetUnpickler.scala b/src/dotty/tools/dotc/core/pickling/OffsetUnpickler.scala
deleted file mode 100644
index 71052bb1b..000000000
--- a/src/dotty/tools/dotc/core/pickling/OffsetUnpickler.scala
+++ /dev/null
@@ -1,37 +0,0 @@
-package dotty.tools
-package dotc
-package core
-package pickling
-
-import java.io.IOException
-
-import Contexts._
-import ast.tpd
-
-/** Unpickler for start or end offset of tree positions */
-class OffsetUnpickler(reader: TastyReader, edge: Edge) {
-  import reader._
-  private val end = readEnd()
-  private var lastOffset = 0
-  private var nextOffset = 0
-  private var nextAddr = 0
-
-  private def next() = {
-    lastOffset = nextOffset
-    if (currentAddr != end) {
-      nextOffset += readInt()
-      nextAddr += readInt()
-    }
-  }  
-      
-  def traverse(trees: List[tpd.Tree])(implicit ctx: Context) = {
-    next()
-    edge.traverseAll(trees) { tree =>
-      if (tree.pos.exists) { // TODO assign positions of tree nodes with NoPosition afterwards based on the context.
-        if (edge.offset(tree.pos) == nextAddr) next()
-        tree.setPosUnchecked(edge.updateOffset(tree.pos, lastOffset))
-      }
-    }
-    // assert(currentAddr == end)
-  }
-}
\ No newline at end of file
diff --git a/src/dotty/tools/dotc/core/pickling/PickleFormat.scala b/src/dotty/tools/dotc/core/pickling/PickleFormat.scala
index d2084dfd9..9e864fc43 100644
--- a/src/dotty/tools/dotc/core/pickling/PickleFormat.scala
+++ b/src/dotty/tools/dotc/core/pickling/PickleFormat.scala
@@ -181,11 +181,14 @@ Note: Tree tags are grouped into 4 categories that determine what follows, and t
 	Category 3 (tags 100-127):	tag Nat AST
 	Category 4 (tags 128-255):  tag Length <payload>
 
-Standard Section: "Positions" startPos_Index endPos_Index
-
-  Index         = Length Assoc*
-  Assoc         = offset_Delta addr_Delta   // largest tree starting/ending at offset
-  Delta         = Int                 	    // difference between consecutive offsets / tree addresses,
+Standard Section: "Positions" Assoc*
+
+  Assoc         = offset_Delta addr_Delta   // offsets and addresses determined by a tree traversal:
+                                            // For each node with positions <start..end>
+                                            //   start <encoding of children> end
+                                            // Offsets and addresses are difference encoded.
+                                            // Entries with same offset as previous entry are omitted.
+  Delta         = Int                 	    // Difference between consecutive offsets / tree addresses,
                                             // First offset/address is always assumed to be 0
 
 **************************************************************************************/
diff --git a/src/dotty/tools/dotc/core/pickling/PositionBuffer.scala b/src/dotty/tools/dotc/core/pickling/PositionBuffer.scala
index f564e364d..88e4733d1 100644
--- a/src/dotty/tools/dotc/core/pickling/PositionBuffer.scala
+++ b/src/dotty/tools/dotc/core/pickling/PositionBuffer.scala
@@ -7,27 +7,15 @@ import TastyBuffer.Addr
 
 class PositionBuffer extends TastyBuffer(100000) { thisBuffer =>
   
-  class PositionRecorder(val edge: Edge) extends TastyBuffer(thisBuffer.bytes.size / 2) {
-    private var lastOffset: Int = -1
-    private var lastIndex: Int = 0
-    def record(addr: Addr, offset: Int): Unit = 
-      if (offset != lastOffset) {
-        if (lastOffset < 0) lastOffset = 0
-        writeInt(offset - lastOffset)
-        writeInt(addr.index - lastIndex)
-        lastOffset = offset
-        lastIndex = addr.index
-      }
-  }
+  private var lastOffset: Int = -1
+  private var lastIndex: Int = 0
   
-  val startPos = new PositionRecorder(Edge.left)
-  val endPos = new PositionRecorder(Edge.right)
-
-  /** Final assembly: copy startPos and endPos into own bytes */
-  override def assemble(): Unit = {
-    writeNat(startPos.length)
-    writeBytes(startPos.bytes, startPos.length)
-    writeNat(endPos.length)
-    writeBytes(endPos.bytes, endPos.length)
-  }
+  def record(addr: Addr, offset: Int, recordAlways: Boolean): Unit =
+    if (offset != lastOffset || recordAlways) {
+      if (lastOffset < 0) lastOffset = 0
+      writeInt(offset - lastOffset)
+      writeInt(addr.index - lastIndex)
+      lastOffset = offset
+      lastIndex = addr.index
+    }
 }
\ No newline at end of file
diff --git a/src/dotty/tools/dotc/core/pickling/PositionPickler.scala b/src/dotty/tools/dotc/core/pickling/PositionPickler.scala
index 58f00926d..7cb7de6e3 100644
--- a/src/dotty/tools/dotc/core/pickling/PositionPickler.scala
+++ b/src/dotty/tools/dotc/core/pickling/PositionPickler.scala
@@ -4,25 +4,27 @@ package core
 package pickling
 
 import ast.tpd._
+import ast.Trees.WithLazyField
 import PickleFormat._
 import core._
 import Contexts._, Symbols._, Types._, Names._, Constants._, Decorators._, Annotations._
 import collection.mutable
 import TastyBuffer._
+import Traversals._
 
 class PositionPickler(pickler: TastyPickler, addrOfTree: Tree => Option[Addr]) {
   val buf = new PositionBuffer
   pickler.newSection("Positions", buf)
   import buf._
     
-  def picklePositions(roots: List[Tree])(implicit ctx: Context) = {
-    def traverseWith(recorder: PositionRecorder) = 
-      recorder.edge.traverseAll(roots) { tree =>
-        if (tree.pos.exists)
-          for (addr <- addrOfTree(tree))
-            recorder.record(addr, recorder.edge.offset(tree.pos))
-      }
-    traverseWith(startPos)
-    traverseWith(endPos)
-  }
+  private def record(tree: Tree, start: Boolean) =
+    if (tree.pos.exists)
+      for (addr <- addrOfTree(tree))
+        buf.record(
+          addr, 
+          offset = if (start) tree.pos.start else tree.pos.end, 
+          recordAlways = !start && tree.isInstanceOf[WithLazyField[_]])
+    
+  def picklePositions(roots: List[Tree])(implicit ctx: Context) =
+    traverse(roots, record(_, true), record(_, false))
 }
\ No newline at end of file
diff --git a/src/dotty/tools/dotc/core/pickling/PositionReader.scala b/src/dotty/tools/dotc/core/pickling/PositionReader.scala
new file mode 100644
index 000000000..5acb3c6d8
--- /dev/null
+++ b/src/dotty/tools/dotc/core/pickling/PositionReader.scala
@@ -0,0 +1,67 @@
+package dotty.tools
+package dotc
+package core
+package pickling
+
+import java.io.IOException
+
+import Contexts._
+import ast.tpd
+import ast.Trees.{Lazy,WithLazyField}
+import TastyBuffer._
+import Traversals._
+import util.Positions._
+
+/** Unpickler for start or end offset of tree positions */
+class PositionReader(reader: TastyReader, startOffset: Int, initNextOffset: Int, initNextAddr: Int) {
+  import reader._
+  import tpd._
+      
+  def this(reader: TastyReader) = {
+    this(reader, 0, 0, 0)
+    next()
+  }
+  
+  var lastOffset: Int = startOffset
+  var nextOffset = initNextOffset
+  var nextAddr = initNextAddr
+  
+  private def fork = 
+    new PositionReader(subReader(currentAddr, endAddr), lastOffset, nextOffset, nextAddr)
+  
+  private def next() = {
+    lastOffset = nextOffset
+    if (!isAtEnd) {
+      nextOffset += readInt()
+      nextAddr += readInt()
+    }
+  }  
+
+  def readStart(tree: Tree) = 
+    if (tree.pos.exists) { // TODO assign positions of tree nodes with NoPosition afterwards based on the context.
+      val addr = tree.pos.end
+      if (addr == nextAddr) next()
+      tree.setPosUnchecked(tree.pos.withStart(lastOffset))
+    }
+  
+  def readEnd(tree: Tree) = {
+    if (tree.pos.exists) {
+      val addr = tree.pos.end
+      tree match {
+        case tree: WithLazyField[_] =>
+          tree.unforced match {
+            case rdr: TreeUnpickler#LazyReader[_] =>
+              rdr.posReader = Some(fork)
+              while (addr != nextAddr) next()
+            case _ =>
+          }
+        case _ =>
+      }
+      if (addr == nextAddr) next()
+      tree.setPosUnchecked(tree.pos.withEnd(lastOffset))
+    }
+  }
+  
+  def unpickle(x: Any) =
+    traverse(x, readStart, readEnd)
+}
\ No newline at end of file
diff --git a/src/dotty/tools/dotc/core/pickling/TastyPrinter.scala b/src/dotty/tools/dotc/core/pickling/TastyPrinter.scala
index 94d57e081..751642a51 100644
--- a/src/dotty/tools/dotc/core/pickling/TastyPrinter.scala
+++ b/src/dotty/tools/dotc/core/pickling/TastyPrinter.scala
@@ -102,26 +102,17 @@ class TastyPrinter(bytes: Array[Byte])(implicit ctx: Context) {
       }
     }
   }
-  
+
   class PositionUnpickler extends SectionUnpickler[Unit]("Positions") {
     def unpickle(reader: TastyReader, tastyName: TastyName.Table): Unit = {
       import reader._
-
-      def unpickleOffsets(edge: Edge): Unit = {
-        var lastOffset = 0
-        var lastAddr = 0
-        val length = readNat()
-        println(s"$length offset bytes")
-        val end = currentAddr + length
-        until(end) {
-          lastOffset += readInt()
-          lastAddr += readInt()
-          println(s"$lastOffset: $lastAddr")
-        }
+      var lastOffset = 0
+      var lastAddr = 0
+      while (!isAtEnd) {
+        lastOffset += readInt()
+        lastAddr += readInt()
+        println(s"$lastOffset: $lastAddr")
       }
-      
-      unpickleOffsets(Edge.left)
-      unpickleOffsets(Edge.right)
     }
   }
 }
\ No newline at end of file
diff --git a/src/dotty/tools/dotc/core/pickling/Traversals.scala b/src/dotty/tools/dotc/core/pickling/Traversals.scala
new file mode 100644
index 000000000..85dc53cd0
--- /dev/null
+++ b/src/dotty/tools/dotc/core/pickling/Traversals.scala
@@ -0,0 +1,54 @@
+package dotty.tools.dotc
+package core
+package pickling
+
+import util.Positions._
+import ast.tpd._
+import ast.Trees.Lazy
+import core.Contexts._
+
+object Traversals {
+    
+  def traverse(x: Any, leftOp: Tree => Unit, rightOp: Tree => Unit) = {
+    
+    def traverseElems(xs: TraversableOnce[Any]) = xs.foreach(traverseElem)
+    
+    def traverseElem(x: Any): Unit = x match {
+      case x: Tree @ unchecked =>
+/** TODO: pickle annotation positions 
+        x match {
+          case x: MemberDef => traverseElems(x.symbol.annotations)
+          case _ =>
+        } */
+        leftOp(x)
+        x match {
+          case x: ValDef => 
+            traverseElem(x.tpt)
+            traverseUnlessLazy(x.unforced)
+          case x: DefDef => 
+            traverseElems(x.tparams)
+            traverseElems(x.vparamss)
+            traverseElem(x.tpt)
+            traverseUnlessLazy(x.unforced)
+          case x: Template =>
+            traverseElem(x.constr)
+            traverseElems(x.parents)
+            traverseElem(x.self)
+            traverseUnlessLazy(x.unforced)
+          case _ =>
+            traverseElems(x.productIterator)
+        }
+        rightOp(x)
+      case xs: List[_] =>
+        traverseElems(xs)
+      case _ =>
+    }
+    
+    def traverseUnlessLazy(x: Any) = x match {
+      case x: Lazy[_] => 
+      case _ => traverseElem(x)
+    }
+ 
+    traverseElem(x)
+  }
+}
\ No newline at end of file
diff --git a/src/dotty/tools/dotc/core/pickling/TreesUnpickler.scala b/src/dotty/tools/dotc/core/pickling/TreeUnpickler.scala
index de999f459..c037b4336 100644
--- a/src/dotty/tools/dotc/core/pickling/TreesUnpickler.scala
+++ b/src/dotty/tools/dotc/core/pickling/TreeUnpickler.scala
@@ -21,7 +21,7 @@ import typer.Mode
  *                        instead of creating a new symbol.
  *  @param readPositions  a flag indicating whether positions should be read
  */
-class TreesUnpickler(reader: TastyReader, tastyName: TastyName.Table, readPositions: Boolean) {
+class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, readPositions: Boolean) {
   import dotty.tools.dotc.core.pickling.PickleFormat._
   import TastyName._
   import tpd._
@@ -258,7 +258,6 @@ class TreesUnpickler(reader: TastyReader, tastyName: TastyName.Table, readPositi
       val tag = readByte()
       val end = readEnd()
       val name = if (tag == TYPEDEF || tag == TYPEPARAM) readName().toTypeName else readName()
-      // println(i"creating symbol $name at $start")
       skipParams()
       val isAbstractType = nextByte == TYPEBOUNDS
       val isClass = nextByte == TEMPLATE
@@ -271,6 +270,7 @@ class TreesUnpickler(reader: TastyReader, tastyName: TastyName.Table, readPositi
           result
         }
       val (givenFlags, annots, privateWithin) = readModifiers(end)
+      println(i"creating symbol $name at $start with flags $givenFlags")
       val lacksDefinition =
         rhsIsEmpty && !name.isConstructorName && !givenFlags.is(ParamOrAccessor) ||
         isAbstractType
@@ -432,7 +432,7 @@ class TreesUnpickler(reader: TastyReader, tastyName: TastyName.Table, readPositi
 
       val name = readName()
       // println(s"reading def of $name at $start")
-      val tree = tag match {
+      val tree: MemberDef = tag match {
         case DEFDEF =>
           val tparams = readParams[TypeDef](TYPEPARAM)(localCtx)
           val vparamss = readParamss(localCtx)
@@ -468,6 +468,10 @@ class TreesUnpickler(reader: TastyReader, tastyName: TastyName.Table, readPositi
           sym.info = readType()
           ValDef(sym.asTerm)
       }
+      val mods = 
+        if (sym.annotations.isEmpty) EmptyModifiers
+        else Modifiers(annotations = sym.annotations.map(_.tree))
+      tree.withMods(mods) // record annotations in tree so that tree positions can be filled in.
       skipTo(end)
       addAddr(start, tree)
     }
@@ -658,7 +662,7 @@ class TreesUnpickler(reader: TastyReader, tastyName: TastyName.Table, readPositi
     def readLater[T](end: Addr, op: TreeReader => T): Trees.Lazy[T] = {
       val localReader = fork
       skipTo(end)
-      new Trees.Lazy[T] { def complete: T = op(localReader) }
+      new LazyReader(localReader, op)
     }
     
 // ------ Hooks for positions ------------------------------------------------
@@ -678,4 +682,14 @@ class TreesUnpickler(reader: TastyReader, tastyName: TastyName.Table, readPositi
       else tree
     }
   }
+
+  class LazyReader[T](reader: TreeReader, op: TreeReader => T) extends Trees.Lazy[T] {
+    var posReader: Option[PositionReader] = None
+    def complete: T = {
+      val res = op(reader)
+      posReader.foreach(_.unpickle(res))
+      res
+    }  
+  }
+    
 }