Don't keep full picklers around until backend.

The memory footprint captured by pickler seems to be about
1/3rd of total footprint. So we gain a lot by not making this
die sooner rather than later.

3 files changed, 17 insertions, 14 deletions

diff --git a/src/dotty/tools/backend/jvm/GenBCode.scala b/src/dotty/tools/backend/jvm/GenBCode.scala
index 902f73ae2..65dcb6c79 100644
--- a/src/dotty/tools/backend/jvm/GenBCode.scala
+++ b/src/dotty/tools/backend/jvm/GenBCode.scala
@@ -203,8 +203,7 @@ class GenBCodePipeline(val entryPoints: List[Symbol], val int: DottyBackendInter
         val plainC = pcb.cnode
 
         if (claszSymbol.isClass) // @DarkDimius is this test needed here?
-          for (pickler <- ctx.compilationUnit.picklers.get(claszSymbol.asClass)) {
-            val binary = pickler.assembleParts()
+          for (binary <- ctx.compilationUnit.pickled.get(claszSymbol.asClass)) {
             val dataAttr = new CustomAttr(nme.TASTYATTR.toString, binary)
             (if (mirrorC ne null) mirrorC else plainC).visitAttribute(dataAttr)
           }
diff --git a/src/dotty/tools/dotc/CompilationUnit.scala b/src/dotty/tools/dotc/CompilationUnit.scala
index 16a59250b..491c2bd9b 100644
--- a/src/dotty/tools/dotc/CompilationUnit.scala
+++ b/src/dotty/tools/dotc/CompilationUnit.scala
@@ -17,12 +17,8 @@ class CompilationUnit(val source: SourceFile) {
 
   def isJava = source.file.name.endsWith(".java")
 
-  /**
-   * Picklers used to create TASTY sections, indexed by toplevel class to which they belong.
-   * Sections: Header, ASTs and Positions are populated by `pickler` phase.
-   * Subsequent phases can add new sections.
-   */
-  var picklers: Map[ClassSymbol, TastyPickler] = Map()
+  /** Pickled TASTY binaries, indexed by class. */
+  var pickled: Map[ClassSymbol, Array[Byte]] = Map()
 
   var unpicklers: Map[ClassSymbol, TastyUnpickler] = Map()
 }
diff --git a/src/dotty/tools/dotc/transform/Pickler.scala b/src/dotty/tools/dotc/transform/Pickler.scala
index fc70ac4f2..71cb355df 100644
--- a/src/dotty/tools/dotc/transform/Pickler.scala
+++ b/src/dotty/tools/dotc/transform/Pickler.scala
@@ -25,7 +25,9 @@ class Pickler extends Phase {
     s.close
   }
 
+  // Maps that keep a record if -Ytest-pickler is set.
   private val beforePickling = new mutable.HashMap[ClassSymbol, String]
+  private val picklers = new mutable.HashMap[ClassSymbol, TastyPickler]
 
   /** Drop any elements of this list that are linked module classes of other elements in the list */
   private def dropCompanionModuleClasses(clss: List[ClassSymbol])(implicit ctx: Context): List[ClassSymbol] = {
@@ -40,9 +42,11 @@ class Pickler extends Phase {
 
     for { cls <- dropCompanionModuleClasses(topLevelClasses(unit.tpdTree))
           tree <- sliceTopLevel(unit.tpdTree, cls) } {
-      if (ctx.settings.YtestPickler.value) beforePickling(cls) = tree.show
       val pickler = new TastyPickler()
-      unit.picklers += (cls -> pickler)
+      if (ctx.settings.YtestPickler.value) {
+        beforePickling(cls) = tree.show
+        picklers(cls) = pickler
+      }
       val treePkl = pickler.treePkl
       treePkl.pickle(tree :: Nil)
       treePkl.compactify()
@@ -51,8 +55,12 @@ class Pickler extends Phase {
       if (tree.pos.exists)
         new PositionPickler(pickler, treePkl.buf.addrsOfTree).picklePositions(tree :: Nil)
 
+      // other pickle sections go here.
+      val pickled = pickler.assembleParts()
+      unit.pickled += (cls -> pickled)
+
       def rawBytes = // not needed right now, but useful to print raw format.
-        pickler.assembleParts().iterator.grouped(10).toList.zipWithIndex.map {
+        pickled.iterator.grouped(10).toList.zipWithIndex.map {
           case (row, i) => s"${i}0: ${row.mkString(" ")}"
         }
       // println(i"rawBytes = \n$rawBytes%\n%") // DEBUG
@@ -66,18 +74,18 @@ class Pickler extends Phase {
   override def runOn(units: List[CompilationUnit])(implicit ctx: Context): List[CompilationUnit] = {
     val result = super.runOn(units)
     if (ctx.settings.YtestPickler.value)
-      testUnpickler(units)(
+      testUnpickler(
           ctx.fresh
             .setPeriod(Period(ctx.runId + 1, FirstPhaseId))
             .addMode(Mode.ReadPositions))
     result
   }
 
-  private def testUnpickler(units: List[CompilationUnit])(implicit ctx: Context): Unit = {
+  private def testUnpickler(implicit ctx: Context): Unit = {
     pickling.println(i"testing unpickler at run ${ctx.runId}")
     ctx.initialize()
     val unpicklers =
-      for (unit <- units; (cls, pickler) <- unit.picklers) yield {
+      for ((cls, pickler) <- picklers) yield {
         val unpickler = new DottyUnpickler(pickler.assembleParts())
         unpickler.enter(roots = Set())
         cls -> unpickler