Make comment pickling context independent

Author: odersky

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

@@ -1,6 +1,6 @@
 package dotty.tools.dotc.core.tasty
 
-import dotty.tools.dotc.ast.tpd
+import dotty.tools.dotc.ast.{tpd, untpd}
 import dotty.tools.dotc.core.Comments.{Comment, CommentsContext, ContextDocstrings}
 import dotty.tools.dotc.core.Contexts._
 
@@ -9,36 +9,36 @@ import TastyBuffer.{Addr, NoAddr}
 
 import java.nio.charset.Charset
 
-class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Addr)(using Context) {
+class CommentPickler(pickler: TastyPickler, addrOfTree: tpd.Tree => Addr, docString: untpd.MemberDef => Option[Comment]):
   private val buf = new TastyBuffer(5000)
   pickler.newSection("Comments", buf)
 
-  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(root: tpd.Tree): Unit = traverse(root)
 
-  def pickleComment(addr: Addr, comment: Option[Comment]): Unit = comment match {
-    case Some(cmt) if addr != NoAddr =>
-      val bytes = cmt.raw.getBytes(Charset.forName("UTF-8"))
+  private def pickleComment(addr: Addr, comment: Comment): Unit =
+    if addr != NoAddr then
+      val bytes = comment.raw.getBytes(Charset.forName("UTF-8"))
       val length = bytes.length
       buf.writeAddr(addr)
       buf.writeNat(length)
       buf.writeBytes(bytes, length)
-      buf.writeLongInt(cmt.span.coords)
-    case other =>
-      ()
-  }
-
-  private class Traverser(docCtx: ContextDocstrings) extends tpd.TreeTraverser {
-    override def traverse(tree: tpd.Tree)(using Context): Unit =
-      tree match {
-        case md: tpd.MemberDef =>
-          val comment = docCtx.docstring(md.symbol)
-          pickleComment(addrOfTree(md), comment)
-          traverseChildren(md)
+      buf.writeLongInt(comment.span.coords)
+
+  private def traverse(x: Any): Unit = x match
+    case x: untpd.Tree @unchecked =>
+      x match
+        case x: tpd.MemberDef @unchecked => // at this point all MembderDefs are t(y)p(e)d.
+          for comment <- docString(x) do pickleComment(addrOfTree(x), comment)
         case _ =>
-          traverseChildren(tree)
-      }
-  }
-}
+      val limit = x.productArity
+      var n = 0
+      while n < limit do
+        traverse(x.productElement(n))
+        n += 1
+    case y :: ys =>
+      traverse(y)
+      traverse(ys)
+    case _ =>
+
+end CommentPickler
+

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

@@ -10,6 +10,7 @@ import ast.Trees._
 import ast.{untpd, tpd}
 import Contexts._, Symbols._, Types._, Names._, Constants._, Decorators._, Annotations._, Flags._
 import Denotations.MultiDenotation
+import Comments.{Comment, CommentsContext}
 import typer.Inliner
 import NameKinds._
 import StdNames.nme
@@ -50,10 +51,20 @@ class TreePickler(pickler: TastyPickler) {
    */
   private val annotTrees = util.HashTable[untpd.MemberDef, mutable.ListBuffer[Tree]]()
 
+  /** A map from member definitions to their doc comments, so that later
+   *  parallel comment pickling does not need to access symbols of trees (which
+   *  would involve accessing symbols of named types and possibly changing phases
+   *  in doing so).
+   */
+  private val docStrings = util.HashTable[untpd.MemberDef, Comment]()
+
   def treeAnnots(tree: untpd.MemberDef): List[Tree] =
     val ts = annotTrees.lookup(tree)
     if ts == null then Nil else ts.toList
 
+  def docString(tree: untpd.MemberDef): Option[Comment] =
+    Option(docStrings.lookup(tree))
+
   private def withLength(op: => Unit) = {
     val lengthAddr = reserveRef(relative = true)
     op
@@ -334,6 +345,11 @@ class TreePickler(pickler: TastyPickler) {
       pickleTreeUnlessEmpty(rhs)
       pickleModifiers(sym, mdef)
     }
+    for
+      docCtx <- ctx.docCtx
+      comment <- docCtx.docstring(sym)
+    do
+      docStrings.enter(mdef, comment)
   }
 
   def pickleParam(tree: Tree)(using Context): Unit = {

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

@@ -62,16 +62,16 @@ class Pickler extends Phase {
       val treePkl = pickler.treePkl
       treePkl.pickle(tree :: Nil)
       val positionWarnings = new mutable.ListBuffer[String]()
-      val parContext = ctx.fresh // TODO refine
-      val pickledF = inContext(parContext) {
+      val pickledF = inContext(ctx.fresh) {
         Future {
           treePkl.compactify()
           if tree.span.exists then
             new PositionPickler(pickler, treePkl.buf.addrOfTree, treePkl.treeAnnots)
               .picklePositions(tree :: Nil, positionWarnings)
 
           if !ctx.settings.YdropComments.value then
-            new CommentPickler(pickler, treePkl.buf.addrOfTree).pickleComment(tree)
+            new CommentPickler(pickler, treePkl.buf.addrOfTree, treePkl.docString)
+              .pickleComment(tree)
 
           val pickled = pickler.assembleParts()