Store source file in TASTY attributes

nicolasstucki · nicolasstucki · commit 6a8eafffa59d · 2023-11-24T12:45:52.000+01:00
diff --git a/compiler/src/dotty/tools/dotc/core/CompilationUnitInfo.scala b/compiler/src/dotty/tools/dotc/core/CompilationUnitInfo.scala
@@ -14,7 +14,8 @@ import dotty.tools.tasty.TastyVersion
 class CompilationUnitInfo(
   val associatedFile: AbstractFile,
   val tastyVersion: Option[TastyVersion],
-  val explicitNulls: Boolean
+  val explicitNulls: Boolean,
+  val sourceFile: Option[String],
 ) {
 
   override def toString(): String =
@@ -28,4 +29,5 @@ object CompilationUnitInfo:
       assocFile,
       tastyVersion = None,
       explicitNulls = explicitNulls,
+      sourceFile = if assocFile.path.endsWith(".scala") then Some(assocFile.path) else None,
     )
diff --git a/compiler/src/dotty/tools/dotc/core/SymbolLoaders.scala b/compiler/src/dotty/tools/dotc/core/SymbolLoaders.scala
@@ -434,6 +434,7 @@ class TastyLoader(val tastyFile: AbstractFile) extends SymbolLoader {
       tastyFile,
       tastyVersion = Some(tastyVersion),
       explicitNulls = attributes.explicitNulls,
+      sourceFile = attributes.sourceFile,
     )
 
   def description(using Context): String = "TASTy file " + tastyFile.toString
diff --git a/compiler/src/dotty/tools/dotc/core/Symbols.scala b/compiler/src/dotty/tools/dotc/core/Symbols.scala
@@ -293,6 +293,22 @@ object Symbols {
       val compUnitInfo = compilationUnitInfo
       compUnitInfo != null && compUnitInfo.explicitNulls
 
+    /** The version of TASTy from which the symbol was loaded, None if not applicable. */
+    def sourceFile(using Context): Option[String] =
+      val compUnitInfo = compilationUnitInfo
+      if compUnitInfo == null then
+        // Try to get it form the 3.0-3.3 @SourceFile annotation
+        atPhaseNoLater(flattenPhase) {
+          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
+        }
+      else
+        // Loaded from TASTy attributes for 3.4+
+        compUnitInfo.sourceFile
+
     /** The class file from which this class was generated, null if not applicable. */
     final def binaryFile(using Context): AbstractFile | Null = {
       val file = associatedFile
@@ -493,13 +509,7 @@ object Symbols {
         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
     }
 
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/AttributePickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/AttributePickler.scala
@@ -12,11 +12,14 @@ object AttributePickler:
     pickler: TastyPickler,
     buf: TastyBuffer
   ): Unit =
-    if attributes.scala2StandardLibrary || attributes.explicitNulls then // or any other attribute is set
-      pickler.newSection(AttributesSection, buf)
+    pickler.newSection(AttributesSection, buf)
 
-      for tag <- attributes.booleanTags do
-        buf.writeByte(tag)
+    for tag <- attributes.booleanTags do
+      buf.writeByte(tag)
+
+    for (tag, value) <- attributes.stringTagValues do
+      buf.writeByte(tag)
+      buf.writeUTF8(value)
 
   end pickleAttributes
 
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/AttributeUnpickler.scala b/compiler/src/dotty/tools/dotc/core/tasty/AttributeUnpickler.scala
@@ -10,11 +10,16 @@ class AttributeUnpickler(reader: TastyReader):
 
   lazy val attributes: Attributes = {
     val booleanTags = List.newBuilder[Int]
+    val stringTagValue = List.newBuilder[(Int, String)]
 
     while !isAtEnd do
-      booleanTags += readByte()
+      val tag = readByte()
+      if tag < TastyFormat.firstStringAttrTag then
+        booleanTags += tag
+      else
+        stringTagValue += ((tag, readUTF8()))
 
-    new Attributes(booleanTags.result())
+    new Attributes(booleanTags.result(), stringTagValue.result())
   }
 
 end AttributeUnpickler
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/Attributes.scala b/compiler/src/dotty/tools/dotc/core/tasty/Attributes.scala
@@ -4,20 +4,29 @@ import dotty.tools.tasty.TastyFormat
 
 class Attributes(
   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
-    new Attributes(booleanTags.result())
+
+    val stringTagValues = List(
+      (TastyFormat.SOURCEFILEattr, sourceFile)
+    )
+
+    new Attributes(booleanTags.result(), stringTagValues)
   end apply
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/CommentPickler.scala b/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
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/CommentUnpickler.scala b/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)
       }
     }
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/DottyUnpickler.scala b/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(new Attributes(booleanTags = Nil))
+    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
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TastyPrinter.scala b/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:
 
@@ -237,6 +238,9 @@ class TastyPrinter(bytes: Array[Byte]) {
 
       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
     }
diff --git a/compiler/src/dotty/tools/dotc/core/tasty/TreeUnpickler.scala b/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
 
diff --git a/compiler/src/dotty/tools/dotc/transform/Pickler.scala b/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,
           )
diff --git a/compiler/src/dotty/tools/dotc/transform/PostTyper.scala b/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)
diff --git a/tasty/src/dotty/tools/tasty/TastyBuffer.scala b/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
diff --git a/tasty/src/dotty/tools/tasty/TastyFormat.scala b/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.
diff --git a/tasty/src/dotty/tools/tasty/TastyReader.scala b/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())
 

Original file line number	Diff line number	Diff line change
`@@ -434,6 +434,7 @@ class TastyLoader(val tastyFile: AbstractFile) extends SymbolLoader {`
`434`	`434`	`tastyFile,`
`435`	`435`	`tastyVersion = Some(tastyVersion),`
`436`	`436`	`explicitNulls = attributes.explicitNulls,`
	`437`	`+ sourceFile = attributes.sourceFile,`
`437`	`438`	`)`
`438`	`439`
`439`	`440`	`def description(using Context): String = "TASTy file " + tastyFile.toString`
Original file line number	Diff line number	Diff line change
`@@ -19,11 +19,9 @@ class CommentUnpickler(reader: TastyReader) {`
`19`	`19`	`val comments = new HashMap[Addr, Comment]`
`20`	`20`	`while (!isAtEnd) {`
`21`	`21`	`val addr = readAddr()`
`22`		`- val length = readNat()`
`23`		`- if (length > 0) {`
`24`		`- val bytes = readBytes(length)`
	`22`	`+ val rawComment = readUTF8()`
	`23`	`+ if (rawComment != "") {`
`25`	`24`	`val position = new Span(readLongInt())`
`26`		`- val rawComment = new String(bytes, StandardCharsets.UTF_8)`
`27`	`25`	`comments(addr) = Comment(position, rawComment)`
`28`	`26`	`}`
`29`	`27`	`}`