Merge pull request #10363 from dotty-staging/add-line-sizes-in-TASTy

Fix #6542: Pickle line sizes in TASTy

Author: nicolasstucki

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

@@ -32,7 +32,21 @@ class PositionPickler(
     (addrDelta << 3) | (toInt(hasStartDelta) << 2) | (toInt(hasEndDelta) << 1) | toInt(hasPoint)
   }
 
-  def picklePositions(roots: List[Tree], warnings: mutable.ListBuffer[String]): Unit = {
+  def picklePositions(source: SourceFile, roots: List[Tree], warnings: mutable.ListBuffer[String]): Unit = {
+    /** Pickle the number of lines followed by the length of each line */
+    def pickleLineOffsets(): Unit = {
+      val content = source.content()
+      buf.writeNat(content.count(_ == '\n') + 1) // number of lines
+      var lastIndex = content.indexOf('\n', 0)
+      buf.writeNat(lastIndex) // size of first line
+      while lastIndex != -1 do
+        val nextIndex = content.indexOf('\n', lastIndex + 1)
+        val end = if nextIndex != -1 then nextIndex else content.length
+        buf.writeNat(end - lastIndex - 1) // size of the next line
+        lastIndex = nextIndex
+    }
+    pickleLineOffsets()
+
     var lastIndex = 0
     var lastSpan = Span(0, 0)
     def pickleDeltas(index: Int, span: Span) = {

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

@@ -15,12 +15,20 @@ import Names.TermName
 class PositionUnpickler(reader: TastyReader, nameAtRef: NameRef => TermName) {
   import reader._
 
+  private var myLineSizes: Array[Int] = _
   private var mySpans: util.HashMap[Addr, Span] = _
   private var mySourcePaths: util.HashMap[Addr, String] = _
   private var isDefined = false
 
   def ensureDefined(): Unit = {
     if (!isDefined) {
+      val lines = readNat()
+      myLineSizes = new Array[Int](lines)
+      var i = 0
+      while i < lines do
+        myLineSizes(i) += readNat()
+        i += 1
+
       mySpans = util.HashMap[Addr, Span]()
       mySourcePaths = util.HashMap[Addr, String]()
       var curIndex = 0
@@ -60,6 +68,11 @@ class PositionUnpickler(reader: TastyReader, nameAtRef: NameRef => TermName) {
     mySourcePaths
   }
 
+  private[tasty] def lineSizes: Array[Int] = {
+    ensureDefined()
+    myLineSizes
+  }
+
   def spanAt(addr: Addr): Span = spans.getOrElse(addr, NoSpan)
   def sourcePathAt(addr: Addr): String = sourcePaths.getOrElse(addr, "")
 }

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

@@ -139,9 +139,14 @@ class TastyPrinter(bytes: Array[Byte]) {
     private val sb: StringBuilder = new StringBuilder
 
     def unpickle(reader: TastyReader, tastyName: NameTable): String = {
+      val posUnpickler = new PositionUnpickler(reader, tastyName)
       sb.append(s" ${reader.endAddr.index - reader.currentAddr.index}")
-      val spans = new PositionUnpickler(reader, tastyName).spans
-      sb.append(s" position bytes:\n")
+      sb.append(" position bytes:\n")
+      val lineSizes = posUnpickler.lineSizes
+      sb.append(s"   lines: ${lineSizes.length}\n")
+      sb.append(posUnpickler.lineSizes.mkString("   line sizes: ", ", ", "\n"))
+      sb.append("   positions:\n")
+      val spans = posUnpickler.spans
       val sorted = spans.toSeq.sortBy(_._1.index)
       for ((addr, pos) <- sorted) {
         sb.append(treeStr("%10d".format(addr.index)))

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

@@ -1363,8 +1363,13 @@ class TreeUnpickler(reader: TastyReader,
     def sourceChangeContext(addr: Addr = currentAddr)(using Context): Context = {
       val path = sourcePathAt(addr)
       if (path.nonEmpty) {
+        val sourceFile = ctx.getSource(path)
+        posUnpicklerOpt match
+          case Some(posUnpickler) =>
+            sourceFile.setLineIndicesFromLineSizes(posUnpickler.lineSizes)
+          case _ =>
         pickling.println(i"source change at $addr: $path")
-        ctx.withSource(ctx.getSource(path))
+        ctx.withSource(sourceFile)
       }
       else ctx
     }

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

@@ -168,7 +168,7 @@ object PickledQuotes {
     if tree.span.exists then
       val positionWarnings = new mutable.ListBuffer[String]()
       new PositionPickler(pickler, treePkl.buf.addrOfTree, treePkl.treeAnnots)
-        .picklePositions(tree :: Nil, positionWarnings)
+        .picklePositions(ctx.compilationUnit.source, tree :: Nil, positionWarnings)
       positionWarnings.foreach(report.warning(_))
 
     val pickled = pickler.assembleParts()

compiler/src/dotty/tools/dotc/reporting/MessageRendering.scala

@@ -147,7 +147,7 @@ trait MessageRendering {
     val sb = mutable.StringBuilder()
     val posString = posStr(pos, diagnosticLevel, msg)
     if (posString.nonEmpty) sb.append(posString).append(EOL)
-    if (pos.exists && pos.source.file.exists) {
+    if (pos.exists) {
       val pos1 = pos.nonInlined
       val (srcBefore, srcAfter, offset) = sourceLines(pos1, diagnosticLevel)
       val marker = columnMarker(pos1, offset, diagnosticLevel)

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

@@ -73,7 +73,7 @@ class Pickler extends Phase {
           treePkl.compactify()
           if tree.span.exists then
             new PositionPickler(pickler, treePkl.buf.addrOfTree, treePkl.treeAnnots)
-              .picklePositions(tree :: Nil, positionWarnings)
+              .picklePositions(unit.source, tree :: Nil, positionWarnings)
 
           if !ctx.settings.YdropComments.value then
             new CommentPickler(pickler, treePkl.buf.addrOfTree, treePkl.docString)

compiler/src/dotty/tools/dotc/util/SourceFile.scala

@@ -82,7 +82,9 @@ class SourceFile(val file: AbstractFile, computeContent: => Array[Char]) extends
 
   def apply(idx: Int): Char = content().apply(idx)
 
-  def length: Int = content().length
+  def length: Int =
+    if lineIndicesCache ne null then lineIndicesCache.last
+    else content().length
 
   /** true for all source files except `NoSource` */
   def exists: Boolean = true
@@ -105,7 +107,8 @@ class SourceFile(val file: AbstractFile, computeContent: => Array[Char]) extends
   def positionInUltimateSource(position: SourcePosition): SourcePosition =
     SourcePosition(underlying, position.span shift start)
 
-  private def calculateLineIndices(cs: Array[Char]) = {
+  private def calculateLineIndicesFromContents() = {
+    val cs = content()
     val buf = new ArrayBuffer[Int]
     buf += 0
     var i = 0
@@ -120,7 +123,22 @@ class SourceFile(val file: AbstractFile, computeContent: => Array[Char]) extends
     buf += cs.length // sentinel, so that findLine below works smoother
     buf.toArray
   }
-  private lazy val lineIndices: Array[Int] = calculateLineIndices(content())
+
+  private var lineIndicesCache: Array[Int] = _
+  private def lineIndices: Array[Int] =
+    if lineIndicesCache eq null then
+      lineIndicesCache = calculateLineIndicesFromContents()
+    lineIndicesCache
+  def setLineIndicesFromLineSizes(sizes: Array[Int]): Unit =
+    val lines = sizes.length
+    val indices = new Array[Int](lines + 1)
+    var i = 0
+    val penultimate = lines - 1
+    while i < penultimate do
+      indices(i + 1) = indices(i) + sizes(i) + 1 // `+1` for the '\n' at the end of the line
+      i += 1
+    indices(lines) = indices(penultimate) + sizes(penultimate) // last line does not end with '\n'
+    lineIndicesCache = indices
 
   /** Map line to offset of first character in line */
   def lineToOffset(index: Int): Int = lineIndices(index)

compiler/src/dotty/tools/dotc/util/SourcePosition.scala

@@ -25,7 +25,7 @@ extends SrcPos, interfaces.SourcePosition, Showable {
 
   def point: Int = span.point
 
-  def line: Int = if (source.length != 0) source.offsetToLine(point) else -1
+  def line: Int = source.offsetToLine(point)
 
   /** Extracts the lines from the underlying source file as `Array[Char]`*/
   def linesSlice: Array[Char] =
@@ -45,16 +45,16 @@ extends SrcPos, interfaces.SourcePosition, Showable {
   def beforeAndAfterPoint: (List[Int], List[Int]) =
     lineOffsets.partition(_ <= point)
 
-  def column: Int = if (source.content().length != 0) source.column(point) else -1
+  def column: Int = source.column(point)
 
   def start: Int = span.start
-  def startLine: Int = if (source.content().length != 0) source.offsetToLine(start) else -1
-  def startColumn: Int = if (source.content().length != 0) source.column(start) else -1
+  def startLine: Int = source.offsetToLine(start)
+  def startColumn: Int = source.column(start)
   def startColumnPadding: String = source.startColumnPadding(start)
 
   def end: Int = span.end
-  def endLine: Int = if (source.content().length != 0) source.offsetToLine(end) else -1
-  def endColumn: Int = if (source.content().length != 0) source.column(end) else -1
+  def endLine: Int = source.offsetToLine(end)
+  def endColumn: Int = source.column(end)
 
   def withOuter(outer: SourcePosition): SourcePosition = SourcePosition(source, span, outer)
   def withSpan(range: Span) = SourcePosition(source, range, outer)

project/scripts/cmdTests

@@ -52,6 +52,7 @@ cp tests/neg/i6371/B_2.scala $OUT/B.scala
 rm $OUT/A.scala
 "$SBT" "scalac -classpath $OUT1 -d $OUT1 $OUT/B.scala" > "$tmp" 2>&1 || echo "ok"
 grep -qe "B.scala:2:7" "$tmp"
+grep -qe "This location contains code that was inlined from A.scala:3" "$tmp"
 
 echo "testing -Ythrough-tasty"
 clear_out "$OUT"