Fix #6542: Pickle line sizes in TASTy

Author: nicolasstucki

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

@@ -17,17 +17,22 @@ object DottyUnpickler {
   /** Exception thrown if classfile is corrupted */
   class BadSignature(msg: String) extends RuntimeException(msg)
 
-  class TreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], commentUnpickler: Option[CommentUnpickler])
+  class TreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], lineSizesUnpickler: Option[LineSizesUnpickler], commentUnpickler: Option[CommentUnpickler])
   extends SectionUnpickler[TreeUnpickler](TreePickler.sectionName) {
     def unpickle(reader: TastyReader, nameAtRef: NameTable): TreeUnpickler =
-      new TreeUnpickler(reader, nameAtRef, posUnpickler, commentUnpickler)
+      new TreeUnpickler(reader, nameAtRef, posUnpickler, lineSizesUnpickler, commentUnpickler)
   }
 
   class PositionsSectionUnpickler extends SectionUnpickler[PositionUnpickler]("Positions") {
     def unpickle(reader: TastyReader, nameAtRef: NameTable): PositionUnpickler =
       new PositionUnpickler(reader, nameAtRef)
   }
 
+  class LineSizesSectionUnpickler extends SectionUnpickler[LineSizesUnpickler]("LineSizes") {
+    def unpickle(reader: TastyReader, nameAtRef: NameTable): LineSizesUnpickler =
+      new LineSizesUnpickler(reader)
+  }
+
   class CommentsSectionUnpickler extends SectionUnpickler[CommentUnpickler]("Comments") {
     def unpickle(reader: TastyReader, nameAtRef: NameTable): CommentUnpickler =
       new CommentUnpickler(reader)
@@ -44,17 +49,18 @@ class DottyUnpickler(bytes: Array[Byte], mode: UnpickleMode = UnpickleMode.TopLe
 
   val unpickler: TastyUnpickler = new TastyUnpickler(bytes)
   private val posUnpicklerOpt = unpickler.unpickle(new PositionsSectionUnpickler)
+  private val lineSizesUnpicklerOpt = unpickler.unpickle(new LineSizesSectionUnpickler)
   private val commentUnpicklerOpt = unpickler.unpickle(new CommentsSectionUnpickler)
-  private val treeUnpickler = unpickler.unpickle(treeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt)).get
+  private val treeUnpickler = unpickler.unpickle(treeSectionUnpickler(posUnpicklerOpt, lineSizesUnpicklerOpt, commentUnpicklerOpt)).get
 
   /** Enter all toplevel classes and objects into their scopes
    *  @param roots          a set of SymDenotations that should be overwritten by unpickling
    */
   def enter(roots: Set[SymDenotation])(using Context): Unit =
     treeUnpickler.enter(roots)
 
-  protected def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler =
-    new TreeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt)
+  protected def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler], lineSizesUnpicklerOpt: Option[LineSizesUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler =
+    new TreeSectionUnpickler(posUnpicklerOpt, lineSizesUnpicklerOpt, commentUnpicklerOpt)
 
   protected def computeRootTrees(using Context): List[Tree] = treeUnpickler.unpickle(mode)
 

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

@@ -0,0 +1,35 @@
+package dotty.tools
+package dotc
+package core
+package tasty
+
+import dotty.tools.tasty.TastyFormat.SOURCE
+import dotty.tools.tasty.TastyBuffer
+import TastyBuffer._
+
+import ast._
+import ast.Trees._
+import ast.Trees.WithLazyField
+import util.{SourceFile, NoSource}
+import core._
+import Contexts._, Symbols._, Annotations._, Decorators._
+import collection.mutable
+import util.Spans._
+
+class LineSizesPickler(pickler: TastyPickler) {
+
+  import ast.tpd._
+  val buf: TastyBuffer = new TastyBuffer(5000)
+  pickler.newSection("LineSizes", buf)
+
+  def pickleLineNumbers(source: SourceFile): Unit =
+    val content = source.content()
+    var lastIndex = content.indexOf('\n', 0)
+    buf.writeInt(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.writeInt(end - lastIndex - 1) // size of the next line
+      lastIndex = nextIndex
+
+}

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

@@ -0,0 +1,45 @@
+package dotty.tools
+package dotc
+package core
+package tasty
+
+import dotty.tools.tasty.{TastyFormat, TastyBuffer, TastyReader}
+import TastyFormat.SOURCE
+import TastyBuffer.{Addr, NameRef}
+
+import util.Spans._
+import collection.{mutable, Map}
+import Names.TermName
+
+/** Unpickler for tree positions */
+class LineSizesUnpickler(reader: TastyReader) {
+  import reader._
+
+  private var mySizes: Array[Int] = _
+  private var isDefined = false
+
+  def ensureDefined(): Unit =
+    if !isDefined then
+      val sizeBuf = Array.newBuilder[Int]
+      // Number of lines if all lines are 127 characters or less
+      var lowSizeBound = endAddr.index - currentAddr.index
+      sizeBuf.sizeHint(lowSizeBound)
+      while !isAtEnd do sizeBuf += readInt()
+      mySizes = sizeBuf.result()
+      isDefined = true
+
+  private[tasty] def sizes: Array[Int] =
+    ensureDefined()
+    mySizes
+
+  private[tasty] def lineIndices: Array[Int] =
+    val szs = sizes
+    val indices = new Array[Int](sizes.length + 1)
+    var i = 0
+    val penultimate = szs.length - 1
+    while i < penultimate do
+      indices(i + 1) = indices(i) + szs(i) + 1 // `+1` for the '\n' at the end of the line
+      i += 1
+    indices(szs.length) = indices(penultimate) + szs(penultimate) // last line does not end with '\n'
+    indices
+}

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

@@ -50,6 +50,11 @@ class TastyPrinter(bytes: Array[Byte]) {
       case _ =>
     }
     sb.append("\n\n")
+    unpickle(new LineSizesSectionUnpickler) match {
+      case Some(s) => sb.append(s)
+      case _ =>
+    }
+    sb.append("\n\n")
     unpickle(new CommentSectionUnpickler) match {
       case Some(s) => sb.append(s)
       case _ =>
@@ -140,7 +145,7 @@ class TastyPrinter(bytes: Array[Byte]) {
     def unpickle(reader: TastyReader, tastyName: NameTable): String = {
       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 sorted = spans.toSeq.sortBy(_._1.index)
       for ((addr, pos) <- sorted) {
         sb.append(treeStr("%10d".format(addr.index)))
@@ -150,6 +155,20 @@ class TastyPrinter(bytes: Array[Byte]) {
     }
   }
 
+  class LineSizesSectionUnpickler extends SectionUnpickler[String]("LineSizes") {
+
+    private val sb: StringBuilder = new StringBuilder
+
+    def unpickle(reader: TastyReader, tastyName: NameTable): String = {
+      sb.append(" ").append(reader.endAddr.index - reader.currentAddr.index)
+      sb.append(" line sizes bytes:\n")
+      val lineSizes = new LineSizesUnpickler(reader)
+      sb.append("   sizes: ")
+      sb.append(lineSizes.sizes.mkString(", "))
+      sb.result
+    }
+  }
+
   class CommentSectionUnpickler extends SectionUnpickler[String]("Comments") {
 
     private val sb: StringBuilder = new StringBuilder

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

@@ -47,13 +47,15 @@ import scala.annotation.constructorOnly
 import scala.annotation.internal.sharable
 
 /** Unpickler for typed trees
- *  @param reader              the reader from which to unpickle
- *  @param posUnpicklerOpt     the unpickler for positions, if it exists
- *  @param commentUnpicklerOpt the unpickler for comments, if it exists
+ *  @param reader                the reader from which to unpickle
+ *  @param posUnpicklerOpt       the unpickler for positions, if it exists
+ *  @param lineSizesUnpicklerOpt the unpickler for line sizes, if it exists
+ *  @param commentUnpicklerOpt   the unpickler for comments, if it exists
  */
 class TreeUnpickler(reader: TastyReader,
                     nameAtRef: NameTable,
                     posUnpicklerOpt: Option[PositionUnpickler],
+                    lineSizesUnpicklerOpt: Option[LineSizesUnpickler],
                     commentUnpicklerOpt: Option[CommentUnpickler]) {
   import TreeUnpickler._
   import tpd._
@@ -1363,8 +1365,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)
+        lineSizesUnpicklerOpt match
+          case Some(lineSizesUnpickler) =>
+            sourceFile.setLineIndices(lineSizesUnpickler.lineIndices)
+          case _ =>
         pickling.println(i"source change at $addr: $path")
-        ctx.withSource(ctx.getSource(path))
+        ctx.withSource(sourceFile)
       }
       else ctx
     }

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

@@ -74,6 +74,8 @@ class Pickler extends Phase {
           if tree.span.exists then
             new PositionPickler(pickler, treePkl.buf.addrOfTree, treePkl.treeAnnots)
               .picklePositions(tree :: Nil, positionWarnings)
+            new LineSizesPickler(pickler)
+              .pickleLineNumbers(unit.source)
 
           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,14 @@ 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 setLineIndices(indices: Array[Int]): Unit =
+    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)

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

@@ -50,6 +50,7 @@ Macro-format:
 Note: Unqualified names in the name table are strings. The context decides whether a name is
 a type-name or a term-name. The same string can represent both.
 
+
 Standard-Section: "ASTs" TopLevelStat*
 
   TopLevelStat  = PACKAGE        Length Path TopLevelStat*                         -- package path { topLevelStats }
@@ -226,6 +227,7 @@ Note: Tree tags are grouped into 5 categories that determine what follows, and t
   Category 4 (tags 110-127):  tag Nat AST
   Category 5 (tags 128-255):  tag Length <payload>
 
+
 Standard Section: "Positions" Assoc*
 
   Assoc         = Header offset_Delta? offset_Delta? point_Delta?
@@ -244,6 +246,12 @@ Standard Section: "Positions" Assoc*
 
 All elements of a position section are serialized as Ints
 
+
+Standard Section: "LineSizes" LineSize*
+
+  LineSize       = Int                      // Size the i-th line not counting the trailing `\n`
+
+
 Standard Section: "Comments" Comment*
 
   Comment       = Length Bytes LongInt      // Raw comment's bytes encoded as UTF-8, followed by the comment's coordinates.