Unify tree loaders

Author: nicolasstucki

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

@@ -99,10 +99,11 @@ object FromTasty extends Driver {
 
   def loadCompilationUnit(clsd: ClassDenotation)(implicit ctx: Context): Option[CompilationUnit] = {
     assert(ctx.settings.XlinkOptimise.value)
-    clsd.dottyUnpickler.flatMap { unpickler =>
+    val tree = clsd.symbol.asClass.unitTree
+    if (tree.isEmpty) None
+    else {
       ctx.log("Loading compilation unit for: " + clsd)
-      val body = unpickler.body(ctx.addMode(Mode.ReadPositions))
-      body.headOption.map(unpickled => mkUnit(clsd, unpickled))
+      Some(mkUnit(clsd, tree))
     }
   }
 

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

@@ -19,8 +19,6 @@ import java.util.WeakHashMap
 import config.Config
 import config.Printers.{completions, incremental, noPrinter}
 
-import dotty.tools.dotc.core.tasty.DottyUnpickler
-
 trait SymDenotations { this: Context =>
   import SymDenotations._
 
@@ -141,8 +139,6 @@ object SymDenotations {
     private[this] var myPrivateWithin: Symbol = initPrivateWithin
     private[this] var myAnnotations: List[Annotation] = Nil
 
-    private[dotc] var dottyUnpickler: Option[DottyUnpickler] = None
-
     /** The owner of the symbol; overridden in NoDenotation */
     def owner: Symbol = ownerIfExists
 

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

@@ -323,11 +323,15 @@ class ClassfileLoader(val classfile: AbstractFile) extends SymbolLoader {
 
   def load(root: SymDenotation)(implicit ctx: Context): Unit = {
     val (classRoot, moduleRoot) = rootDenots(root.asClass)
-    (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 _ =>
+    val classfileParser = new ClassfileParser(classfile, classRoot, moduleRoot)(ctx)
+    val result = classfileParser.run()
+    if (ctx.settings.YretainTrees.value || ctx.settings.XlinkOptimise.value) {
+      result match {
+        case Some(unpickler: tasty.DottyUnpickler) =>
+          classRoot.symbol.asClass.unpickler = unpickler
+          moduleRoot.symbol.asClass.unpickler = unpickler
+        case _ =>
+      }
     }
   }
 }

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

@@ -558,32 +558,44 @@ object Symbols {
     /** If this is a top-level class, and if `-Yretain-trees` is set, return the TypeDef tree
      *  for this class, otherwise EmptyTree. This will force the info of the class.
      */
-    def tree(implicit ctx: Context): tpd.Tree /* tpd.TypeDef | tpd.EmptyTree */ = {
+    def tree(implicit ctx: Context): tpd.Tree /* tpd.TypeDef | tpd.EmptyTree */  = {
+      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 t = unitTree
+      if (t.isEmpty) t
+      else findTree(tree).get
+    }
+
+    /** If this is a top-level class, and if `-Yretain-trees` or `-YlinkOptimise` is set,
+      * return the PackageDef tree for this class, otherwise EmptyTree.
+      * This will force the info of the class.
+      */
+    def unitTree(implicit ctx: Context): tpd.Tree /* tpd.PackageDef | tpd.TypeDef | tpd.EmptyTree */ = {
       denot.info
       // 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))
+        val body = unpickler.body(ctx.addMode(Mode.ReadPositions))
+        myTree = body.headOption.getOrElse(tpd.EmptyTree)
         unpickler = null
-
-        myTree = findTree(unpickledTree).get
       }
       myTree
     }
-    private[dotc] var myTree: tpd.Tree = tpd.EmptyTree
+    private var myTree: tpd.Tree /* tpd.PackageDef | tpd.TypeDef | tpd.EmptyTree */ = tpd.EmptyTree
     private[dotc] var unpickler: tasty.DottyUnpickler = _
 
+    private[dotc] def registerTree(tree: tpd.TypeDef)(implicit ctx: Context): Unit = {
+      if (ctx.settings.YretainTrees.value)
+        myTree = tree
+    }
+
     /** 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/classfile/ClassfileParser.scala

@@ -15,7 +15,6 @@ import scala.annotation.switch
 import typer.Checking.checkNonCyclic
 import io.AbstractFile
 import scala.util.control.NonFatal
-import tasty.DottyUnpickler
 
 object ClassfileParser {
   /** Marker trait for unpicklers that can be embedded in classfiles. */
@@ -163,16 +162,6 @@ class ClassfileParser(
       }
     }
 
-    if (ctx.settings.XlinkOptimise.value) {
-      // Save references to DottyUnpickler to be able to load compilation units from tasty
-      result match {
-        case result@Some(_: DottyUnpickler) =>
-          classRoot.dottyUnpickler = result.asInstanceOf[Option[DottyUnpickler]]
-          moduleRoot.dottyUnpickler = result.asInstanceOf[Option[DottyUnpickler]]
-        case _ =>
-      }
-    }
-
     // eager load java enum definitions for exhaustivity check of pattern match
     if (isEnum) {
       instanceScope.toList.map(_.ensureCompleted())

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

@@ -1429,9 +1429,8 @@ 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
-      }
+      cls.registerTree(cdef1)
+
       cdef1
 
       // todo later: check that