Add Comments section in TASTY

compiler/src/dotty/tools/dotc/Driver.scala

@@ -2,7 +2,9 @@ package dotty.tools.dotc
 
 import dotty.tools.FatalError
 import config.CompilerCommand
+import core.Comments.{ContextDoc, ContextDocstrings}
 import core.Contexts.{Context, ContextBase}
+import core.Mode
 import util.DotClass
 import reporting._
 import scala.util.control.NonFatal
@@ -40,10 +42,22 @@ class Driver extends DotClass {
 
   protected def sourcesRequired = true
 
+  /**
+   * Should the `ContextDocstrings` be set for this context? The `ContextDocstrings` is used
+   * to store doc comments unless `-Ydrop-comments` is set, or when TASTY is configured to
+   * unpickle the doc comments.
+   */
+  protected def shouldAddDocContext(implicit ctx: Context): Boolean = {
+    !ctx.settings.YdropComments.value || ctx.mode.is(Mode.ReadComments)
+  }
+
   def setup(args: Array[String], rootCtx: Context): (List[String], Context) = {
     val ctx = rootCtx.fresh
     val summary = CompilerCommand.distill(args)(ctx)
     ctx.setSettings(summary.sstate)
+
+    if (shouldAddDocContext(ctx)) ctx.setProperty(ContextDoc, new ContextDocstrings)
+
     val fileNames = CompilerCommand.checkUsage(summary, sourcesRequired)(ctx)
     (fileNames, ctx)
   }

compiler/src/dotty/tools/dotc/config/ScalaSettings.scala

@@ -116,7 +116,7 @@ class ScalaSettings extends Settings.SettingGroup {
   val YshowPrintErrors = BooleanSetting("-Yshow-print-errors", "don't suppress exceptions thrown during tree printing.")
   val YtestPickler = BooleanSetting("-Ytest-pickler", "self-test for pickling functionality; should be used with -Ystop-after:pickler")
   val YcheckReentrant = BooleanSetting("-Ycheck-reentrant", "check that compiled program does not contain vars that can be accessed from a global root.")
-  val YkeepComments = BooleanSetting("-Ykeep-comments", "Keep comments when scanning source files.")
+  val YdropComments = BooleanSetting("-Ydrop-comments", "Drop comments when scanning source files.")
   val YcookComments = BooleanSetting("-Ycook-comments", "Cook the comments (type check `@usecase`, etc.)")
   val YforceSbtPhases = BooleanSetting("-Yforce-sbt-phases", "Run the phases used by sbt for incremental compilation (ExtractDependencies and ExtractAPI) even if the compiler is ran outside of sbt, for debugging.")
   val YdumpSbtInc = BooleanSetting("-Ydump-sbt-inc", "For every compiled foo.scala, output the API representation and dependencies used for sbt incremental compilation in foo.inc, implies -Yforce-sbt-phases.")

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

@@ -99,7 +99,7 @@ object Comments {
   }
 
   object Comment {
-    def apply(pos: Position, raw: String, expanded: Boolean = false, usc: List[UseCase] = Nil)(implicit ctx: Context): Comment =
+    def apply(pos: Position, raw: String, expanded: Boolean = false, usc: List[UseCase] = Nil): Comment =
       new Comment(pos, raw) {
         val isExpanded = expanded
         val usecases = usc

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

@@ -93,4 +93,8 @@ object Mode {
 
   /** We are in the IDE */
   val Interactive = newMode(20, "Interactive")
+
+  /** Read comments from definitions when unpickling from TASTY */
+  val ReadComments = newMode(21, "ReadComments")
+
 }

compiler/src/dotty/tools/dotc/core/quoted/TastyUnpickler.scala

@@ -4,10 +4,10 @@ import dotty.tools.dotc.core.tasty._
 import dotty.tools.dotc.core.tasty.TastyUnpickler.NameTable
 
 object TastyUnpickler {
-  class QuotedTreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], splices: Seq[Any])
-    extends DottyUnpickler.TreeSectionUnpickler(posUnpickler) {
+  class QuotedTreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], commentUnpickler: Option[CommentUnpickler], splices: Seq[Any])
+    extends DottyUnpickler.TreeSectionUnpickler(posUnpickler, commentUnpickler) {
     override def unpickle(reader: TastyReader, nameAtRef: NameTable) =
-      new TreeUnpickler(reader, nameAtRef, posUnpickler, splices)
+      new TreeUnpickler(reader, nameAtRef, posUnpickler, commentUnpickler, splices)
   }
 }
 
@@ -19,6 +19,6 @@ class TastyUnpickler(bytes: Array[Byte], splices: Seq[Any]) extends DottyUnpickl
   import DottyUnpickler._
   import TastyUnpickler._
 
-  protected override def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler]): TreeSectionUnpickler =
-    new QuotedTreeSectionUnpickler(posUnpicklerOpt, splices)
+  protected override def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler =
+    new QuotedTreeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt, splices)
 }

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

@@ -0,0 +1,43 @@
+package dotty.tools.dotc.core.tasty
+
+import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.core.Comments.{Comment, CommentsContext}
+import dotty.tools.dotc.core.Contexts.Context
+import dotty.tools.dotc.core.tasty.TastyBuffer.Addr
+
+import java.nio.charset.Charset
+
+class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr])(implicit ctx: Context) {
+  val buf = new TastyBuffer(5000)
+  pickler.newSection("Comments", buf)
+
+  def pickleComment(root: tpd.Tree): Unit =
+    new Traverser().traverse(root)
+
+  def pickleComment(addrOfTree: Option[Addr], comment: Option[Comment]): Unit = (addrOfTree, comment) match {
+    case (Some(addr), Some(cmt)) =>
+      val bytes = cmt.raw.getBytes(Charset.forName("UTF-8"))
+      val length = bytes.length
+      buf.writeAddr(addr)
+      buf.writeNat(length)
+      buf.writeBytes(bytes, length)
+      buf.writeBoolean(cmt.isExpanded)
+    case other =>
+      ()
+  }
+
+  private class Traverser extends tpd.TreeTraverser {
+    override def traverse(tree: tpd.Tree)(implicit ctx: Context): Unit =
+      tree match {
+        case md: tpd.MemberDef =>
+          ctx.docCtx.foreach { docCtx =>
+            val comment = docCtx.docstring(md.symbol)
+            pickleComment(addrOfTree(md), comment)
+          }
+          traverseChildren(md)
+        case _ =>
+          traverseChildren(tree)
+      }
+  }
+
+}

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

@@ -0,0 +1,33 @@
+package dotty.tools.dotc.core.tasty
+
+import dotty.tools.dotc.core.Comments.Comment
+import dotty.tools.dotc.core.Symbols.Symbol
+import dotty.tools.dotc.core.tasty.TastyBuffer.Addr
+import dotty.tools.dotc.util.Positions
+
+import scala.collection.mutable.HashMap
+
+import java.nio.charset.Charset
+
+class CommentUnpickler(reader: TastyReader) {
+  import reader._
+
+  private[tasty] lazy val comments = {
+    val comments = new HashMap[Addr, Comment]
+    while (!isAtEnd) {
+      val addr = readAddr()
+      val length = readNat()
+      if (length > 0) {
+        val bytes = readBytes(length)
+        val expanded = readBoolean()
+        val rawComment = new String(bytes, Charset.forName("UTF-8"))
+        comments(addr) = Comment(Positions.NoPosition, rawComment, expanded = expanded)
+      }
+    }
+    comments.toMap
+  }
+
+  def commentAt(addr: Addr): Option[Comment] =
+    comments.get(addr)
+
+}

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

@@ -17,16 +17,21 @@ object DottyUnpickler {
   /** Exception thrown if classfile is corrupted */
   class BadSignature(msg: String) extends RuntimeException(msg)
 
-  class TreeSectionUnpickler(posUnpickler: Option[PositionUnpickler])
+  class TreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], commentUnpickler: Option[CommentUnpickler])
   extends SectionUnpickler[TreeUnpickler]("ASTs") {
     def unpickle(reader: TastyReader, nameAtRef: NameTable) =
-      new TreeUnpickler(reader, nameAtRef, posUnpickler, Seq.empty)
+      new TreeUnpickler(reader, nameAtRef, posUnpickler, commentUnpickler, Seq.empty)
   }
 
   class PositionsSectionUnpickler extends SectionUnpickler[PositionUnpickler]("Positions") {
     def unpickle(reader: TastyReader, nameAtRef: NameTable) =
       new PositionUnpickler(reader)
   }
+
+  class CommentsSectionUnpickler extends SectionUnpickler[CommentUnpickler]("Comments") {
+    def unpickle(reader: TastyReader, nameAtRef: NameTable): CommentUnpickler =
+      new CommentUnpickler(reader)
+  }
 }
 
 /** A class for unpickling Tasty trees and symbols.
@@ -38,7 +43,8 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded with t
 
   val unpickler = new TastyUnpickler(bytes)
   private val posUnpicklerOpt = unpickler.unpickle(new PositionsSectionUnpickler)
-  private val treeUnpickler = unpickler.unpickle(treeSectionUnpickler(posUnpicklerOpt)).get
+  private val commentUnpicklerOpt = unpickler.unpickle(new CommentsSectionUnpickler)
+  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
@@ -52,8 +58,8 @@ class DottyUnpickler(bytes: Array[Byte]) extends ClassfileParser.Embedded with t
   def unpickleTypeTree()(implicit ctx: Context): Tree =
     treeUnpickler.unpickleTypeTree()
 
-  protected def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler]): TreeSectionUnpickler = {
-    new TreeSectionUnpickler(posUnpicklerOpt)
+  protected def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler = {
+    new TreeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt)
   }
 
   protected def computeTrees(implicit ctx: Context) = treeUnpickler.unpickle()

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

@@ -48,6 +48,10 @@ class TastyBuffer(initialSize: Int) {
 
   // -- Output routines --------------------------------------------
 
+  /** Write a boolean value. */
+  def writeBoolean(b: Boolean): Unit =
+    writeByte(if (b) 1 else 0)
+
   /** Write a byte of data. */
   def writeByte(b: Int): Unit = {
     if (length >= bytes.length)

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

@@ -222,6 +222,10 @@ Standard Section: "Positions" Assoc*
                                             // same offset in the previously recorded node (or 0 for the first recorded node)
   Delta         = Int                       // Difference between consecutive offsets,
 
+Standard Section: "Comments" Comment*
+
+  Comment       = Length Bytes              // Raw comment's bytes encoded as UTF-8
+
 **************************************************************************************/
 
 object TastyFormat {

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

@@ -41,6 +41,10 @@ class TastyReader(val bytes: Array[Byte], start: Int, end: Int, val base: Int =
   def subReader(start: Addr, end: Addr): TastyReader =
     new TastyReader(bytes, index(start), index(end), base)
 
+  /** Read a boolean value. */
+  def readBoolean(): Boolean =
+    readByte() == 1
+
   /** Read a byte of data. */
   def readByte(): Int = {
     val result = bytes(bp) & 0xff

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

@@ -3,6 +3,7 @@ package dotc
 package core
 package tasty
 
+import Comments.CommentsContext
 import Contexts._, Symbols._, Types._, Scopes._, SymDenotations._, Names._, NameOps._
 import StdNames._, Denotations._, Flags._, Constants._, Annotations._
 import NameKinds._
@@ -25,13 +26,15 @@ import scala.quoted.Types.TreeType
 import scala.quoted.Exprs.TastyTreeExpr
 
 /** Unpickler for typed trees
- *  @param reader          the reader from which to unpickle
- *  @param tastyName       the nametable
- *  @param posUNpicklerOpt the unpickler for positions, if it exists
+ *  @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 splices
  */
 class TreeUnpickler(reader: TastyReader,
                     nameAtRef: NameRef => TermName,
                     posUnpicklerOpt: Option[PositionUnpickler],
+                    commentUnpicklerOpt: Option[CommentUnpickler],
                     splices: Seq[Any]) {
   import TastyFormat._
   import TreeUnpickler._
@@ -820,6 +823,16 @@ class TreeUnpickler(reader: TastyReader,
         // Child annotations for local classes and enum values are not pickled, so
         // need to be re-established here.
         sym.registerIfChild(late = true)
+
+      if (ctx.mode.is(Mode.ReadComments)) {
+        for { docCtx <- ctx.docCtx
+              commentUnpickler <- commentUnpicklerOpt } {
+          val comment = commentUnpickler.commentAt(start)
+          docCtx.addDocstring(tree.symbol, comment)
+          tree.setComment(comment)
+        }
+      }
+
       tree
     }
 

compiler/src/dotty/tools/dotc/parsing/Scanners.scala

@@ -172,7 +172,7 @@ object Scanners {
   }
 
   class Scanner(source: SourceFile, override val startFrom: Offset = 0)(implicit ctx: Context) extends ScannerCommon(source)(ctx) {
-    val keepComments = ctx.settings.YkeepComments.value
+    val keepComments = !ctx.settings.YdropComments.value
 
     /** All doc comments kept by their end position in a `Map` */
     private[this] var docstringMap: SortedMap[Int, Comment] = SortedMap.empty

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

@@ -60,6 +60,9 @@ class Pickler extends Phase {
       if (tree.pos.exists)
         new PositionPickler(pickler, treePkl.buf.addrOfTree).picklePositions(tree :: Nil)
 
+      if (!ctx.settings.YdropComments.value)
+        new CommentPickler(pickler, treePkl.buf.addrOfTree).pickleComment(tree)
+
       // other pickle sections go here.
       val pickled = pickler.assembleParts()
       unit.pickled += (cls -> pickled)
@@ -84,6 +87,7 @@ class Pickler extends Phase {
             .setPeriod(Period(ctx.runId + 1, FirstPhaseId))
             .setReporter(new ThrowingReporter(ctx.reporter))
             .addMode(Mode.ReadPositions)
+            .addMode(Mode.ReadComments)
             .addMode(Mode.PrintShowExceptions))
     result
   }