Add TreeOps.setDefTree

Author: nicolasstucki

compiler/src/dotty/tools/dotc/ast/tpd.scala

@@ -1001,6 +1001,13 @@ object tpd extends Trees.Instance[Type] with TypedTreeInfo {
       foreachSubTree { tree => if (f(tree)) buf += tree }
       buf.toList
     }
+
+    /** Set this tree as the `defTree` of its symbol and return this tree */
+    def setDefTree(implicit ctx: Context): ThisTree = {
+      val sym = tree.symbol
+      if (sym.exists) sym.defTree = tree
+      tree
+    }
   }
 
   /** Map Inlined nodes, NamedArgs, Blocks with no statements and local references to underlying arguments.

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

@@ -851,8 +851,6 @@ class TreeUnpickler(reader: TastyReader,
         sym.info = ta.avoidPrivateLeaks(sym, tree.sourcePos)
       }
 
-      sym.defTree = tree
-
       if (ctx.mode.is(Mode.ReadComments)) {
         assert(ctx.docCtx.isDefined, "Mode is `ReadComments`, but no `docCtx` is set.")
         commentUnpicklerOpt.foreach { commentUnpickler =>
@@ -862,7 +860,7 @@ class TreeUnpickler(reader: TastyReader,
         }
       }
 
-      tree
+      tree.setDefTree
     }
 
     private def readTemplate(implicit ctx: Context): Template = {

compiler/src/dotty/tools/dotc/typer/EtaExpansion.scala

@@ -47,9 +47,7 @@ abstract class Lifter {
       var liftedType = expr.tpe.widen
       if (liftedFlags.is(Method)) liftedType = ExprType(liftedType)
       val lifted = ctx.newSymbol(ctx.owner, name, liftedFlags | Synthetic, liftedType, coord = spanCoord(expr.span))
-      val ddef = liftedDef(lifted, expr).withSpan(expr.span)
-      lifted.defTree = ddef
-      defs += ddef
+      defs += liftedDef(lifted, expr).withSpan(expr.span).setDefTree
       ref(lifted.termRef).withSpan(expr.span.focus)
     }
 

compiler/src/dotty/tools/dotc/typer/Inliner.scala

@@ -260,8 +260,7 @@ class Inliner(call: tpd.Tree, rhsToInline: tpd.Tree)(implicit ctx: Context) {
       if (isByName) DefDef(boundSym, arg.changeOwner(ctx.owner, boundSym))
       else ValDef(boundSym, arg)
     }.withSpan(boundSym.span)
-    boundSym.defTree = binding
-    bindingsBuf += binding
+    bindingsBuf += binding.setDefTree
     binding
   }
 
@@ -313,9 +312,8 @@ class Inliner(call: tpd.Tree, rhsToInline: tpd.Tree)(implicit ctx: Context) {
           ref(rhsClsSym.sourceModule)
         else
           inlineCallPrefix
-      val binding = ValDef(selfSym.asTerm, rhs).withSpan(selfSym.span)
+      val binding = ValDef(selfSym.asTerm, rhs).withSpan(selfSym.span).setDefTree
       bindingsBuf += binding
-      selfSym.defTree = binding
       inlining.println(i"proxy at $level: $selfSym = ${bindingsBuf.last}")
       lastSelf = selfSym
       lastLevel = level
@@ -490,11 +488,9 @@ class Inliner(call: tpd.Tree, rhsToInline: tpd.Tree)(implicit ctx: Context) {
 
       // The normalized bindings collected in `bindingsBuf`
       bindingsBuf.transform { binding =>
-        val transformedBinding = reducer.normalizeBinding(binding)(inlineCtx)
         // Set trees to symbols allow macros to see the definition tree.
         // This is used by `underlyingArgument`.
-        transformedBinding.symbol.defTree = transformedBinding
-        transformedBinding
+        reducer.normalizeBinding(binding)(inlineCtx).setDefTree
       }
 
       // Run a typing pass over the inlined tree. See InlineTyper for details.

compiler/src/dotty/tools/dotc/typer/Typer.scala

@@ -1489,11 +1489,9 @@ class Typer extends Namer
     val vdef1 = assignType(cpy.ValDef(vdef)(name, tpt1, rhs1), sym)
     if (sym.is(Inline, butNot = DeferredOrTermParamOrAccessor))
       checkInlineConformant(rhs1, isFinal = sym.is(Final), em"right-hand side of inline $sym")
-    if (sym.exists)
-      sym.defTree = vdef1
     patchIfLazy(vdef1)
     patchFinalVals(vdef1)
-    vdef1
+    vdef1.setDefTree
   }
 
   /** Add a @volitile to lazy vals when rewriting from Scala2 */
@@ -1564,9 +1562,7 @@ class Typer extends Namer
       for (param <- tparams1 ::: vparamss1.flatten)
         checkRefsLegal(param, sym.owner, (name, sym) => sym.is(TypeParam), "secondary constructor")
 
-    val ddef1 = assignType(cpy.DefDef(ddef)(name, tparams1, vparamss1, tpt1, rhs1), sym)
-    sym.defTree = ddef1
-    ddef1
+    assignType(cpy.DefDef(ddef)(name, tparams1, vparamss1, tpt1, rhs1), sym).setDefTree
       //todo: make sure dependent method types do not depend on implicits or by-name params
   }