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.

1 files changed, 28 insertions, 0 deletions

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