Store source file in TASTY attributes

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

@@ -471,6 +471,7 @@ object Symbols {
         else
           mySource = defn.patchSource(this)
           if !mySource.exists then
+            // TODO try to get source from TASTY attributes, otherwise load from annotation
             mySource = atPhaseNoLater(flattenPhase) {
               denot.topLevelClass.unforcedAnnotation(defn.SourceFileAnnot) match
                 case Some(sourceAnnot) => sourceAnnot.argumentConstant(0) match

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

@@ -14,12 +14,14 @@ object AttributePickler:
     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

@@ -10,24 +10,18 @@ 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/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

@@ -107,6 +107,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

@@ -418,6 +418,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*
@@ -270,8 +270,16 @@ Standard Section: "Comments" Comment*
 
 Standard Section: "Attributes" Attribute*
 ```none
-  Attribute     = SCALA2STANDARDLIBRARYattr
-                  EXPLICITNULLSattr
+  Attribute     = SOURCEFILEattr UTF8
+                  SOURCEFILEattr
+                  SCALA2STANDARDLIBRARYatt2
+                  EXPLICITNULLSatt3
+```
+
+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 +444,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 +493,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 +513,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 +531,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 +544,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 +607,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 +844,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.

tasty/src/dotty/tools/tasty/TastyReader.scala

@@ -3,6 +3,7 @@ package dotty.tools.tasty
 import collection.mutable
 
 import TastyBuffer._
+import java.nio.charset.StandardCharsets
 
 /** A byte array buffer that can be filled with bytes or natural numbers in TASTY format,
  *  and that supports reading and patching addresses represented as natural numbers.
@@ -104,6 +105,13 @@ class TastyReader(val bytes: Array[Byte], start: Int, end: Int, val base: Int =
     x
   }
 
+  /** Read a UTF8 string encoded as `Length UTF8-CodePoint*` */
+  def readUTF8(): String = {
+    val length = readNat()
+    if (length == 0) ""
+    else new String(readBytes(length), StandardCharsets.UTF_8)
+  }
+
   /** Read a natural number and return as a NameRef */
   def readNameRef(): NameRef = NameRef(readNat())