TASTy reflect load complete definition trees

Currently only works with -Yretain-tree.

Also fix an issue with TASTy reflect

Author: nicolasstucki

compiler/src/dotty/tools/dotc/Compiler.scala

@@ -52,6 +52,7 @@ class Compiler {
   /** Phases dealing with TASTY tree pickling and unpickling */
   protected def picklerPhases: List[List[Phase]] =
     List(new Pickler) ::            // Generate TASTY info
+    List(new RetainTrees) ::        // Retain trees of the compilation units if -Yretain-trees is set
     List(new ReifyQuotes) ::        // Turn quoted trees into explicit run-time data structures
     Nil
 

compiler/src/dotty/tools/dotc/tastyreflect/FromSymbol.scala

@@ -10,6 +10,7 @@ import dotty.tools.dotc.core.Types._
 object FromSymbol {
 
   def definitionFromSym(sym: Symbol)(implicit ctx: Context): tpd.Tree = {
+    assert(sym.exists)
     if (sym.is(Package)) packageDefFromSym(sym)
     else if (sym.isClass) classDef(sym.asClass)
     else if (sym.isType) typeDefFromSym(sym.asType)
@@ -19,21 +20,52 @@ object FromSymbol {
 
   def packageDefFromSym(sym: Symbol)(implicit ctx: Context): PackageDefinition = PackageDefinitionImpl(sym)
 
-  def classDef(cls: ClassSymbol)(implicit ctx: Context): tpd.TypeDef = {
-    val constrSym = cls.unforcedDecls.find(_.isPrimaryConstructor).orElse(
-      // Dummy constructor for classes such as `<refinement>`
-      ctx.newSymbol(cls, nme.CONSTRUCTOR, EmptyFlags, NoType)
-    )
-    val constr = tpd.DefDef(constrSym.asTerm)
-    val parents = cls.classParents.map(tpd.TypeTree(_))
-    val body = cls.unforcedDecls.filter(!_.isPrimaryConstructor).map(s => definitionFromSym(s))
-    tpd.ClassDefWithParents(cls, constr, parents, body)
+  def classDef(cls: ClassSymbol)(implicit ctx: Context): tpd.TypeDef = cls.tree match {
+    case tree: tpd.TypeDef => tree
+    case tpd.EmptyTree =>
+      val constrSym = cls.unforcedDecls.find(_.isPrimaryConstructor).orElse(
+        // Dummy constructor for classes such as `<refinement>`
+        ctx.newSymbol(cls, nme.CONSTRUCTOR, EmptyFlags, NoType)
+      )
+      val constr = tpd.DefDef(constrSym.asTerm)
+      val parents = cls.classParents.map(tpd.TypeTree(_))
+      val body = cls.unforcedDecls.filter(!_.isPrimaryConstructor).map(s => definitionFromSym(s))
+      tpd.ClassDefWithParents(cls, constr, parents, body)
   }
 
-  def typeDefFromSym(sym: TypeSymbol)(implicit ctx: Context): tpd.TypeDef = tpd.TypeDef(sym)
+  // TODO: this logic could be moved inside sym.tree
 
-  def defDefFromSym(sym: TermSymbol)(implicit ctx: Context): tpd.DefDef = tpd.DefDef(sym)
+  def typeDefFromSym(sym: TypeSymbol)(implicit ctx: Context): tpd.TypeDef = tree(sym) match {
+    case tree: tpd.TypeDef => tree
+    case tpd.EmptyTree => tpd.TypeDef(sym)
+  }
+
+  def defDefFromSym(sym: TermSymbol)(implicit ctx: Context): tpd.DefDef = tree(sym) match {
+    case tree: tpd.DefDef => tree
+    case tpd.EmptyTree => tpd.DefDef(sym)
+  }
+
+  def valDefFromSym(sym: TermSymbol)(implicit ctx: Context): tpd.ValDef = tree(sym) match {
+    case tree: tpd.ValDef => tree
+    case tpd.EmptyTree => tpd.ValDef(sym)
+  }
+
+  private def tree(sym: Symbol)(implicit ctx: Context): tpd.Tree = sym.topLevelClass match {
+    case top: ClassSymbol =>
+      val findTree = new tpd.TreeAccumulator[tpd.Tree] {
+        def apply(x: tpd.Tree, tree: tpd.Tree)(implicit ctx: Context): tpd.Tree = {
+          if (!x.isEmpty) x
+          else tree match {
+            case tree: tpd.DefTree if tree.symbol == sym => tree
+            // TODO avoid searching in all the tree, look at sym.ownerIterator.toSet
+            case _ => foldOver(x, tree)
+          }
 
-  def valDefFromSym(sym: TermSymbol)(implicit ctx: Context): tpd.ValDef = tpd.ValDef(sym)
+        }
+      }
+      findTree(tpd.EmptyTree, top.tree)
+
+    case _ => tpd.EmptyTree
+  }
 
 }

compiler/src/dotty/tools/dotc/tastyreflect/TreeOpsImpl.scala

@@ -344,7 +344,7 @@ trait TreeOpsImpl extends scala.tasty.reflect.TreeOps with TastyCoreImpl with He
     }
 
     object Inlined extends InlinedExtractor {
-      def unapply(x: Term)(implicit ctx: Context): Option[(Option[Term], List[Statement], Term)] = x match {
+      def unapply(x: Term)(implicit ctx: Context): Option[(Option[Parent], List[Statement], Term)] = x match {
         case x: tpd.Inlined =>
           Some((optional(x.call), x.bindings, x.expansion))
         case _ => None

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

@@ -0,0 +1,26 @@
+package dotty.tools.dotc
+package transform
+
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.ast.Trees._
+import dotty.tools.dotc.core.Contexts._
+import dotty.tools.dotc.core.Phases._
+
+/** Retain trees of the compilation units if -Yretain-trees is set */
+class RetainTrees extends Phase {
+  import tpd._
+
+  def phaseName: String = "RetainTrees"
+
+  def run(implicit ctx: Context): Unit = {
+    if (ctx.settings.YretainTrees.value)
+      retainTopLevelClassTrees(ctx.compilationUnit.tpdTree)
+  }
+
+  private def retainTopLevelClassTrees(tree: Tree)(implicit ctx: Context): Unit = tree match {
+    case PackageDef(_, stats) => stats.foreach(retainTopLevelClassTrees)
+    case tree: TypeDef => tree.symbol.asClass.treeOrProvider = tree
+    case _ =>
+  }
+
+}

compiler/test/dotty/tools/dotc/CompilationTests.scala

@@ -168,6 +168,7 @@ class CompilationTests extends ParallelTesting {
 
   @Test def runAll: Unit = {
     implicit val testGroup: TestGroup = TestGroup("runAll")
+    compileFilesInDir("tests/run-custom-args/Yretain-trees", defaultOptions and "-Yretain-trees") +
     compileFilesInDir("tests/run", defaultOptions)
   }.checkRuns()
 

library/src/scala/tasty/reflect/TreeOps.scala

@@ -291,7 +291,7 @@ trait TreeOps extends TastyCore {
 
     val Inlined: InlinedExtractor
     abstract class InlinedExtractor {
-      def unapply(tree: Tree)(implicit ctx: Context): Option[(Option[Term], List[Definition], Term)]
+      def unapply(tree: Tree)(implicit ctx: Context): Option[(Option[Parent], List[Definition], Term)]
     }
 
     val SelectOuter: SelectOuterExtractor

library/src/scala/tasty/util/ShowExtractors.scala

@@ -67,7 +67,9 @@ class ShowExtractors[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
       case Term.Repeated(elems) =>
         this += "Term.Repeated(" ++= elems += ")"
       case Term.Inlined(call, bindings, expansion) =>
-        this += "Term.Inlined(" += call += ", " ++= bindings += ", " += expansion += ")"
+        this += "Term.Inlined("
+        visitOption(call, visitParent)
+        this += ", " ++= bindings += ", " += expansion += ")"
       case ValDef(name, tpt, rhs) =>
         this += "ValDef(\"" += name += "\", " += tpt += ", " += rhs += ")"
       case DefDef(name, typeParams, paramss, returnTpt, rhs) =>
@@ -76,10 +78,7 @@ class ShowExtractors[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
         this += "TypeDef(\"" += name += "\", " += rhs += ")"
       case ClassDef(name, constr, parents, self, body) =>
         this += "ClassDef(\"" += name += "\", " += constr += ", "
-        visitList[Parent](parents, {
-          case IsTerm(parent) => this += parent
-          case IsTypeTree(parent) => this += parent
-        })
+        visitList[Parent](parents, visitParent)
         this += ", " += self += ", " ++= body += ")"
       case PackageDef(name, owner) =>
         this += "PackageDef(\"" += name += "\", " += owner += ")"
@@ -140,6 +139,11 @@ class ShowExtractors[T <: Tasty with Singleton](tasty0: T) extends Show[T](tasty
         this += "Pattern.TypeTest(" += tpt += ")"
     }
 
+    def visitParent(x: Parent): Buffer = x match {
+      case IsTerm(parent) => this += parent
+      case IsTypeTree(parent) => this += parent
+    }
+
     def visitConstant(x: Constant): Buffer = x match {
       case Constant.Unit() => this += "Constant.Unit()"
       case Constant.Null() => this += "Constant.Null()"

tests/run-custom-args/Yretain-trees/tasty-definitions-2.check

@@ -0,0 +1,2 @@
+DefDef("foo", Nil, Nil, TypeTree.TypeIdent("Int"), Some(Term.Apply(Term.Select(Term.Literal(Constant.Int(1)), "+", Some(Signature(List(scala.Int), scala.Int))), List(Term.Literal(Constant.Int(2))))))
+ValDef("bar", TypeTree.TypeIdent("Int"), Some(Term.Apply(Term.Select(Term.Literal(Constant.Int(2)), "+", Some(Signature(List(scala.Int), scala.Int))), List(Term.Literal(Constant.Int(3))))))