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/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,32 @@ 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.defTree 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)
+  def typeDefFromSym(sym: TypeSymbol)(implicit ctx: Context): tpd.TypeDef = sym.defTree match {
+    case tree: tpd.TypeDef => tree
+    case tpd.EmptyTree => tpd.TypeDef(sym)
+  }
 
-  def defDefFromSym(sym: TermSymbol)(implicit ctx: Context): tpd.DefDef = tpd.DefDef(sym)
+  def defDefFromSym(sym: TermSymbol)(implicit ctx: Context): tpd.DefDef = sym.defTree match {
+    case tree: tpd.DefDef => tree
+    case tpd.EmptyTree => tpd.DefDef(sym)
+  }
 
-  def valDefFromSym(sym: TermSymbol)(implicit ctx: Context): tpd.ValDef = tpd.ValDef(sym)
+  def valDefFromSym(sym: TermSymbol)(implicit ctx: Context): tpd.ValDef = sym.defTree match {
+    case tree: tpd.ValDef => tree
+    case tpd.EmptyTree => tpd.ValDef(sym)
+  }
 
 }

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

@@ -313,7 +313,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/test/dotty/tools/dotc/CompilationTests.scala

@@ -169,6 +169,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

@@ -285,7 +285,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))))))

tests/run/tasty-definitions-2/Macro_1.scala renamed to tests/run-custom-args/Yretain-trees/tasty-definitions-2/Macro_1.scala

@@ -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))))))