Store source file in TASTY attributes (#18948)

Add source file to TASTy attributes. This is a first step towards
removing the `@SourceFile` annotation.

#### Release notes
Tools that read TASTy need to know that the source file must be
unpickled from TASTy attributes.

Author: bishabosha

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

@@ -244,9 +244,6 @@ object Annotations {
         }
         else None
     }
-
-    def makeSourceFile(path: String, span: Span)(using Context): Annotation =
-      apply(defn.SourceFileAnnot, Literal(Constant(path)), span)
   }
 
   @sharable val EmptyAnnotation = Annotation(EmptyTree)

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/Symbols.scala

@@ -486,6 +486,12 @@ object Symbols extends SymUtils {
           mySource = ctx.getSource(file)
         else
           mySource = defn.patchSource(this)
+          if !mySource.exists then
+            val compUnitInfo = compilationUnitInfo
+            if compUnitInfo != null then
+              compUnitInfo.tastyInfo.flatMap(_.attributes.sourceFile) match
+                case Some(path) => mySource = ctx.getSource(path)
+                case _ =>
           if !mySource.exists then
             mySource = atPhaseNoLater(flattenPhase) {
               denot.topLevelClass.unforcedAnnotation(defn.SourceFileAnnot) match

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

@@ -3,7 +3,7 @@ package dotty.tools.dotc.core.tasty
 import dotty.tools.dotc.ast.{tpd, untpd}
 
 import dotty.tools.tasty.TastyBuffer
-import dotty.tools.tasty.TastyFormat, TastyFormat.AttributesSection
+import dotty.tools.tasty.TastyFormat.*
 
 object AttributePickler:
 
@@ -12,11 +12,26 @@ object AttributePickler:
     pickler: TastyPickler,
     buf: TastyBuffer
   ): Unit =
-    if attributes.booleanTags.nonEmpty then
-      pickler.newSection(AttributesSection, buf)
+    pickler.newSection(AttributesSection, buf)
+
+    var lastTag = -1
+    def assertTagOrder(tag: Int): Unit =
+      assert(tag != lastTag, s"duplicate attribute tag: $tag")
+      assert(tag > lastTag, s"attribute tags are not ordered: $tag after $lastTag")
+      lastTag = tag
 
     for tag <- attributes.booleanTags do
+      assert(isBooleanAttrTag(tag), "Not a boolean attribute tag: " + tag)
+      assertTagOrder(tag)
+      buf.writeByte(tag)
+
+    assert(attributes.stringTagValues.exists(_._1 == SOURCEFILEattr))
+    for (tag, value) <- attributes.stringTagValues do
+      assert(isStringAttrTag(tag), "Not a string attribute tag: " + tag)
+      assertTagOrder(tag)
+      val utf8Ref = pickler.nameBuffer.utf8Index(value)
       buf.writeByte(tag)
+      buf.writeNat(utf8Ref.index)
 
   end pickleAttributes
 

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

@@ -3,19 +3,36 @@ package core.tasty
 
 import scala.language.unsafeNulls
 import scala.collection.immutable.BitSet
+import scala.collection.immutable.TreeMap
 
-import dotty.tools.tasty.{TastyFormat, TastyReader, TastyBuffer}
+import dotty.tools.tasty.{TastyFormat, TastyReader, TastyBuffer}, TastyFormat.{isBooleanAttrTag, isStringAttrTag}
+import dotty.tools.dotc.core.tasty.TastyUnpickler.NameTable
 
-class AttributeUnpickler(reader: TastyReader):
+class AttributeUnpickler(reader: TastyReader, nameAtRef: NameTable):
   import reader._
 
   lazy val attributes: Attributes = {
     val booleanTags = BitSet.newBuilder
+    val stringTagValue = List.newBuilder[(Int, String)]
 
+    var lastTag = -1
     while !isAtEnd do
-      booleanTags += readByte()
+      val tag = readByte()
+      if isBooleanAttrTag(tag) then
+        booleanTags += tag
+      else if isStringAttrTag(tag) then
+        val utf8Ref = readNameRef()
+        val value = nameAtRef(utf8Ref).toString
+        stringTagValue += tag -> value
+      else
+        assert(false, "unknown attribute tag: " + tag)
 
-    new Attributes(booleanTags.result())
+      assert(tag != lastTag, s"duplicate attribute tag: $tag")
+      assert(tag > lastTag, s"attribute tags are not ordered: $tag after $lastTag")
+      lastTag = tag
+    end while
+
+    new Attributes(booleanTags.result(), stringTagValue.result())
   }
 
 end AttributeUnpickler

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

@@ -3,20 +3,24 @@ package dotty.tools.dotc.core.tasty
 import dotty.tools.tasty.TastyFormat.*
 
 import scala.collection.immutable.BitSet
+import scala.collection.immutable.TreeMap
 
 class Attributes private[tasty](
   private[tasty] val booleanTags: BitSet,
+  private[tasty] val stringTagValues: List[(Int, String)],
 ) {
   def scala2StandardLibrary: Boolean = booleanTags(SCALA2STANDARDLIBRARYattr)
   def explicitNulls: Boolean = booleanTags(EXPLICITNULLSattr)
   def captureChecked: Boolean = booleanTags(CAPTURECHECKEDattr)
   def withPureFuns: Boolean = booleanTags(WITHPUREFUNSattr)
   def isJava: Boolean = booleanTags(JAVAattr)
   def isOutline: Boolean = booleanTags(OUTLINEattr)
+  def sourceFile: Option[String] = stringTagValues.find(_._1 == SOURCEFILEattr).map(_._2)
 }
 
 object Attributes:
   def apply(
+    sourceFile: String,
     scala2StandardLibrary: Boolean,
     explicitNulls: Boolean,
     captureChecked: Boolean,
@@ -31,8 +35,12 @@ object Attributes:
     if withPureFuns then booleanTags += WITHPUREFUNSattr
     if isJava then booleanTags += JAVAattr
     if isOutline then booleanTags += OUTLINEattr
-    new Attributes(booleanTags.result())
+
+    val stringTagValues = List.newBuilder[(Int, String)]
+    stringTagValues += SOURCEFILEattr -> sourceFile
+
+    new Attributes(booleanTags.result(), stringTagValues.result())
   end apply
 
   val empty: Attributes =
-    new Attributes(BitSet.empty)
+    new Attributes(BitSet.empty, Nil)

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) {
@@ -39,26 +41,33 @@ object DottyUnpickler {
 
   class AttributesSectionUnpickler extends SectionUnpickler[AttributeUnpickler](AttributesSection) {
     def unpickle(reader: TastyReader, nameAtRef: NameTable): AttributeUnpickler =
-      new AttributeUnpickler(reader)
+      new AttributeUnpickler(reader, nameAtRef)
   }
 }
 
 /** 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/NameBuffer.scala

@@ -49,6 +49,10 @@ class NameBuffer extends TastyBuffer(10000) {
     }
   }
 
+  def utf8Index(value: String): NameRef =
+    import Decorators.toTermName
+    nameIndex(value.toTermName)
+
   private inline def withLength(inline op: Unit, lengthWidth: Int = 1): Unit = {
     val lengthAddr = currentAddr
     var i = 0

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

@@ -228,11 +228,17 @@ class TastyPrinter(bytes: Array[Byte]) {
     import dotty.tools.tasty.TastyFormat.*
     def unpickle(reader: TastyReader, tastyName: NameTable): Unit = {
       import reader.*
-      val attributes = new AttributeUnpickler(reader).attributes
       sb.append(s"\n\nAttributes (${reader.endAddr.index - reader.startAddr.index} bytes, starting from $base):\n")
-
-      for tag <- attributes.booleanTags do
-        sb.append("  ").append(attributeTagToString(tag)).append("\n")
+      while !isAtEnd do
+        val tag = readByte()
+        sb.append("  ").append(attributeTagToString(tag))
+        if isBooleanAttrTag(tag) then ()
+        else if isStringAttrTag(tag) then
+          val utf8Ref = readNameRef()
+          val value = nameAtRef(utf8Ref).toString
+          sb.append(nameStr(s" ${utf8Ref.index} [$value]"))
+        sb.append("\n")
+      sb.result
     }
   }
 

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

@@ -96,12 +96,16 @@ class Pickler extends Phase {
     do
       if ctx.settings.YtestPickler.value then beforePickling(cls) = tree.show
 
+      val sourceRelativePath =
+        val reference = ctx.settings.sourceroot.value
+        util.SourceFile.relativePath(unit.source, reference)
       val isJavaAttr = unit.isJava // we must always set JAVAattr when pickling Java sources
       if isJavaAttr then
         // assert that Java sources didn't reach Pickler without `-Yjava-tasty`.
         assert(ctx.settings.YjavaTasty.value, "unexpected Java source file without -Yjava-tasty")
       val isOutline = isJavaAttr // TODO: later we may want outline for Scala sources too
       val attributes = Attributes(
+        sourceFile = sourceRelativePath,
         scala2StandardLibrary = ctx.settings.YcompileScala2Library.value,
         explicitNulls = ctx.settings.YexplicitNulls.value,
         captureChecked = Feature.ccEnabled,
@@ -231,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/src/dotty/tools/dotc/transform/PostTyper.scala

@@ -417,12 +417,13 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
                   if illegalRefs.nonEmpty then
                     report.error(
                       em"The type of a class parent cannot refer to constructor parameters, but ${parent.tpe} refers to ${illegalRefs.map(_.name.show).mkString(",")}", parent.srcPos)
-            // Add SourceFile annotation to top-level classes
             if sym.owner.is(Package) then
+              // Add SourceFile annotation to top-level classes
+              // TODO remove this annotation once the reference compiler uses the TASTy source file attribute.
               if ctx.compilationUnit.source.exists && sym != defn.SourceFileAnnot then
                 val reference = ctx.settings.sourceroot.value
                 val relativePath = util.SourceFile.relativePath(ctx.compilationUnit.source, reference)
-                sym.addAnnotation(Annotation.makeSourceFile(relativePath, tree.span))
+                sym.addAnnotation(Annotation(defn.SourceFileAnnot, Literal(Constants.Constant(relativePath)), tree.span))
           else
             if !sym.is(Param) && !sym.owner.isOneOf(AbstractOrTrait) then
               Checking.checkGoodBounds(tree.symbol)