Address review comments

Author: Duhemm

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

@@ -42,21 +42,14 @@ 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)
+    if (!ctx.settings.YdropComments.value(ctx) || ctx.mode.is(Mode.ReadComments)) {
+      ctx.setProperty(ContextDoc, new ContextDocstrings)
+    }
 
     val fileNames = CompilerCommand.checkUsage(summary, sourcesRequired)(ctx)
     (fileNames, ctx)

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

@@ -4,14 +4,14 @@ import dotty.tools.dotc.core.tasty._
 import dotty.tools.dotc.core.tasty.TastyUnpickler.NameTable
 
 object TastyUnpickler {
-  class QuotedTreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], commentUnpickler: Option[CommentUnpickler], splices: Seq[Any])
-    extends DottyUnpickler.TreeSectionUnpickler(posUnpickler, commentUnpickler) {
+  class QuotedTreeSectionUnpickler(posUnpickler: Option[PositionUnpickler], splices: Seq[Any])
+    extends DottyUnpickler.TreeSectionUnpickler(posUnpickler, None) {
     override def unpickle(reader: TastyReader, nameAtRef: NameTable) =
-      new TreeUnpickler(reader, nameAtRef, posUnpickler, commentUnpickler, splices)
+      new TreeUnpickler(reader, nameAtRef, posUnpickler, None, splices)
   }
 }
 
-/** A class for unpickling quoted Tasty trees and symbols.
+/** A class for unpickling quoted Tasty trees and symbols. Comments are never unpickled.
  *  @param bytes         the bytearray containing the Tasty file from which we unpickle
  *  @param splices       splices that will fill the holes in the quote
  */
@@ -20,5 +20,5 @@ class TastyUnpickler(bytes: Array[Byte], splices: Seq[Any]) extends DottyUnpickl
   import TastyUnpickler._
 
   protected override def treeSectionUnpickler(posUnpicklerOpt: Option[PositionUnpickler], commentUnpicklerOpt: Option[CommentUnpickler]): TreeSectionUnpickler =
-    new QuotedTreeSectionUnpickler(posUnpicklerOpt, commentUnpicklerOpt, splices)
+    new QuotedTreeSectionUnpickler(posUnpicklerOpt, splices)
 }

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

@@ -1,7 +1,7 @@
 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.Comments.{Comment, CommentsContext, ContextDocstrings}
 import dotty.tools.dotc.core.Contexts.Context
 import dotty.tools.dotc.core.tasty.TastyBuffer.Addr
 
@@ -11,8 +11,10 @@ class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr]
   private[this] val buf = new TastyBuffer(5000)
   pickler.newSection("Comments", buf)
 
-  def pickleComment(root: tpd.Tree): Unit =
-    new Traverser().traverse(root)
+  def pickleComment(root: tpd.Tree): Unit = {
+    assert(ctx.docCtx.isDefined, "Trying to pickle comments, but there's no `docCtx`.")
+    new Traverser(ctx.docCtx.get).traverse(root)
+  }
 
   def pickleComment(addrOfTree: Option[Addr], comment: Option[Comment]): Unit = (addrOfTree, comment) match {
     case (Some(addr), Some(cmt)) =>
@@ -26,14 +28,12 @@ class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Option[Addr]
       ()
   }
 
-  private class Traverser extends tpd.TreeTraverser {
+  private class Traverser(docCtx: ContextDocstrings) 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)
-          }
+          val comment = docCtx.docstring(md.symbol)
+          pickleComment(addrOfTree(md), comment)
           traverseChildren(md)
         case _ =>
           traverseChildren(tree)

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

@@ -825,10 +825,10 @@ class TreeUnpickler(reader: TastyReader,
         sym.registerIfChild(late = true)
 
       if (ctx.mode.is(Mode.ReadComments)) {
-        for { docCtx <- ctx.docCtx
-              commentUnpickler <- commentUnpicklerOpt } {
+        assert(ctx.docCtx.isDefined, "Mode is `ReadComments`, but no `docCtx` is set.")
+        commentUnpicklerOpt.foreach { commentUnpickler =>
           val comment = commentUnpickler.commentAt(start)
-          docCtx.addDocstring(tree.symbol, comment)
+          ctx.docCtx.get.addDocstring(tree.symbol, comment)
           tree.setComment(comment)
         }
       }

compiler/test/dotty/tools/dotc/core/tasty/CommentPicklingTest.scala

@@ -2,7 +2,7 @@ package dotty.tools.dotc.core.tasty
 
 import dotty.tools.dotc.ast.tpd
 import dotty.tools.dotc.ast.tpd.TreeOps
-import dotty.tools.dotc.Driver
+import dotty.tools.dotc.{Driver, Main}
 import dotty.tools.dotc.core.Comments.CommentsContext
 import dotty.tools.dotc.core.Contexts.Context
 import dotty.tools.dotc.core.Decorators.PreNamedString
@@ -78,16 +78,11 @@ class CommentPicklingTest {
     }
   }
 
-  private def findTreeNamed(name: Name)(trees: List[tpd.Tree], ctx: Context): Option[tpd.NameTree] = {
-    val acc = new tpd.TreeAccumulator[Option[tpd.NameTree]] {
-      override def apply(x: Option[tpd.NameTree], tree: tpd.Tree)(implicit ctx: Context): Option[tpd.NameTree] = {
-        x.orElse(tree match {
-          case md: tpd.NameTree if md.name == name => Some(md)
-          case other => foldOver(None, other)
-        })
-      }
-    }
-    acc(None, trees)(ctx)
+  private def findTreeNamed(name: Name)(trees: List[tpd.Tree], ctx: Context): Option[tpd.MemberDef] = {
+    implicit val _ctx: Context = ctx
+    trees.flatMap { _.find { case md: tpd.MemberDef => md.name == name; case _ => false }
+      .map(_.asInstanceOf[tpd.MemberDef]).toList
+    }.headOption
   }
 
   private def compileAndUnpickle[T](sources: List[String])(fn: (List[tpd.Tree], Context) => T) = {
@@ -103,9 +98,8 @@ class CommentPicklingTest {
       Files.createDirectories(out)
 
       val options = compileOptions.and("-d", out.toAbsolutePath.toString).and(sourceFiles: _*)
-      val driver = new Driver
       val reporter = TestReporter.reporter(System.out, logLevel = ERROR)
-      driver.process(options.all, reporter)
+      Main.process(options.all, reporter)
       assertFalse("Compilation failed.", reporter.hasErrors)
 
       val tastyFiles = getAll(tmp, "glob:**.tasty")