Join logic for unpickling top level and term tasty

Author: nicolasstucki

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

@@ -99,15 +99,14 @@ object PickledQuotes {
 
   /** Unpickle TASTY bytes into it's tree */
   private def unpickle(bytes: Array[Byte], splices: Seq[Any], isType: Boolean)(implicit ctx: Context): Tree = {
-    val unpickler = new TastyUnpickler(bytes, splices)
     if (pickling ne noPrinter) {
       println(i"**** unpickling quote from TASTY")
       new TastyPrinter(bytes).printContents()
     }
 
-    val tree =
-      if (isType) unpickler.unpickleTypeTree()
-      else unpickler.unpickleExpr()
+    val unpickler = new TastyUnpickler(bytes, splices, isType)
+    unpickler.enter(Set.empty)
+    val tree = unpickler.tree
 
     if (pickling ne noPrinter)
       println(i"**** unpickle quote ${tree.show}")

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

@@ -15,10 +15,13 @@ object TastyUnpickler {
  *  @param bytes         the bytearray containing the Tasty file from which we unpickle
  *  @param splices       splices that will fill the holes in the quote
  */
-class TastyUnpickler(bytes: Array[Byte], splices: Seq[Any]) extends DottyUnpickler(bytes) {
+class TastyUnpickler(bytes: Array[Byte], splices: Seq[Any], isTypeTree: Boolean) extends DottyUnpickler(bytes) {
   import DottyUnpickler._
   import TastyUnpickler._
 
+  override protected def mode: TreeUnpickler.UnpickleMode =
+    if (isTypeTree) TreeUnpickler.UnpickleMode.TypeTree else TreeUnpickler.UnpickleMode.Term
+
   protected override def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler =
     new QuotedTreeSectionUnpickler(posUnpicklerOpt, splices)
 }

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

@@ -46,14 +46,13 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded with t
   private val commentUnpicklerOpt = unpickler.unpickle(new CommentsSectionUnpickler)
   private val treeUnpickler = unpickler.unpickle(treeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt)).get
 
+  protected def mode: TreeUnpickler.UnpickleMode = TreeUnpickler.UnpickleMode.TopLevel
+
   /** 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()
+    treeUnpickler.enter(roots)
 
   def unpickleTypeTree()(implicit ctx: Context): Tree =
     treeUnpickler.unpickleTypeTree()
@@ -62,7 +61,7 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded with t
     new TreeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt)
   }
 
-  protected def computeTrees(implicit ctx: Context) = treeUnpickler.unpickle()
+  protected def computeTrees(implicit ctx: Context) = treeUnpickler.unpickle(mode)
 
   private[this] var ids: Array[String] = null
 

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

@@ -74,18 +74,12 @@ 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])(implicit ctx: Context): Unit = {
     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 (rdr.isTopLevel)
+      rdr.indexStats(reader.endAddr)
   }
 
   def unpickleTypeTree()(implicit ctx: Context): Tree = {
@@ -96,9 +90,14 @@ class TreeUnpickler(reader: TastyReader,
   }
 
   /** The unpickled trees */
-  def unpickle()(implicit ctx: Context): List[Tree] = {
+  def unpickle(mode: UnpickleMode)(implicit ctx: Context): List[Tree] = {
     assert(roots != null, "unpickle without previous enterTopLevel")
-    new TreeReader(reader).readTopLevel()
+    val rdr = new TreeReader(reader)
+    mode match {
+      case UnpickleMode.TopLevel => rdr.readTopLevel()
+      case UnpickleMode.Term => rdr.readTerm() :: Nil
+      case UnpickleMode.TypeTree => rdr.readTpt() :: Nil
+    }
   }
 
   class Completer(owner: Symbol, reader: TastyReader) extends LazyType {
@@ -885,21 +884,22 @@ class TreeUnpickler(reader: TastyReader,
     }
 
     def skipToplevel()(implicit ctx: Context): Unit= {
-      if (!isAtEnd)
-        nextByte match {
-          case IMPORT | PACKAGE =>
-            skipTree()
-            skipToplevel()
-          case _ =>
-        }
+      if (!isAtEnd && isTopLevel) {
+        skipTree()
+        skipToplevel()
+      }
     }
 
+    def isTopLevel(implicit ctx: Context): Boolean =
+      nextByte == IMPORT || nextByte == PACKAGE
+
     def readTopLevel()(implicit ctx: Context): List[Tree] = {
-      @tailrec def read(acc: ListBuffer[Tree]): List[Tree] = nextByte match {
-        case IMPORT | PACKAGE =>
+      @tailrec def read(acc: ListBuffer[Tree]): List[Tree] = {
+        if (isTopLevel) {
           acc += readIndexedStat(NoSymbol)
           if (!isAtEnd) read(acc) else acc.toList
-        case _ => // top-level trees which are not imports or packages are not part of tree
+        }
+        else // top-level trees which are not imports or packages are not part of tree
           acc.toList
       }
       read(new ListBuffer[tpd.Tree])
@@ -1298,6 +1298,13 @@ class TreeUnpickler(reader: TastyReader,
 
 object TreeUnpickler {
 
+  sealed trait UnpickleMode
+  object UnpickleMode {
+    object TopLevel extends UnpickleMode
+    object Term extends UnpickleMode
+    object TypeTree extends UnpickleMode
+  }
+
   /** A marker value used to detect cyclic reference while unpickling definitions. */
   @sharable val PoisonTree: tpd.Tree = Thicket(Nil)
 

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

@@ -98,7 +98,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.empty)
         cls -> unpickler
       }
     pickling.println("************* entered toplevel ***********")