Retain tree at pickling time

Author: nicolasstucki

compiler/src/dotty/tools/dotc/core/tasty/TreePickler.scala

@@ -291,7 +291,9 @@ class TreePickler(pickler: TastyPickler) {
   def pickleTreeUnlessEmpty(tree: Tree)(implicit ctx: Context): Unit =
     if (!tree.isEmpty) pickleTree(tree)
 
-  def pickleDef(tag: Int, sym: Symbol, tpt: Tree, rhs: Tree = EmptyTree, pickleParams: => Unit = ())(implicit ctx: Context) = {
+  def pickleDef(tag: Int, tree: MemberDef, tpt: Tree, rhs: Tree = EmptyTree, pickleParams: => Unit = ())(implicit ctx: Context) = {
+    val sym = tree.symbol
+    sym.defTree = tree
     assert(symRefs(sym) == NoAddr, sym)
     registerDef(sym)
     writeByte(tag)
@@ -310,9 +312,9 @@ class TreePickler(pickler: TastyPickler) {
   def pickleParam(tree: Tree)(implicit ctx: Context): Unit = {
     registerTreeAddr(tree)
     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)
+      case tree: ValDef => pickleDef(PARAM, tree, tree.tpt)
+      case tree: DefDef => pickleDef(PARAM, tree, tree.tpt, tree.rhs)
+      case tree: TypeDef => pickleDef(TYPEPARAM, tree, tree.rhs)
     }
   }
 
@@ -467,7 +469,7 @@ class TreePickler(pickler: TastyPickler) {
             patterns.foreach(pickleTree)
           }
         case tree: ValDef =>
-          pickleDef(VALDEF, tree.symbol, tree.tpt, tree.rhs)
+          pickleDef(VALDEF, tree, tree.tpt, tree.rhs)
         case tree: DefDef =>
           def pickleAllParams = {
             pickleParams(tree.tparams)
@@ -476,9 +478,9 @@ class TreePickler(pickler: TastyPickler) {
               withLength { pickleParams(vparams) }
             }
           }
-          pickleDef(DEFDEF, tree.symbol, tree.tpt, tree.rhs, pickleAllParams)
+          pickleDef(DEFDEF, tree, tree.tpt, tree.rhs, pickleAllParams)
         case tree: TypeDef =>
-          pickleDef(TYPEDEF, tree.symbol, tree.rhs)
+          pickleDef(TYPEDEF, tree, tree.rhs)
         case tree: Template =>
           registerDef(tree.symbol)
           writeByte(TEMPLATE)

compiler/src/dotty/tools/dotc/transform/PostTyper.scala

@@ -105,7 +105,6 @@ class PostTyper extends MacroTransform with IdentityDenotTransformer { thisPhase
       val sym = tree.symbol
       sym.registerIfChild()
       sym.transformAnnotations(transformAnnot)
-      sym.defTree = tree
     }
 
     private def transformSelect(tree: Select, targs: List[Tree])(implicit ctx: Context): Tree = {