Fix missing CompilationUnit on class symbol loaded in TreeUnpickler

Also refactor accesses to the tasty attributes to get them directly from
the compilation unit info instead of the unpickler.

Author: nicolasstucki

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

@@ -420,20 +420,9 @@ class TastyLoader(val tastyFile: AbstractFile) extends SymbolLoader {
   private val unpickler: tasty.DottyUnpickler =
     handleUnpicklingExceptions:
       val tastyBytes = tastyFile.toByteArray
-      new tasty.DottyUnpickler(tastyBytes) // reads header and name table
-
-  val compilationUnitInfo: CompilationUnitInfo | Null =
-    val tastyHeader = unpickler.unpickler.header
-    val tastyVersion = TastyVersion(
-      tastyHeader.majorVersion,
-      tastyHeader.minorVersion,
-      tastyHeader.experimentalVersion,
-    )
-    val attributes = unpickler.tastyAttributes
-    new CompilationUnitInfo(
-      tastyFile,
-      tastyInfo = Some(TastyInfo(tastyVersion, attributes)),
-    )
+      new tasty.DottyUnpickler(tastyFile, tastyBytes) // reads header and name table
+
+  val compilationUnitInfo: CompilationUnitInfo | Null = unpickler.compilationUnitInfo
 
   def description(using Context): String = "TASTy file " + tastyFile.toString
 

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

@@ -16,15 +16,17 @@ import dotty.tools.tasty.TastyReader
 import dotty.tools.tasty.TastyFormat.{ASTsSection, PositionsSection, CommentsSection, AttributesSection}
 import dotty.tools.tasty.TastyVersion
 
+import dotty.tools.io.AbstractFile
+
 object DottyUnpickler {
 
   /** Exception thrown if classfile is corrupted */
   class BadSignature(msg: String) extends RuntimeException(msg)
 
-  class TreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], commentUnpickler: Option[CommentUnpickler], attributeUnpickler: Option[AttributeUnpickler])
+  class TreeSectionUnpickler(compilationUnitInfo: CompilationUnitInfo, posUnpickler: Option[PositionUnpickler], commentUnpickler: Option[CommentUnpickler])
   extends SectionUnpickler[TreeUnpickler](ASTsSection) {
     def unpickle(reader: TastyReader, nameAtRef: NameTable): TreeUnpickler =
-      new TreeUnpickler(reader, nameAtRef, posUnpickler, commentUnpickler, attributeUnpickler)
+      new TreeUnpickler(reader, nameAtRef, compilationUnitInfo, posUnpickler, commentUnpickler)
   }
 
   class PositionsSectionUnpickler extends SectionUnpickler[PositionUnpickler](PositionsSection) {
@@ -44,21 +46,28 @@ object DottyUnpickler {
 }
 
 /** A class for unpickling Tasty trees and symbols.
+ *  @param tastyFile     tasty file from which we unpickle (used for CompilationUnitInfo)
  *  @param bytes         the bytearray containing the Tasty file from which we unpickle
  *  @param mode          the tasty file contains package (TopLevel), an expression (Term) or a type (TypeTree)
  */
-class DottyUnpickler(bytes: Array[Byte], mode: UnpickleMode = UnpickleMode.TopLevel) extends ClassfileParser.Embedded with tpd.TreeProvider {
+class DottyUnpickler(tastyFile: AbstractFile, bytes: Array[Byte], mode: UnpickleMode = UnpickleMode.TopLevel) extends ClassfileParser.Embedded with tpd.TreeProvider {
   import tpd.*
   import DottyUnpickler.*
 
   val unpickler: TastyUnpickler = new TastyUnpickler(bytes)
+
+  val tastyAttributes: Attributes =
+    unpickler.unpickle(new AttributesSectionUnpickler)
+      .map(_.attributes).getOrElse(Attributes.empty)
+  val compilationUnitInfo: CompilationUnitInfo =
+    import unpickler.header.{majorVersion, minorVersion, experimentalVersion}
+    val tastyVersion = TastyVersion(majorVersion, minorVersion, experimentalVersion)
+    val tastyInfo = TastyInfo(tastyVersion, tastyAttributes)
+    new CompilationUnitInfo(tastyFile, Some(tastyInfo))
+
   private val posUnpicklerOpt = unpickler.unpickle(new PositionsSectionUnpickler)
   private val commentUnpicklerOpt = unpickler.unpickle(new CommentsSectionUnpickler)
-  private val attributeUnpicklerOpt = unpickler.unpickle(new AttributesSectionUnpickler)
-  private val treeUnpickler = unpickler.unpickle(treeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt, attributeUnpicklerOpt)).get
-
-  def tastyAttributes: Attributes =
-    attributeUnpicklerOpt.map(_.attributes).getOrElse(Attributes.empty)
+  private val treeUnpickler = unpickler.unpickle(treeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt)).get
 
   /** Enter all toplevel classes and objects into their scopes
    *  @param roots          a set of SymDenotations that should be overwritten by unpickling
@@ -69,9 +78,8 @@ class DottyUnpickler(bytes: Array[Byte], mode: UnpickleMode = UnpickleMode.TopLe
   protected def treeSectionUnpickler(
     posUnpicklerOpt: Option[PositionUnpickler],
     commentUnpicklerOpt: Option[CommentUnpickler],
-    attributeUnpicklerOpt: Option[AttributeUnpickler]
   ): TreeSectionUnpickler =
-    new TreeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt, attributeUnpicklerOpt)
+    new TreeSectionUnpickler(compilationUnitInfo, posUnpicklerOpt, commentUnpicklerOpt)
 
   protected def computeRootTrees(using Context): List[Tree] = treeUnpickler.unpickle(mode)
 

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

@@ -49,15 +49,16 @@ import scala.compiletime.uninitialized
 
 /** Unpickler for typed trees
  *  @param reader              the reader from which to unpickle
+ *  @param compilationUnitInfo  the compilation unit info of the TASTy
  *  @param posUnpicklerOpt     the unpickler for positions, if it exists
  *  @param commentUnpicklerOpt the unpickler for comments, if it exists
  *  @param attributeUnpicklerOpt the unpickler for attributes, if it exists
  */
 class TreeUnpickler(reader: TastyReader,
                     nameAtRef: NameTable,
+                    compilationUnitInfo: CompilationUnitInfo,
                     posUnpicklerOpt: Option[PositionUnpickler],
-                    commentUnpicklerOpt: Option[CommentUnpickler],
-                    attributeUnpicklerOpt: Option[AttributeUnpickler]) {
+                    commentUnpicklerOpt: Option[CommentUnpickler]) {
   import TreeUnpickler.*
   import tpd.*
 
@@ -92,22 +93,21 @@ class TreeUnpickler(reader: TastyReader,
   /** The root owner tree. See `OwnerTree` class definition. Set by `enterTopLevel`. */
   private var ownerTree: OwnerTree = uninitialized
 
+  /** TASTy attributes */
+  private val attributes: Attributes = compilationUnitInfo.tastyInfo.get.attributes
+
   /** Was unpickled class compiled with capture checks? */
-  private val withCaptureChecks: Boolean =
-    attributeUnpicklerOpt.exists(_.attributes.captureChecked)
+  private val withCaptureChecks: Boolean = attributes.captureChecked
 
-  private val unpicklingScala2Library =
-    attributeUnpicklerOpt.exists(_.attributes.scala2StandardLibrary)
+  private val unpicklingScala2Library = attributes.scala2StandardLibrary
 
   /** This dependency was compiled with explicit nulls enabled */
   // TODO Use this to tag the symbols of this dependency as compiled with explicit nulls (see use of unpicklingScala2Library).
-  private val explicitNulls =
-    attributeUnpicklerOpt.exists(_.attributes.explicitNulls)
+  private val explicitNulls = attributes.explicitNulls
 
-  private val unpicklingJava =
-    attributeUnpicklerOpt.exists(_.attributes.isJava)
+  private val unpicklingJava = attributes.isJava
 
-  private val isOutline = attributeUnpicklerOpt.exists(_.attributes.isOutline)
+  private val isOutline = attributes.isOutline
 
   private def registerSym(addr: Addr, sym: Symbol) =
     symAtAddr(addr) = sym
@@ -636,8 +636,8 @@ class TreeUnpickler(reader: TastyReader,
             rootd.symbol
           case _ =>
             val completer = adjustIfModule(new Completer(subReader(start, end)))
-            if (isClass)
-              newClassSymbol(ctx.owner, name.asTypeName, flags, completer, privateWithin, coord)
+            if isClass then
+              newClassSymbol(ctx.owner, name.asTypeName, flags, completer, privateWithin, coord, compilationUnitInfo)
             else
               newSymbol(ctx.owner, name, flags, completer, privateWithin, coord)
         }

compiler/src/dotty/tools/dotc/fromtasty/ReadTasty.scala

@@ -11,6 +11,7 @@ import Denotations.staticRef
 import NameOps.*
 import ast.Trees.Tree
 import Phases.Phase
+import core.tasty.Attributes
 
 /** Load trees from TASTY files */
 class ReadTasty extends Phase {

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

@@ -21,6 +21,7 @@ import scala.quoted.runtime.impl.*
 import scala.collection.mutable
 
 import QuoteUtils.*
+import dotty.tools.io.NoAbstractFile
 
 object PickledQuotes {
   import tpd.*
@@ -268,7 +269,7 @@ object PickledQuotes {
           quotePickling.println(s"**** unpickling quote from TASTY\n${TastyPrinter.showContents(bytes, ctx.settings.color.value == "never")}")
 
           val mode = if (isType) UnpickleMode.TypeTree else UnpickleMode.Term
-          val unpickler = new DottyUnpickler(bytes, mode)
+          val unpickler = new DottyUnpickler(NoAbstractFile, bytes, mode)
           unpickler.enter(Set.empty)
 
           val tree = unpickler.tree

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

@@ -235,7 +235,7 @@ class Pickler extends Phase {
     ctx.initialize()
     val unpicklers =
       for ((cls, (unit, bytes)) <- pickledBytes) yield {
-        val unpickler = new DottyUnpickler(bytes)
+        val unpickler = new DottyUnpickler(unit.source.file, bytes)
         unpickler.enter(roots = Set.empty)
         cls -> (unit, unpickler)
       }

compiler/test/dotty/tools/dotc/core/tasty/CommentPicklingTest.scala

@@ -18,6 +18,7 @@ import dotty.tools.vulpix.TestConfiguration
 
 import org.junit.Test
 import org.junit.Assert.{assertEquals, assertFalse, fail}
+import dotty.tools.io.AbstractFile
 
 class CommentPicklingTest {
 
@@ -116,7 +117,7 @@ class CommentPicklingTest {
       implicit val ctx: Context = setup(args, initCtx).map(_._2).getOrElse(initCtx)
       ctx.initialize()
       val trees = files.flatMap { f =>
-        val unpickler = new DottyUnpickler(f.toByteArray())
+        val unpickler = new DottyUnpickler(AbstractFile.getFile(f.jpath), f.toByteArray())
         unpickler.enter(roots = Set.empty)
         unpickler.rootTrees(using ctx)
       }