Add ClassSymbol#tree to get the tree of a top-level class

When typechecking a class from source or when unpickling it from tasty,
we normally do not keep a reference to the class tree, but there's a
wealth of information only present in the tree that we may wish to use.
In an IDE for example, this makes it easy to provide features like "jump
to definition". This commit unlocks this potential by adding a
`ClassSymbol#tree` method that, when used together with the
`-Yretain-trees` flag, returns the tree corresponding to a top-level
class.

This also replaces the `unpicklers` map in `CompilationUnit` by
`ClassSymbol#unpickler`.

Author: smarter

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

@@ -19,6 +19,4 @@ class CompilationUnit(val source: SourceFile) {
 
   /** Pickled TASTY binaries, indexed by class. */
   var pickled: Map[ClassSymbol, Array[Byte]] = Map()
-
-  var unpicklers: Map[ClassSymbol, TastyUnpickler] = Map()
 }

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

@@ -83,17 +83,15 @@ object FromTasty extends Driver {
           case clsd: ClassDenotation =>
             clsd.infoOrCompleter match {
               case info: ClassfileLoader =>
-                info.load(clsd) match {
-                  case Some(unpickler: DottyUnpickler) =>
-                    val List(unpickled) = unpickler.body(ctx.addMode(Mode.ReadPositions))
+                info.load(clsd)
+                val unpickled = clsd.symbol.asClass.tree
+                if (unpickled != null) {
                     val unit1 = new CompilationUnit(new SourceFile(clsd.symbol.sourceFile, Seq()))
                     unit1.tpdTree = unpickled
-                    unit1.unpicklers += (clsd.classSymbol -> unpickler.unpickler)
                     force.traverse(unit1.tpdTree)
                     unit1
-                  case _ =>
+                } else
                     cannotUnpickle(s"its class file ${info.classfile} does not have a TASTY attribute")
-                }
               case info =>
                 cannotUnpickle(s"its info of type ${info.getClass} is not a ClassfileLoader")
             }

compiler/src/dotty/tools/dotc/config/ScalaSettings.scala

@@ -98,6 +98,7 @@ class ScalaSettings extends Settings.SettingGroup {
   val YforceSbtPhases = BooleanSetting("-Yforce-sbt-phases", "Run the phases used by sbt for incremental compilation (ExtractDependencies and ExtractAPI) even if the compiler is ran outside of sbt, for debugging.")
   val YdumpSbtInc = BooleanSetting("-Ydump-sbt-inc", "For every compiled foo.scala, output the API representation and dependencies used for sbt incremental compilation in foo.inc, implies -Yforce-sbt-phases.")
   val YcheckAllPatmat = BooleanSetting("-Ycheck-all-patmat", "Check exhaustivity and redundancy of all pattern matching (used for testing the algorithm)")
+  val YretainTrees = BooleanSetting("-Yretain-trees", "Retain trees for top-level classes, accessible from ClassSymbol#tree")
 
   /** Area-specific debug output */
   val Yexplainlowlevel = BooleanSetting("-Yexplain-lowlevel", "When explaining type errors, show types at a lower level.")

compiler/src/dotty/tools/dotc/core/SymbolLoaders.scala

@@ -17,6 +17,7 @@ import Decorators.{PreNamedString, StringInterpolators}
 import classfile.ClassfileParser
 import util.Stats
 import scala.util.control.NonFatal
+import ast.Trees._
 
 object SymbolLoaders {
   /** A marker trait for a completer that replaces the original
@@ -320,9 +321,14 @@ class ClassfileLoader(val classfile: AbstractFile) extends SymbolLoader {
   override def doComplete(root: SymDenotation)(implicit ctx: Context): Unit =
     load(root)
 
-  def load(root: SymDenotation)(implicit ctx: Context): Option[ClassfileParser.Embedded] = {
+  def load(root: SymDenotation)(implicit ctx: Context): Unit = {
     val (classRoot, moduleRoot) = rootDenots(root.asClass)
-    new ClassfileParser(classfile, classRoot, moduleRoot)(ctx).run()
+    (new ClassfileParser(classfile, classRoot, moduleRoot)(ctx)).run() match {
+      case Some(unpickler: tasty.DottyUnpickler) if ctx.settings.YretainTrees.value =>
+        classRoot.symbol.asClass.unpickler = unpickler
+        moduleRoot.symbol.asClass.unpickler = unpickler
+      case _ =>
+    }
   }
 }
 

compiler/src/dotty/tools/dotc/core/Symbols.scala

@@ -20,7 +20,8 @@ import DenotTransformers._
 import StdNames._
 import NameOps._
 import NameKinds.LazyImplicitName
-import ast.tpd.Tree
+import ast.tpd
+import tpd.Tree
 import ast.TreeTypeMap
 import Constants.Constant
 import reporting.diagnostic.Message
@@ -552,6 +553,34 @@ object Symbols {
 
     type ThisName = TypeName
 
+    /** If this is a top-level class, and if `-Yretain-trees` is set, return the TypeDef tree
+     *  for this class, otherwise EmptyTree.
+     */
+    def tree(implicit ctx: Context): tpd.Tree /* tpd.TypeDef | tpd.EmptyTree */ = {
+      // TODO: Consider storing this tree like we store lazy trees for inline functions
+      if (unpickler != null && !denot.isAbsent) {
+        assert(myTree.isEmpty)
+
+        import ast.Trees._
+
+        def findTree(tree: tpd.Tree): Option[tpd.TypeDef] = tree match {
+          case PackageDef(_, stats) =>
+            stats.flatMap(findTree).headOption
+          case tree: tpd.TypeDef if tree.symbol == this =>
+            Some(tree)
+          case _ =>
+              None
+        }
+        val List(unpickledTree) = unpickler.body(ctx.addMode(Mode.ReadPositions))
+        unpickler = null
+
+        myTree = findTree(unpickledTree).get
+      }
+      myTree
+    }
+    private[dotc] var myTree: tpd.Tree = tpd.EmptyTree
+    private[dotc] var unpickler: tasty.DottyUnpickler = _
+
     /** The source or class file from which this class was generated, null if not applicable. */
     override def associatedFile(implicit ctx: Context): AbstractFile =
       if (assocFile != null || (this.owner is PackageClass) || this.isEffectiveRoot) assocFile

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

@@ -46,8 +46,15 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded {
   def enter(roots: Set[SymDenotation])(implicit ctx: Context): Unit =
     treeUnpickler.enterTopLevel(roots)
 
+  /** Only used if `-Yretain-trees` is set. */
+  private[this] var myBody: List[Tree] = _
   /** The unpickled trees, and the source file they come from. */
   def body(implicit ctx: Context): List[Tree] = {
-    treeUnpickler.unpickle()
+    def computeBody() = treeUnpickler.unpickle()
+    if (ctx.settings.YretainTrees.value) {
+      if (myBody == null)
+        myBody = computeBody()
+      myBody
+    } else computeBody()
   }
 }

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

@@ -1341,6 +1341,9 @@ class Typer extends Namer with TypeAssigner with Applications with Implicits wit
     // check value class constraints
     checkDerivedValueClass(cls, body1)
 
+    if (ctx.settings.YretainTrees.value) {
+      cls.myTree = cdef1
+    }
     cdef1
 
     // todo later: check that