Put line sizes in the Positions section

Author: nicolasstucki

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

@@ -17,22 +17,17 @@ object DottyUnpickler {
   /** Exception thrown if classfile is corrupted */
   class BadSignature(msg: String) extends RuntimeException(msg)
 
-  class TreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], lineSizesUnpickler: Option[LineSizesUnpickler], commentUnpickler: Option[CommentUnpickler])
+  class TreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], commentUnpickler: Option[CommentUnpickler])
   extends SectionUnpickler[TreeUnpickler](TreePickler.sectionName) {
     def unpickle(reader: TastyReader, nameAtRef: NameTable): TreeUnpickler =
-      new TreeUnpickler(reader, nameAtRef, posUnpickler, lineSizesUnpickler, commentUnpickler)
+      new TreeUnpickler(reader, nameAtRef, posUnpickler, 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)
@@ -49,18 +44,17 @@ 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, lineSizesUnpicklerOpt, commentUnpicklerOpt)).get
+  private val treeUnpickler = unpickler.unpickle(treeSectionUnpickler(posUnpicklerOpt, 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], lineSizesUnpicklerOpt: Option[LineSizesUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler =
-    new TreeSectionUnpickler(posUnpicklerOpt, lineSizesUnpicklerOpt, commentUnpicklerOpt)
+  protected def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler =
+    new TreeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt)
 
   protected def computeRootTrees(using Context): List[Tree] = treeUnpickler.unpickle(mode)
 

compiler/src/dotty/tools/dotc/core/tasty/LineSizesPickler.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 pickleLineOffsetts(): Unit = {
+      val content = source.content()
+      buf.writeInt(content.count(_ == '\n') + 1) // number of lines
+      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
+    }
+    pickleLineOffsetts()
+
     var lastIndex = 0
     var lastSpan = Span(0, 0)
     def pickleDeltas(index: Int, span: Span) = {

compiler/src/dotty/tools/dotc/core/tasty/LineSizesUnpickler.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 = readInt()
+      myLineSizes = new Array[Int](lines)
+      var i = 0
+      while i < lines do
+        myLineSizes(i) += readInt()
+        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/PositionPickler.scala

@@ -50,11 +50,6 @@ 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 _ =>
@@ -143,9 +138,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(" 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)))
@@ -155,20 +155,6 @@ 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/PositionUnpickler.scala

@@ -49,13 +49,11 @@ 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 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._
@@ -1366,9 +1364,9 @@ class TreeUnpickler(reader: TastyReader,
       val path = sourcePathAt(addr)
       if (path.nonEmpty) {
         val sourceFile = ctx.getSource(path)
-        lineSizesUnpicklerOpt match
-          case Some(lineSizesUnpickler) =>
-            sourceFile.setLineIndices(lineSizesUnpickler.lineIndices)
+        posUnpicklerOpt match
+          case Some(posUnpickler) =>
+            sourceFile.setLineIndicesFromLineSizes(posUnpickler.lineSizes)
           case _ =>
         pickling.println(i"source change at $addr: $path")
         ctx.withSource(sourceFile)

compiler/src/dotty/tools/dotc/core/tasty/TastyPrinter.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/core/tasty/TreeUnpickler.scala

@@ -73,9 +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)
-            new LineSizesPickler(pickler)
-              .pickleLineNumbers(unit.source)
+              .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/quoted/PickledQuotes.scala

@@ -129,7 +129,15 @@ class SourceFile(val file: AbstractFile, computeContent: => Array[Char]) extends
     if lineIndicesCache eq null then
       lineIndicesCache = calculateLineIndicesFromContents()
     lineIndicesCache
-  def setLineIndices(indices: Array[Int]): Unit =
+  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 */

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

@@ -228,7 +228,9 @@ Note: Tree tags are grouped into 5 categories that determine what follows, and t
   Category 5 (tags 128-255):  tag Length <payload>
 
 
-Standard Section: "Positions" Assoc*
+Standard Section: "Positions" LinesSizes Assoc*
+
+  LinesSizes    = Int Int*                 // Number of lines followed by the size of each line not counting the trailing `\n`
 
   Assoc         = Header offset_Delta? offset_Delta? point_Delta?
                 | SOURCE nameref_Int
@@ -247,11 +249,6 @@ 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.
@@ -262,7 +259,7 @@ Standard Section: "Comments" Comment*
 object TastyFormat {
 
   final val header: Array[Int] = Array(0x5C, 0xA1, 0xAB, 0x1F)
-  val MajorVersion: Int = 25
+  val MajorVersion: Int = 26
   val MinorVersion: Int = 0
 
   /** Tags used to serialize names, should update [[nameTagToString]] if a new constant is added */