Join logic for unpickling top level and term tasty

Author: nicolasstucki

compiler/src/dotty/tools/dotc/core/classfile/ClassfileParser.scala

@@ -758,7 +758,7 @@ class ClassfileParser(
 
       def unpickleTASTY(bytes: Array[Byte]): Some[Embedded]  = {
         val unpickler = new tasty.DottyUnpickler(bytes)
-        unpickler.enter(roots = Set(classRoot, moduleRoot, moduleRoot.sourceModule))
+        unpickler.enter(roots = Set(classRoot, moduleRoot, moduleRoot.sourceModule), isTopLevel = true)
         Some(unpickler)
       }
 

compiler/src/dotty/tools/dotc/core/quoted/PickledQuotes.scala

@@ -110,7 +110,8 @@ object PickledQuotes {
   /** Unpickle TASTY bytes into it's tree */
   private def unpickle(bytes: Array[Byte], splices: Seq[Any])(implicit ctx: Context): Tree = {
     val unpickler = new TastyUnpickler(bytes, splices)
-    val tree = unpickler.unpickleExpr()
+    unpickler.enter(Set(ctx.owner), isTopLevel = false)
+    val tree = unpickler.tree
     if (pickling ne noPrinter) {
       println(i"**** unpickled quote for \n${tree.show}")
       new TastyPrinter(bytes).printContents()

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

@@ -43,11 +43,8 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded with t
   /** Enter all toplevel classes and objects into their scopes
    *  @param roots          a set of SymDenotations that should be overwritten by unpickling
    */
-  def enter(roots: Set[SymDenotation])(implicit ctx: Context): Unit =
-    treeUnpickler.enterTopLevel(roots)
-
-  def unpickleExpr()(implicit ctx: Context): Tree =
-    treeUnpickler.unpickleExpr()
+  def enter(roots: Set[SymDenotation], isTopLevel: Boolean)(implicit ctx: Context): Unit =
+    treeUnpickler.enter(roots, isTopLevel)
 
   protected def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler]): TreeSectionUnpickler = {
     new TreeSectionUnpickler(posUnpicklerOpt)

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

@@ -52,6 +52,9 @@ class TreeUnpickler(reader: TastyReader,
    */
   private val typeAtAddr = new mutable.HashMap[Addr, Type]
 
+  /** If the Tasty corresponds to a top level tree. `flase` is used used for quotes */
+  private[this] var isTopLevel: Boolean = true
+
   /** The root symbol denotation which are defined by the Tasty file associated with this
    *  TreeUnpickler. Set by `enterTopLevel`.
    */
@@ -71,24 +74,21 @@ class TreeUnpickler(reader: TastyReader,
   /** Enter all toplevel classes and objects into their scopes
    *  @param roots          a set of SymDenotations that should be overwritten by unpickling
    */
-  def enterTopLevel(roots: Set[SymDenotation])(implicit ctx: Context): Unit = {
+  def enter(roots: Set[SymDenotation], isTopLevel: Boolean)(implicit ctx: Context): Unit = {
+    this.isTopLevel = isTopLevel
     this.roots = roots
     val rdr = new TreeReader(reader).fork
     ownerTree = new OwnerTree(NoAddr, 0, rdr.fork, reader.endAddr)
-    rdr.indexStats(reader.endAddr)
-  }
-
-  def unpickleExpr()(implicit ctx: Context): Tree = {
-    this.roots = Set(ctx.owner)
-    val rdr = new TreeReader(reader).fork
-    ownerTree = new OwnerTree(NoAddr, 0, rdr.fork, reader.endAddr)
-    rdr.readTerm()
+    if (isTopLevel)
+      rdr.indexStats(reader.endAddr)
   }
 
   /** The unpickled trees */
   def unpickle()(implicit ctx: Context): List[Tree] = {
     assert(roots != null, "unpickle without previous enterTopLevel")
-    new TreeReader(reader).readTopLevel()
+    val rdr = new TreeReader(reader)
+    if (isTopLevel) rdr.readTopLevel()
+    else rdr.readTerm() :: Nil
   }
 
   class Completer(owner: Symbol, reader: TastyReader) extends LazyType {

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

@@ -94,7 +94,7 @@ class Pickler extends Phase {
     val unpicklers =
       for ((cls, pickler) <- picklers) yield {
         val unpickler = new DottyUnpickler(pickler.assembleParts())
-        unpickler.enter(roots = Set())
+        unpickler.enter(roots = Set(), isTopLevel = true)
         cls -> unpickler
       }
     pickling.println("************* entered toplevel ***********")