Make position pickling context-independent

Author: odersky

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

@@ -18,6 +18,7 @@ import printing.Printer
 import printing.Texts._
 import util.SourceFile
 import annotation.constructorOnly
+import collection.mutable
 
 object TreePickler {
 
@@ -44,6 +45,15 @@ 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]]()
+
+  def treeAnnots(tree: untpd.MemberDef): List[Tree] =
+    val ts = annotTrees.lookup(tree)
+    if ts == null then Nil else ts.toList
+
   private def withLength(op: => Unit) = {
     val lengthAddr = reserveRef(relative = true)
     op
@@ -309,7 +319,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 +332,16 @@ class TreePickler(pickler: TastyPickler) {
         case _ if tpt.isType => pickleTpt(tpt)
       }
       pickleTreeUnlessEmpty(rhs)
-      pickleModifiers(sym)
+      pickleModifiers(sym, mdef)
     }
   }
 
   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 +531,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 +542,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 +659,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 +671,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 +734,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()}")

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

@@ -51,41 +51,48 @@ class Pickler extends Phase {
     val unit = ctx.compilationUnit
     pickling.println(i"unpickling in run ${ctx.runId}")
 
-    for {
+    for
       cls <- dropCompanionModuleClasses(topLevelClasses(unit.tpdTree))
       tree <- sliceTopLevel(unit.tpdTree, cls)
-    }
-    {
+    do
       val pickler = new TastyPickler(cls)
       if ctx.settings.YtestPickler.value then
         beforePickling(cls) = tree.show
         picklers(cls) = pickler
       val treePkl = pickler.treePkl
       treePkl.pickle(tree :: Nil)
-      val pickledF = Future {
-        treePkl.compactify()
-        if tree.span.exists then
-          new PositionPickler(pickler, treePkl.buf.addrOfTree).picklePositions(tree :: Nil)
-
-        if !ctx.settings.YdropComments.value then
-          new CommentPickler(pickler, treePkl.buf.addrOfTree).pickleComment(tree)
-
-        val pickled = pickler.assembleParts()
-
-        def rawBytes = // not needed right now, but useful to print raw format.
-          pickled.iterator.grouped(10).toList.zipWithIndex.map {
-            case (row, i) => s"${i}0: ${row.mkString(" ")}"
-          }
-
-        // println(i"rawBytes = \n$rawBytes%\n%") // DEBUG
-        if pickling ne noPrinter then
-          pickling.synchronized {
-            println(i"**** pickled info of $cls")
-            println(new TastyPrinter(pickled).printContents())
-          }
-        pickled
-      }(using ExecutionContext.global)
-      def force(): Array[Byte] = Await.result(pickledF, Duration.Inf)
+      val positionWarnings = new mutable.ListBuffer[String]()
+      val parContext = ctx.fresh // TODO refine
+      val pickledF = inContext(parContext) {
+        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)
+
+          val pickled = pickler.assembleParts()
+
+          def rawBytes = // not needed right now, but useful to print raw format.
+            pickled.iterator.grouped(10).toList.zipWithIndex.map {
+              case (row, i) => s"${i}0: ${row.mkString(" ")}"
+            }
+
+          // println(i"rawBytes = \n$rawBytes%\n%") // DEBUG
+          if pickling ne noPrinter then
+            pickling.synchronized {
+              println(i"**** pickled info of $cls")
+              println(new TastyPrinter(pickled).printContents())
+            }
+          pickled
+        }(using ExecutionContext.global)
+      }
+      def force(): Array[Byte] =
+        val result = Await.result(pickledF, Duration.Inf)
+        positionWarnings.foreach(report.warning(_))
+        result
 
       // Turn off parallelism because it lead to non-deterministic CI failures:
       // - https://github.com/lampepfl/dotty/runs/1029579877?check_suite_focus=true#step:10:967
@@ -98,7 +105,7 @@ class Pickler extends Phase {
       if !ctx.settings.YparallelPickler.value then force()
 
       unit.pickled += (cls -> force)
-    }
+    end for
   }
 
   override def runOn(units: List[CompilationUnit])(using Context): List[CompilationUnit] = {