diff options
Diffstat (limited to 'src/dotty/tools/dotc/ast/TreeInfo.scala')
-rw-r--r-- | src/dotty/tools/dotc/ast/TreeInfo.scala | 28 |
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 |