Produce one TASTY pickling per top-level class.

If a unit has several top-level classes or object (which are not linked classes
of each other) each gets its own pickle information, which contains
any enclosing package clauses and imports and then just the top-level
class/object and its companion object.

3 files changed, 67 insertions, 24 deletions

diff --git a/src/dotty/tools/dotc/CompilationUnit.scala b/src/dotty/tools/dotc/CompilationUnit.scala
index 60e16ec3e..4f8c30aab 100644
--- a/src/dotty/tools/dotc/CompilationUnit.scala
+++ b/src/dotty/tools/dotc/CompilationUnit.scala
@@ -19,11 +19,11 @@ class CompilationUnit(val source: SourceFile) {
   def isJava = source.file.name.endsWith(".java")
 
   /**
-   * Pickler used to create TASTY sections.
+   * 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.
    */
-  lazy val pickler: TastyPickler = new TastyPickler()
+  var picklers: Map[ClassSymbol, TastyPickler] = Map()
 
   /**
    * Addresses in TASTY file of trees, stored by pickling.
diff --git a/src/dotty/tools/dotc/ast/TreeInfo.scala b/src/dotty/tools/dotc/ast/TreeInfo.scala
index a7f89337c..c6d7e10f7 100644
--- a/src/dotty/tools/dotc/ast/TreeInfo.scala
+++ b/src/dotty/tools/dotc/ast/TreeInfo.scala
@@ -461,6 +461,34 @@ trait TypedTreeInfo extends TreeInfo[Type] { self: Trees.Instance[Type] =>
     accum(Nil, root)
   }
 
+
+  /** The top level classes in this tree, including only those module classes that
+   *  are not a linked class of some other class in the result.
+   */
+  def topLevelClasses(tree: Tree)(implicit ctx: Context): List[ClassSymbol] = tree match {
+    case PackageDef(_, stats) => stats.flatMap(topLevelClasses)
+    case tdef: TypeDef if tdef.symbol.isClass => tdef.symbol.asClass :: Nil
+    case _ => Nil
+  }
+
+  /** The tree containing only the top-level classes and objects matching either `cls` or its companion object */
+  def sliceTopLevel(tree: Tree, cls: ClassSymbol)(implicit ctx: Context): List[Tree] = tree match {
+    case PackageDef(pid, stats) =>
+      cpy.PackageDef(tree)(pid, stats.flatMap(sliceTopLevel(_, cls))) :: Nil
+    case tdef: TypeDef =>
+      val sym = tdef.symbol
+      assert(sym.isClass)
+      if (cls == sym || cls == sym.linkedClass) tdef :: Nil
+      else Nil
+    case vdef: ValDef =>
+      val sym = vdef.symbol
+      assert(sym is Module)
+      if (cls == sym.companionClass || cls == sym.moduleClass) vdef :: Nil
+      else Nil
+    case tree =>
+      tree :: Nil
+  }
+
   /** The statement sequence that contains a definition of `sym`, or Nil
    *  if none was found.
    *  For a tree to be found, The symbol must have a position and its definition
diff --git a/src/dotty/tools/dotc/transform/Pickler.scala b/src/dotty/tools/dotc/transform/Pickler.scala
index 515dad2aa..c09ba6889 100644
--- a/src/dotty/tools/dotc/transform/Pickler.scala
+++ b/src/dotty/tools/dotc/transform/Pickler.scala
@@ -9,6 +9,8 @@ import config.Printers.{noPrinter, pickling}
 import java.io.PrintStream
 import Periods._
 import Phases._
+import Symbols._
+import Flags.Module
 import collection.mutable
 
 /** This phase pickles trees */
@@ -23,28 +25,41 @@ class Pickler extends Phase {
     s.close
   }
 
-  private val beforePickling = new mutable.HashMap[CompilationUnit, String]
+  private val beforePickling = new mutable.HashMap[ClassSymbol, String]
+
+  /** 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] = {
+    val companionModuleClasses =
+      clss.filterNot(_ is Module).map(_.linkedClass).filterNot(_.isAbsent)
+    clss.filterNot(companionModuleClasses.contains)
+  }
 
   override def run(implicit ctx: Context): Unit = {
     val unit = ctx.compilationUnit
-    val tree = unit.tpdTree
     pickling.println(i"unpickling in run ${ctx.runId}")
-    if (ctx.settings.YtestPickler.value) beforePickling(unit) = tree.show
 
-    val pickler = unit.pickler
-    val treePkl = new TreePickler(pickler)
-    treePkl.pickle(tree :: Nil)
-    unit.addrOfTree = treePkl.buf.addrOfTree
-    unit.addrOfSym = treePkl.addrOfSym
-    if (tree.pos.exists)
-      new PositionPickler(pickler, treePkl.buf.addrOfTree).picklePositions(tree :: Nil, tree.pos)
+    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)
+      val treePkl = new TreePickler(pickler)
+      treePkl.pickle(tree :: Nil)
+      unit.addrOfTree = treePkl.buf.addrOfTree
+      unit.addrOfSym = treePkl.addrOfSym
+      if (tree.pos.exists)
+        new PositionPickler(pickler, treePkl.buf.addrOfTree).picklePositions(tree :: Nil, tree.pos)
 
-    def rawBytes = // not needed right now, but useful to print raw format.
-      unit.pickler.assembleParts().iterator.grouped(10).toList.zipWithIndex.map {
-        case (row, i) => s"${i}0: ${row.mkString(" ")}"
+      def rawBytes = // not needed right now, but useful to print raw format.
+        pickler.assembleParts().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) {
+        println(i"**** pickled info of $cls")
+        new TastyPrinter(pickler.assembleParts()).printContents()
       }
-    // println(i"rawBytes = \n$rawBytes%\n%") // DEBUG
-    if (pickling ne noPrinter) new TastyPrinter(pickler.assembleParts()).printContents()
+    }
   }
 
   override def runOn(units: List[CompilationUnit])(implicit ctx: Context): List[CompilationUnit] = {
@@ -58,23 +73,23 @@ class Pickler extends Phase {
     pickling.println(i"testing unpickler at run ${ctx.runId}")
     ctx.definitions.init
     val unpicklers =
-      for (unit <- units) yield {
-        val unpickler = new DottyUnpickler(unit.pickler.assembleParts())
+      for (unit <- units; (cls, pickler) <- unit.picklers) yield {
+        val unpickler = new DottyUnpickler(pickler.assembleParts())
         unpickler.enter(roots = Set())
-        unpickler
+        cls -> unpickler
       }
     pickling.println("************* entered toplevel ***********")
-    for ((unpickler, unit) <- unpicklers zip units) {
+    for ((cls, unpickler) <- unpicklers) {
       val unpickled = unpickler.body(readPositions = false)
-      testSame(i"$unpickled%\n%", beforePickling(unit), unit)
+      testSame(i"$unpickled%\n%", beforePickling(cls), cls)
     }
   }
 
-  private def testSame(unpickled: String, previous: String, unit: CompilationUnit)(implicit ctx: Context) =
+  private def testSame(unpickled: String, previous: String, cls: ClassSymbol)(implicit ctx: Context) =
     if (previous != unpickled) {
       output("before-pickling.txt", previous)
       output("after-pickling.txt", unpickled)
-      ctx.error(s"""pickling difference for $unit, for details:
+      ctx.error(s"""pickling difference for ${cls.fullName}, for details:
                    |
                    |  diff before-pickling.txt after-pickling.txt""".stripMargin)
     }