Refactor TASTy Attributes (#19091)

We generalize the internal encoding of `Attributes` to be a list of
tags. Then we add add helper methods to have simpler ways to interact
with this abstraction using booleans. This implies that the
pickling/unpickling can be agnostic of the semantics of each tag.
Therefore reducing the number of places that need to be updated when we
add a new tag.

Useful for #19074, #19033, and #18948.

Author: nicolasstucki

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

@@ -287,6 +287,22 @@ object Symbols extends SymUtils {
       if compUnitInfo == null then None
       else compUnitInfo.tastyInfo
 
+    /** The source file from which the symbol was compiled, None if not applicable. */
+    def sourceFile(using Context): Option[String] =
+      def sourceFromSourceFileAnnot = atPhaseNoLater(flattenPhase) {
+        // Annotation present in sources compiled with 3.0-3.3
+        denot.topLevelClass.unforcedAnnotation(defn.SourceFileAnnot) match
+          case Some(sourceAnnot) => sourceAnnot.argumentConstant(0) match
+            case Some(Constant(path: String)) => Some(path)
+            case none => None
+          case none => None
+      }
+      val compUnitInfo = compilationUnitInfo
+      if compUnitInfo == null then sourceFromSourceFileAnnot
+      else compUnitInfo.tastyInfo match
+        case Some(tastyInfo) => tastyInfo.attributes.sourceFile.orElse(sourceFromSourceFileAnnot)
+        case _ => sourceFromSourceFileAnnot
+
     /** The class file from which this class was generated, null if not applicable. */
     final def binaryFile(using Context): AbstractFile | Null = {
       val file = associatedFile
@@ -487,13 +503,7 @@ object Symbols extends SymUtils {
         else
           mySource = defn.patchSource(this)
           if !mySource.exists then
-            mySource = atPhaseNoLater(flattenPhase) {
-              denot.topLevelClass.unforcedAnnotation(defn.SourceFileAnnot) match
-                case Some(sourceAnnot) => sourceAnnot.argumentConstant(0) match
-                  case Some(Constant(path: String)) => ctx.getSource(path)
-                  case none => NoSource
-                case none => NoSource
-            }
+            mySource = sourceFile.map(ctx.getSource).getOrElse(NoSource)
       mySource
     }
 

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

@@ -5,21 +5,21 @@ import dotty.tools.dotc.ast.{tpd, untpd}
 import dotty.tools.tasty.TastyBuffer
 import dotty.tools.tasty.TastyFormat, TastyFormat.AttributesSection
 
-import java.nio.charset.StandardCharsets
-
 object AttributePickler:
 
   def pickleAttributes(
     attributes: Attributes,
     pickler: TastyPickler,
     buf: TastyBuffer
   ): Unit =
-    if attributes.scala2StandardLibrary || attributes.explicitNulls then // or any other attribute is set
-      pickler.newSection(AttributesSection, buf)
-      // Pickle attributes
-      if attributes.scala2StandardLibrary then buf.writeNat(TastyFormat.SCALA2STANDARDLIBRARYattr)
-      if attributes.explicitNulls then buf.writeNat(TastyFormat.EXPLICITNULLSattr)
-    end if
+    pickler.newSection(AttributesSection, buf)
+
+    for tag <- attributes.booleanTags do
+      buf.writeByte(tag)
+
+    for (tag, value) <- attributes.stringTagValues do
+      buf.writeByte(tag)
+      buf.writeUtf8(value)
 
   end pickleAttributes
 

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

@@ -5,29 +5,21 @@ import scala.language.unsafeNulls
 
 import dotty.tools.tasty.{TastyFormat, TastyReader, TastyBuffer}
 
-import java.nio.charset.StandardCharsets
-
 class AttributeUnpickler(reader: TastyReader):
   import reader._
 
-  lazy val attributeTags: List[Int] =
-    val listBuilder = List.newBuilder[Int]
-    while !isAtEnd do listBuilder += readNat()
-    listBuilder.result()
-
   lazy val attributes: Attributes = {
-    var scala2StandardLibrary = false
-    var explicitNulls = false
-    for attributeTag <- attributeTags do
-      attributeTag match
-        case TastyFormat.SCALA2STANDARDLIBRARYattr => scala2StandardLibrary = true
-        case TastyFormat.EXPLICITNULLSattr => explicitNulls = true
-        case attribute =>
-          assert(false, "Unexpected attribute value: " + attribute)
-    Attributes(
-      scala2StandardLibrary,
-      explicitNulls,
-    )
+    val booleanTags = List.newBuilder[Int]
+    val stringTagValue = List.newBuilder[(Int, String)]
+
+    while !isAtEnd do
+      val tag = readByte()
+      if tag < TastyFormat.firstStringAttrTag then
+        booleanTags += tag
+      else
+        stringTagValue += ((tag, readUtf8()))
+
+    new Attributes(booleanTags.result(), stringTagValue.result())
   }
 
 end AttributeUnpickler

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

@@ -1,6 +1,32 @@
 package dotty.tools.dotc.core.tasty
 
+import dotty.tools.tasty.TastyFormat
+
 class Attributes(
-  val scala2StandardLibrary: Boolean,
-  val explicitNulls: Boolean,
-)
+  val booleanTags: List[Int],
+  val stringTagValues: List[(Int, String)],
+) {
+  def scala2StandardLibrary: Boolean =
+    booleanTags.contains(TastyFormat.SCALA2STANDARDLIBRARYattr)
+  def explicitNulls: Boolean =
+    booleanTags.contains(TastyFormat.EXPLICITNULLSattr)
+  def sourceFile: Option[String] =
+    stringTagValues.find(_._1 == TastyFormat.SOURCEFILEattr).map(_._2)
+}
+
+object Attributes:
+  def apply(
+    sourceFile: String,
+    scala2StandardLibrary: Boolean,
+    explicitNulls: Boolean,
+  ): Attributes =
+    val booleanTags = List.newBuilder[Int]
+    if scala2StandardLibrary then booleanTags += TastyFormat.SCALA2STANDARDLIBRARYattr
+    if explicitNulls then booleanTags += TastyFormat.EXPLICITNULLSattr
+
+    val stringTagValues = List(
+      (TastyFormat.SOURCEFILEattr, sourceFile)
+    )
+
+    new Attributes(booleanTags.result(), stringTagValues)
+  end apply

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

@@ -22,11 +22,8 @@ object CommentPickler:
 
     def pickleComment(addr: Addr, comment: Comment): Unit =
       if addr != NoAddr then
-        val bytes = comment.raw.getBytes(StandardCharsets.UTF_8).nn
-        val length = bytes.length
         buf.writeAddr(addr)
-        buf.writeNat(length)
-        buf.writeBytes(bytes, length)
+        buf.writeUtf8(comment.raw)
         buf.writeLongInt(comment.span.coords)
 
     def traverse(x: Any): Unit = x match

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

@@ -19,11 +19,9 @@ class CommentUnpickler(reader: TastyReader) {
     val comments = new HashMap[Addr, Comment]
     while (!isAtEnd) {
       val addr = readAddr()
-      val length = readNat()
-      if (length > 0) {
-        val bytes = readBytes(length)
+      val rawComment = readUtf8()
+      if (rawComment != "") {
         val position = new Span(readLongInt())
-        val rawComment = new String(bytes, StandardCharsets.UTF_8)
         comments(addr) = Comment(position, rawComment)
       }
     }

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

@@ -57,7 +57,7 @@ class DottyUnpickler(bytes: Array[Byte], mode: UnpickleMode = UnpickleMode.TopLe
   private val treeUnpickler = unpickler.unpickle(treeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt, attributeUnpicklerOpt)).get
 
   def tastyAttributes: Attributes =
-    attributeUnpicklerOpt.map(_.attributes).getOrElse(Attributes(false, false))
+    attributeUnpicklerOpt.map(_.attributes).getOrElse(new Attributes(booleanTags = Nil, stringTagValues = Nil))
 
   /** Enter all toplevel classes and objects into their scopes
    *  @param roots          a set of SymDenotations that should be overwritten by unpickling

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

@@ -12,6 +12,7 @@ import util.Spans.offsetToInt
 import dotty.tools.tasty.TastyFormat.{ASTsSection, PositionsSection, CommentsSection, AttributesSection}
 import java.nio.file.{Files, Paths}
 import dotty.tools.io.{JarArchive, Path}
+import java.nio.charset.StandardCharsets
 
 object TastyPrinter:
 
@@ -225,15 +226,22 @@ class TastyPrinter(bytes: Array[Byte]) {
   }
 
   class AttributesSectionUnpickler extends SectionUnpickler[String](AttributesSection) {
-  import dotty.tools.tasty.TastyFormat.attributeTagToString
+    import dotty.tools.tasty.TastyFormat.*
+
     private val sb: StringBuilder = new StringBuilder
 
     def unpickle(reader: TastyReader, tastyName: NameTable): String = {
+      import reader.*
       sb.append(s" ${reader.endAddr.index - reader.currentAddr.index}")
-      val attributeTags = new AttributeUnpickler(reader).attributeTags
+      val attributes = new AttributeUnpickler(reader).attributes
       sb.append(s"  attributes bytes:\n")
-      for attributeTag <- attributeTags do
-        sb.append("    ").append(attributeTagToString(attributeTag)).append("\n")
+
+      for tag <- attributes.booleanTags do
+        sb.append("   ").append(attributeTagToString(tag)).append("\n")
+      for (tag, value) <- attributes.stringTagValues do
+        sb.append("   ").append(attributeTagToString(tag))
+          .append(" ").append(value).append("\n")
+
       sb.result
     }
   }

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

@@ -106,6 +106,10 @@ class TreeUnpickler(reader: TastyReader,
   private val explicitNulls =
     attributeUnpicklerOpt.exists(_.attributes.explicitNulls)
 
+  /** Source file in the TASTy attributes */
+  private val sourceFile: Option[String] =
+    attributeUnpicklerOpt.flatMap(_.attributes.sourceFile)
+
   private def registerSym(addr: Addr, sym: Symbol) =
     symAtAddr(addr) = sym
 

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

@@ -108,7 +108,11 @@ class Pickler extends Phase {
                 pickler, treePkl.buf.addrOfTree, treePkl.docString, tree,
                 scratch.commentBuffer)
 
+          val sourceRelativePath =
+            val reference = ctx.settings.sourceroot.value
+            util.SourceFile.relativePath(unit.source, reference)
           val attributes = Attributes(
+            sourceFile = sourceRelativePath,
             scala2StandardLibrary = ctx.settings.YcompileScala2Library.value,
             explicitNulls = ctx.settings.YexplicitNulls.value,
           )

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

@@ -419,6 +419,7 @@ class PostTyper extends MacroTransform with InfoTransformer { thisPhase =>
                       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
+              // TODO do not add SourceFile annotation
               if ctx.compilationUnit.source.exists && sym != defn.SourceFileAnnot then
                 val reference = ctx.settings.sourceroot.value
                 val relativePath = util.SourceFile.relativePath(ctx.compilationUnit.source, reference)

tasty/src/dotty/tools/tasty/TastyBuffer.scala

@@ -1,6 +1,7 @@
 package dotty.tools.tasty
 
 import util.Util.dble
+import java.nio.charset.StandardCharsets
 
 object TastyBuffer {
 
@@ -115,6 +116,14 @@ class TastyBuffer(initialSize: Int) {
     writeBytes(bytes, 8)
   }
 
+  /** Write a UTF8 string encoded as `Length UTF8-CodePoint*` */
+  def writeUtf8(x: String): Unit = {
+    val bytes = x.getBytes(StandardCharsets.UTF_8)
+    val length = bytes.length
+    writeNat(length)
+    writeBytes(bytes, length)
+  }
+
   // -- Address handling --------------------------------------------
 
   /** Write natural number `x` right-adjusted in a field of `width` bytes

tasty/src/dotty/tools/tasty/TastyFormat.scala

@@ -234,11 +234,11 @@ Note: The signature of a SELECTin or TERMREFin node is the signature of the sele
 
 Note: Tree tags are grouped into 5 categories that determine what follows, and thus allow to compute the size of the tagged tree in a generic way.
 ```none
-  Category 1 (tags 1-59)   :  tag
-  Category 2 (tags 60-89)  :  tag Nat
-  Category 3 (tags 90-109) :  tag AST
-  Category 4 (tags 110-127):  tag Nat AST
-  Category 5 (tags 128-255):  tag Length <payload>
+  Tree Category 1 (tags 1-59)   :  tag
+  Tree Category 2 (tags 60-89)  :  tag Nat
+  Tree Category 3 (tags 90-109) :  tag AST
+  Tree Category 4 (tags 110-127):  tag Nat AST
+  Tree Category 5 (tags 128-255):  tag Length <payload>
 ```
 
 Standard-Section: "Positions" LinesSizes Assoc*
@@ -272,6 +272,13 @@ Standard Section: "Attributes" Attribute*
 ```none
   Attribute     = SCALA2STANDARDLIBRARYattr
                   EXPLICITNULLSattr
+                  SOURCEFILEattr UTF8
+```
+
+Note: Attribute tags are grouped into 2 categories that determine what follows, and thus allow to compute the size of the tagged tree in a generic way.
+```none
+  Attribute Category 1 (tags 1-127)  :  tag
+  Attribute Category 2 (tags 128-255):  tag UTF8
 ```
 
 **************************************************************************************/
@@ -436,9 +443,8 @@ object TastyFormat {
   final val SOURCE = 4
 
  // AST tags
-  // Cat. 1:    tag
+  // Tree Cat. 1:    tag
 
-  final val firstSimpleTreeTag = UNITconst
   // final val ??? = 1
   final val UNITconst = 2
   final val FALSEconst = 3
@@ -486,7 +492,7 @@ object TastyFormat {
   final val EMPTYCLAUSE = 45
   final val SPLITCLAUSE = 46
 
-  // Cat. 2:    tag Nat
+  // Tree  Cat. 2:    tag Nat
 
   final val SHAREDterm = 60
   final val SHAREDtype = 61
@@ -506,7 +512,7 @@ object TastyFormat {
   final val IMPORTED = 75
   final val RENAMED = 76
 
-  // Cat. 3:    tag AST
+  // Tree  Cat. 3:    tag AST
 
   final val THIS = 90
   final val QUALTHIS = 91
@@ -524,7 +530,7 @@ object TastyFormat {
   final val EXPLICITtpt = 103
 
 
-  // Cat. 4:    tag Nat AST
+  // Tree  Cat. 4:    tag Nat AST
 
   final val IDENT = 110
   final val IDENTtpt = 111
@@ -537,7 +543,7 @@ object TastyFormat {
   final val SELFDEF = 118
   final val NAMEDARG = 119
 
-  // Cat. 5:    tag Length ...
+  // Tree  Cat. 5:    tag Length ...
 
   final val PACKAGE = 128
   final val VALDEF = 129
@@ -600,17 +606,25 @@ object TastyFormat {
 
   final val HOLE = 255
 
+  final val firstSimpleTreeTag = UNITconst
   final val firstNatTreeTag = SHAREDterm
   final val firstASTTreeTag = THIS
   final val firstNatASTTreeTag = IDENT
   final val firstLengthTreeTag = PACKAGE
 
 
-  // Attributes tags
+// Attributes tags
 
+  // Attr Cat. 1: tag
   final val SCALA2STANDARDLIBRARYattr = 1
   final val EXPLICITNULLSattr = 2
 
+  // Attr Cat. 2: tag UTF8
+  final val SOURCEFILEattr = 128
+
+  final val firstBooleanAttrTag = SCALA2STANDARDLIBRARYattr
+  final val firstStringAttrTag = SOURCEFILEattr
+
   /** Useful for debugging */
   def isLegalTag(tag: Int): Boolean =
     firstSimpleTreeTag <= tag && tag <= SPLITCLAUSE ||
@@ -829,6 +843,7 @@ object TastyFormat {
   def attributeTagToString(tag: Int): String = tag match {
     case SCALA2STANDARDLIBRARYattr => "SCALA2STANDARDLIBRARYattr"
     case EXPLICITNULLSattr => "EXPLICITNULLSattr"
+    case SOURCEFILEattr => "SOURCEFILEattr"
   }
 
   /** @return If non-negative, the number of leading references (represented as nats) of a length/trees entry.