Re-enable parallel pickling

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/PositionPickler.scala

@@ -16,10 +16,14 @@ import Contexts._, Symbols._, Annotations._, Decorators._
 import collection.mutable
 import util.Spans._
 
-class PositionPickler(pickler: TastyPickler, addrOfTree: PositionPickler.TreeToAddr) {
+class PositionPickler(
+    pickler: TastyPickler,
+    addrOfTree: PositionPickler.TreeToAddr,
+    treeAnnots: untpd.MemberDef => List[tpd.Tree]) {
+
+  import ast.tpd._
   val buf: TastyBuffer = new TastyBuffer(5000)
   pickler.newSection("Positions", buf)
-  import ast.tpd._
 
   private val pickledIndices = new mutable.BitSet
 
@@ -28,7 +32,7 @@ class PositionPickler(pickler: TastyPickler, addrOfTree: PositionPickler.TreeToA
     (addrDelta << 3) | (toInt(hasStartDelta) << 2) | (toInt(hasEndDelta) << 1) | toInt(hasPoint)
   }
 
-  def picklePositions(roots: List[Tree])(using Context): Unit = {
+  def picklePositions(roots: List[Tree], warnings: mutable.ListBuffer[String]): Unit = {
     var lastIndex = 0
     var lastSpan = Span(0, 0)
     def pickleDeltas(index: Int, span: Span) = {
@@ -55,7 +59,7 @@ class PositionPickler(pickler: TastyPickler, addrOfTree: PositionPickler.TreeToA
           val cwd = java.nio.file.Paths.get("").toAbsolutePath().normalize()
           try cwd.relativize(path)
           catch case _: IllegalArgumentException =>
-            report.warning("Could not relativize path for pickling: " + originalPath)
+            warnings += "Could not relativize path for pickling: " + originalPath
             originalPath
         else
           originalPath
@@ -100,10 +104,9 @@ class PositionPickler(pickler: TastyPickler, addrOfTree: PositionPickler.TreeToA
               pickleDeltas(addr.index, x.span)
           }
         }
-        x match {
-          case x: untpd.MemberDef @unchecked => traverse(x.symbol.annotations, x.source)
+        x match
+          case x: untpd.MemberDef => traverse(treeAnnots(x), x.source)
           case _ =>
-        }
         val limit = x.productArity
         var n = 0
         while (n < limit) {
@@ -113,8 +116,6 @@ class PositionPickler(pickler: TastyPickler, addrOfTree: PositionPickler.TreeToA
       case y :: ys =>
         traverse(y, current)
         traverse(ys, current)
-      case x: Annotation =>
-        traverse(x.tree, current)
       case _ =>
     }
     for (root <- roots)

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

@@ -177,12 +177,11 @@ class TreeBuffer extends TastyBuffer(50000) {
     //println(s"offsets: ${offsets.take(numOffsets).deep}")
     //println(s"deltas: ${delta.take(numOffsets).deep}")
     var saved = 0
-    while ({
+    while
       saved = adjustDeltas()
       pickling.println(s"adjusting deltas, saved = $saved")
       saved > 0 && length / saved < 100
-    })
-    ()
+    do ()
     adjustOffsets()
     adjustTreeAddrs()
     val wasted = compress()

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
@@ -18,6 +19,7 @@ import printing.Printer
 import printing.Texts._
 import util.SourceFile
 import annotation.constructorOnly
+import collection.mutable
 
 object TreePickler {
 
@@ -44,6 +46,25 @@ class TreePickler(pickler: TastyPickler) {
   private val forwardSymRefs = Symbols.newMutableSymbolMap[List[Addr]]
   private val pickledTypes = new java.util.IdentityHashMap[Type, Any] // Value type is really Addr, but that's not compatible with null
 
+  /** A list of annotation trees for every member definition, so that later
+   *  parallel position pickling does not need to access and force symbols.
+   */
+  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
@@ -309,7 +330,8 @@ class TreePickler(pickler: TastyPickler) {
     if (!tree.isEmpty) pickleTree(tree)
   }
 
-  def pickleDef(tag: Int, sym: Symbol, tpt: Tree, rhs: Tree = EmptyTree, pickleParams: => Unit = ())(using Context): Unit = {
+  def pickleDef(tag: Int, mdef: MemberDef, tpt: Tree, rhs: Tree = EmptyTree, pickleParams: => Unit = ())(using Context): Unit = {
+    val sym = mdef.symbol
     assert(symRefs(sym) == NoAddr, sym)
     registerDef(sym)
     writeByte(tag)
@@ -321,16 +343,21 @@ class TreePickler(pickler: TastyPickler) {
         case _ if tpt.isType => pickleTpt(tpt)
       }
       pickleTreeUnlessEmpty(rhs)
-      pickleModifiers(sym)
+      pickleModifiers(sym, mdef)
     }
+    for
+      docCtx <- ctx.docCtx
+      comment <- docCtx.docstring(sym)
+    do
+      docStrings.enter(mdef, comment)
   }
 
   def pickleParam(tree: Tree)(using Context): Unit = {
     registerTreeAddr(tree)
     tree match {
-      case tree: ValDef => pickleDef(PARAM, tree.symbol, tree.tpt)
-      case tree: DefDef => pickleDef(PARAM, tree.symbol, tree.tpt, tree.rhs)
-      case tree: TypeDef => pickleDef(TYPEPARAM, tree.symbol, tree.rhs)
+      case tree: ValDef => pickleDef(PARAM, tree, tree.tpt)
+      case tree: DefDef => pickleDef(PARAM, tree, tree.tpt, tree.rhs)
+      case tree: TypeDef => pickleDef(TYPEPARAM, tree, tree.rhs)
     }
   }
 
@@ -520,7 +547,7 @@ class TreePickler(pickler: TastyPickler) {
             patterns.foreach(pickleTree)
           }
         case tree: ValDef =>
-          pickleDef(VALDEF, tree.symbol, tree.tpt, tree.rhs)
+          pickleDef(VALDEF, tree, tree.tpt, tree.rhs)
         case tree: DefDef =>
           def pickleParamss(paramss: List[List[ValDef]]): Unit = paramss match
             case Nil =>
@@ -531,9 +558,9 @@ class TreePickler(pickler: TastyPickler) {
           def pickleAllParams =
             pickleParams(tree.tparams)
             pickleParamss(tree.vparamss)
-          pickleDef(DEFDEF, tree.symbol, tree.tpt, tree.rhs, pickleAllParams)
+          pickleDef(DEFDEF, tree, tree.tpt, tree.rhs, pickleAllParams)
         case tree: TypeDef =>
-          pickleDef(TYPEDEF, tree.symbol, tree.rhs)
+          pickleDef(TYPEDEF, tree, tree.rhs)
         case tree: Template =>
           registerDef(tree.symbol)
           writeByte(TEMPLATE)
@@ -648,7 +675,7 @@ class TreePickler(pickler: TastyPickler) {
     pickleName(id.name)
   }
 
-  def pickleModifiers(sym: Symbol)(using Context): Unit = {
+  def pickleModifiers(sym: Symbol, mdef: MemberDef)(using Context): Unit = {
     import Flags._
     var flags = sym.flags
     val privateWithin = sym.privateWithin
@@ -660,7 +687,7 @@ class TreePickler(pickler: TastyPickler) {
     if (flags.is(ParamAccessor) && sym.isTerm && !sym.isSetter)
       flags = flags &~ ParamAccessor // we only generate a tag for parameter setters
     pickleFlags(flags, sym.isTerm)
-    sym.annotations.foreach(pickleAnnotation(sym, _))
+    val annots = sym.annotations.foreach(pickleAnnotation(sym, mdef, _))
   }
 
   def pickleFlags(flags: FlagSet, isTerm: Boolean)(using Context): Unit = {
@@ -723,12 +750,15 @@ class TreePickler(pickler: TastyPickler) {
       ann.symbol == defn.BodyAnnot // inline bodies are reconstituted automatically when unpickling
   }
 
-  def pickleAnnotation(owner: Symbol, ann: Annotation)(using Context): Unit = {
-    if (!isUnpicklable(owner, ann)) {
+  def pickleAnnotation(owner: Symbol, mdef: MemberDef, ann: Annotation)(using Context): Unit =
+    if !isUnpicklable(owner, ann) then
       writeByte(ANNOTATION)
       withLength { pickleType(ann.symbol.typeRef); pickleTree(ann.tree) }
-    }
-  }
+      var treeBuf = annotTrees.lookup(mdef)
+      if treeBuf == null then
+        treeBuf = new mutable.ListBuffer[Tree]
+        annotTrees.enter(mdef, treeBuf)
+      treeBuf += ann.tree
 
 // ---- main entry points ---------------------------------------
 

compiler/src/dotty/tools/dotc/quoted/PickledQuotes.scala

@@ -15,13 +15,15 @@ import dotty.tools.dotc.core.tasty.TreePickler.Hole
 import dotty.tools.dotc.core.tasty.{ PositionPickler, TastyPickler, TastyPrinter }
 import dotty.tools.dotc.core.tasty.DottyUnpickler
 import dotty.tools.dotc.core.tasty.TreeUnpickler.UnpickleMode
+import dotty.tools.dotc.report
 
 import dotty.tools.tasty.TastyString
 
 import scala.reflect.ClassTag
 
 import scala.internal.quoted.Unpickler._
 import scala.quoted.QuoteContext
+import scala.collection.mutable
 
 object PickledQuotes {
   import tpd._
@@ -167,8 +169,11 @@ object PickledQuotes {
     val treePkl = pickler.treePkl
     treePkl.pickle(tree :: Nil)
     treePkl.compactify()
-    if (tree.span.exists)
-      new PositionPickler(pickler, treePkl.buf.addrOfTree).picklePositions(tree :: Nil)
+    if tree.span.exists then
+      val positionWarnings = new mutable.ListBuffer[String]()
+      new PositionPickler(pickler, treePkl.buf.addrOfTree, treePkl.treeAnnots)
+        .picklePositions(tree :: Nil, positionWarnings)
+      positionWarnings.foreach(report.warning(_))
 
     val pickled = pickler.assembleParts()
     quotePickling.println(s"**** pickled quote\n${new TastyPrinter(pickled).printContents()}")