Preserve all positions in Tasty.

The goal is that pickled and unpickled trees should print
the same with -Yprintpos. There are several reasons why this
is not the case so far. Some of them are fixed in this
commit.

5 files changed, 86 insertions, 45 deletions

diff --git a/src/dotty/tools/dotc/core/tasty/PositionPickler.scala b/src/dotty/tools/dotc/core/tasty/PositionPickler.scala
index 63bb00a71..65e9d12be 100644
--- a/src/dotty/tools/dotc/core/tasty/PositionPickler.scala
+++ b/src/dotty/tools/dotc/core/tasty/PositionPickler.scala
@@ -37,23 +37,27 @@ class PositionPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr
       lastIndex = index
       lastPos = pos
     }
-    def traverse(x: Any, parentPos: Position): Unit = x match {
+    def traverse(x: Any): Unit = x match {
       case x: Tree @unchecked =>
-        if (x.pos.exists && x.pos.toSynthetic != parentPos.toSynthetic) {
+        if (x.pos.exists /*&& x.pos.toSynthetic != x.initialPos.toSynthetic*/) {
           addrOfTree(x) match {
-            case Some(addr) => pickleDeltas(addr.index, x.pos)
+            case Some(addr) =>
+              //println(i"pickling $x")
+              pickleDeltas(addr.index, x.pos)
             case _ =>
+              //println(i"no address for $x")
           }
         }
+        //else println(i"skipping $x")
         x match {
-          case x: MemberDef @unchecked => traverse(x.symbol.annotations, x.pos)
+          case x: MemberDef @unchecked => traverse(x.symbol.annotations)
           case _ =>
         }
-        traverse(x.productIterator, x.pos)
+        traverse(x.productIterator)
       case xs: TraversableOnce[_] =>
-        xs.foreach(traverse(_, parentPos))
+        xs.foreach(traverse)
       case _ =>
     }
-    traverse(roots, NoPosition)
+    traverse(roots)
   }
 }
diff --git a/src/dotty/tools/dotc/core/tasty/TastyFormat.scala b/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
index 8e8d58b47..fc551658b 100644
--- a/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
+++ b/src/dotty/tools/dotc/core/tasty/TastyFormat.scala
@@ -57,7 +57,8 @@ Standard-Section: "ASTs" TopLevelStat*
                   TYPEDEF        Length NameRef (Type | Template) Modifier*
                   IMPORT         Length qual_Term Selector*
   Selector      = IMPORTED              name_NameRef
-                  RENAMED        Length from_NameRef to_NameRef
+                  RENAMED               to_NameRef
+
                                  // Imports are for scala.meta, they are not used in the backend
 
   TypeParam     = TYPEPARAM      Length NameRef Type Modifier*
@@ -264,6 +265,7 @@ object TastyFormat {
   final val DOUBLEconst = 76
   final val STRINGconst = 77
   final val IMPORTED = 78
+  final val RENAMED = 79
 
   final val THIS = 96
   final val CLASSconst = 97
@@ -291,11 +293,10 @@ object TastyFormat {
   final val TYPEPARAM = 133
   final val PARAMS = 134
   final val PARAM = 136
-  final val RENAMED = 138
+
   final val APPLY = 139
   final val TYPEAPPLY = 140
 
-
   final val TYPED = 143
   final val NAMEDARG = 144
   final val ASSIGN = 145
diff --git a/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala b/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala
index d741c42c3..f8f7c330f 100644
--- a/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala
+++ b/src/dotty/tools/dotc/core/tasty/TreeBuffer.scala
@@ -6,7 +6,7 @@ package tasty
 import util.Util.{bestFit, dble}
 import TastyBuffer.{Addr, AddrWidth}
 import config.Printers.pickling
-import ast.tpd.Tree
+import ast.untpd.Tree
 
 class TreeBuffer extends TastyBuffer(50000) {
 
@@ -19,6 +19,8 @@ class TreeBuffer extends TastyBuffer(50000) {
 
   private[tasty] val pickledTrees = new java.util.IdentityHashMap[Tree, Any] // Value type is really Addr, but that's not compatible with null
 
+  def registerTreeAddr(tree: Tree) = pickledTrees.put(tree, currentAddr)
+
   def addrOfTree(tree: Tree): Option[Addr] = pickledTrees.get(tree) match {
     case null => None
     case n => Some(n.asInstanceOf[Addr])
diff --git a/src/dotty/tools/dotc/core/tasty/TreePickler.scala b/src/dotty/tools/dotc/core/tasty/TreePickler.scala
index 8889e8a5c..215cef295 100644
--- a/src/dotty/tools/dotc/core/tasty/TreePickler.scala
+++ b/src/dotty/tools/dotc/core/tasty/TreePickler.scala
@@ -4,6 +4,7 @@ package core
 package tasty
 
 import ast.Trees._
+import ast.untpd
 import TastyFormat._
 import Contexts._, Symbols._, Types._, Names._, Constants._, Decorators._, Annotations._, StdNames.tpnme, NameOps._
 import collection.mutable
@@ -299,7 +300,10 @@ class TreePickler(pickler: TastyPickler) {
     pickled
   }
 
-  def pickleTpt(tpt: Tree)(implicit ctx: Context): Unit = pickleType(tpt.tpe) // TODO correlate with original when generating positions
+  def pickleTpt(tpt: Tree)(implicit ctx: Context): Unit = {
+    pickledTrees.put(tpt, currentAddr)
+    pickleType(tpt.tpe) // TODO correlate with original when generating positions
+  }
 
   def pickleTreeUnlessEmpty(tree: Tree)(implicit ctx: Context): Unit =
     if (!tree.isEmpty) pickleTree(tree)
@@ -313,17 +317,20 @@ class TreePickler(pickler: TastyPickler) {
       pickleParams
       tpt match {
         case tpt: TypeTree => pickleTpt(tpt)
-        case _ => pickleTree(tpt)
+        case templ: Template => pickleTree(tpt)
       }
       pickleTreeUnlessEmpty(rhs)
       pickleModifiers(sym)
     }
   }
 
-  def pickleParam(tree: Tree)(implicit ctx: Context): Unit = tree match {
-    case tree: ValDef => pickleDef(PARAM, tree.symbol, tree.tpt)
-    case tree: DefDef => pickleDef(PARAM, tree.symbol, tree.tpt, tree.rhs)
-    case tree: TypeDef => pickleDef(TYPEPARAM, tree.symbol, tree.rhs)
+  def pickleParam(tree: Tree)(implicit ctx: Context): Unit = {
+    pickledTrees.put(tree, currentAddr)
+    tree match {
+      case tree: ValDef => pickleDef(PARAM, tree.symbol, tree.tpt)
+      case tree: DefDef => pickleDef(PARAM, tree.symbol, tree.tpt, tree.rhs)
+      case tree: TypeDef => pickleDef(TYPEPARAM, tree.symbol, tree.rhs)
+    }
   }
 
   def pickleParams(trees: List[Tree])(implicit ctx: Context): Unit = {
@@ -500,12 +507,11 @@ class TreePickler(pickler: TastyPickler) {
         withLength {
           pickleTree(expr)
           selectors foreach {
-            case Thicket(Ident(from) :: Ident(to) :: Nil) =>
-              writeByte(RENAMED)
-              withLength { pickleName(from); pickleName(to) }
-            case Ident(name) =>
-              writeByte(IMPORTED)
-              pickleName(name)
+            case Thicket((from @ Ident(_)) :: (to @ Ident(_)) :: Nil) =>
+              pickleSelector(IMPORTED, from)
+              pickleSelector(RENAMED, to)
+            case id @ Ident(_) =>
+              pickleSelector(IMPORTED, id)
           }
         }
       case PackageDef(pid, stats) =>
@@ -518,6 +524,12 @@ class TreePickler(pickler: TastyPickler) {
       throw ex
   }
 
+  def pickleSelector(tag: Int, id: untpd.Ident)(implicit ctx: Context): Unit = {
+    pickledTrees.put(id, currentAddr)
+    writeByte(tag)
+    pickleName(id.name)
+  }
+
   def qualifiedName(sym: Symbol)(implicit ctx: Context): TastyName =
     if (sym.isRoot || sym.owner.isRoot) TastyName.Simple(sym.name.toTermName)
     else TastyName.Qualified(nameIndex(qualifiedName(sym.owner)), nameIndex(sym.name))
diff --git a/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala b/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
index f67159808..7038f6e0a 100644
--- a/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
+++ b/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala
@@ -626,7 +626,7 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
      *  or else read definition.
      */
     def readIndexedDef()(implicit ctx: Context): Tree = treeAtAddr.remove(currentAddr) match {
-      case Some(tree) => skipTree(); tree
+      case Some(tree) => skipTree(); setPos(currentAddr, tree)
       case none => readNewDef()
     }
 
@@ -655,12 +655,18 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
 
       def localCtx = localContext(sym)
 
+      def ValDef(tpt: Tree) =
+        ta.assignType(untpd.ValDef(sym.name.asTermName, tpt, readRhs(localCtx)), sym)
+
       def DefDef(tparams: List[TypeDef], vparamss: List[List[ValDef]], tpt: Tree) =
          ta.assignType(
             untpd.DefDef(
               sym.name.asTermName, tparams, vparamss, tpt, readRhs(localCtx)),
             sym)
 
+      def TypeDef(rhs: Tree) =
+        ta.assignType(untpd.TypeDef(sym.name.asTypeName, rhs), sym)
+
       def ta =  ctx.typeAssigner
 
       val name = readName()
@@ -682,8 +688,9 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
           }
           DefDef(tparams, vparamss, tpt)
         case VALDEF =>
-          sym.info = readType()
-          ValDef(sym.asTerm, readRhs(localCtx))
+          val tpt = readTpt()
+          sym.info = tpt.tpe
+          ValDef(tpt)
         case TYPEDEF | TYPEPARAM =>
           if (sym.isClass) {
             val companion = sym.scalacLinkedClass
@@ -699,22 +706,23 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
               if (sym is Flags.ModuleClass) sym.registerCompanionMethod(nme.COMPANION_CLASS_METHOD, companion)
               else sym.registerCompanionMethod(nme.COMPANION_MODULE_METHOD, companion)
             }
-            ta.assignType(untpd.TypeDef(sym.name.asTypeName, readTemplate(localCtx)), sym)
+            TypeDef(readTemplate(localCtx))
           } else {
-            sym.info = readType()
-            TypeDef(sym.asType)
+            val rhs = readTpt()
+            sym.info = rhs.tpe
+            TypeDef(rhs)
           }
         case PARAM =>
-          val info = readType()
+          val tpt = readTpt()
           if (noRhs(end)) {
-            sym.info = info
-            ValDef(sym.asTerm)
+            sym.info = tpt.tpe
+            ValDef(tpt)
           }
           else {
             sym.setFlag(Method)
-            sym.info = ExprType(info)
+            sym.info = ExprType(tpt.tpe)
             pickling.println(i"reading param alias $name -> $currentAddr")
-            DefDef(Nil, Nil, TypeTree(info))
+            DefDef(Nil, Nil, tpt)
           }
       }
       val mods =
@@ -752,6 +760,7 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
       val parentRefs = ctx.normalizeToClassRefs(parents.map(_.tpe), cls, cls.unforcedDecls)
       val self =
         if (nextByte == SELFDEF) {
+          val selfStart = currentAddr
           readByte()
           untpd.ValDef(readName(), readTpt(), EmptyTree).withType(NoType)
         }
@@ -816,21 +825,28 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
     }
 
     def readImport()(implicit ctx: Context): Tree = {
+      val start = currentAddr
       readByte()
       readEnd()
       val expr = readTerm()
       def readSelectors(): List[untpd.Tree] = nextByte match {
-        case RENAMED =>
-          readByte()
-          readEnd()
-          untpd.Thicket(untpd.Ident(readName()), untpd.Ident(readName())) :: readSelectors()
         case IMPORTED =>
+          val start = currentAddr
           readByte()
-          untpd.Ident(readName()) :: readSelectors()
-        case _ =>
-          Nil
+          val from = setPos(start, untpd.Ident(readName()))
+          nextByte match {
+            case RENAMED =>
+              val start2 = currentAddr
+              readByte()
+              val to = setPos(start2, untpd.Ident(readName()))
+              untpd.Thicket(from, to) :: readSelectors()
+            case _ =>
+              from :: readSelectors()
+          }
+          case _ =>
+            Nil
       }
-      Import(expr, readSelectors())
+      setPos(start, Import(expr, readSelectors()))
     }
 
     def readIndexedStats(exprOwner: Symbol, end: Addr)(implicit ctx: Context): List[Tree] =
@@ -1002,11 +1018,17 @@ class TreeUnpickler(reader: TastyReader, tastyName: TastyName.Table, posUnpickle
 // ------ Setting positions ------------------------------------------------
 
     /** Set position of `tree` at given `addr`. */
-    def setPos[T <: Tree](addr: Addr, tree: T)(implicit ctx: Context): tree.type =
+    def setPos[T <: untpd.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
+          case Some(posUnpickler) =>
+            //println(i"setPos $tree / ${tree.getClass} to ${posUnpickler.posAt(addr)}")
+            val pos = posUnpickler.posAt(addr)
+            if (pos.exists) tree.setPosUnchecked(pos)
+            tree
+          case _  =>
+            //println(i"no pos $tree")
+            tree
         }
       }
       else tree