Introduce SourceImportTree and SourceNamedTree

Author: Duhemm

compiler/src/dotty/tools/dotc/interactive/Interactive.scala

@@ -251,7 +251,7 @@ object Interactive {
    *  source code.
    */
   def namedTrees(trees: List[SourceTree], include: Include.Set, sym: Symbol)
-   (implicit ctx: Context): List[SourceTree] =
+   (implicit ctx: Context): List[SourceNamedTree] =
     if (!sym.exists)
       Nil
     else
@@ -262,38 +262,40 @@ object Interactive {
    *  @param includeReferences  If true, include references and not just definitions
    */
   def namedTrees(trees: List[SourceTree], includeReferences: Boolean, nameSubstring: String)
-   (implicit ctx: Context): List[SourceTree] =
+   (implicit ctx: Context): List[SourceNamedTree] =
     namedTrees(trees, includeReferences, _.show.toString.contains(nameSubstring))
 
   /** Find named trees with a non-empty position satisfying `treePredicate` in `trees`.
    *
    *  @param includeReferences  If true, include references and not just definitions
    */
   def namedTrees(trees: List[SourceTree], includeReferences: Boolean, treePredicate: NameTree => Boolean)
-    (implicit ctx: Context): List[SourceTree] = safely {
-    val buf = new mutable.ListBuffer[SourceTree]
-
-    trees foreach { case SourceTree(_, topTree, source) =>
-      new untpd.TreeTraverser {
-        override def traverse(tree: untpd.Tree)(implicit ctx: Context) = {
-          tree match {
-            case utree: untpd.NameTree if tree.hasType =>
-              val tree = utree.asInstanceOf[tpd.NameTree]
-              if (tree.symbol.exists
-                   && !tree.symbol.is(Synthetic)
-                   && tree.pos.exists
-                   && !tree.pos.isZeroExtent
-                   && (includeReferences || isDefinition(tree))
-                   && treePredicate(tree))
-                buf += SourceTree(Nil, tree, source)
-              traverseChildren(tree)
-            case tree: untpd.Inlined =>
-              traverse(tree.call)
-            case _ =>
-              traverseChildren(tree)
+    (implicit ctx: Context): List[SourceNamedTree] = safely {
+    val buf = new mutable.ListBuffer[SourceNamedTree]
+
+    trees foreach {
+      case SourceNamedTree(topTree, source) =>
+        new untpd.TreeTraverser {
+          override def traverse(tree: untpd.Tree)(implicit ctx: Context) = {
+            tree match {
+              case utree: untpd.NameTree if tree.hasType =>
+                val tree = utree.asInstanceOf[tpd.NameTree]
+                if (tree.symbol.exists
+                  && !tree.symbol.is(Synthetic)
+                  && tree.pos.exists
+                  && !tree.pos.isZeroExtent
+                  && (includeReferences || isDefinition(tree))
+                  && treePredicate(tree))
+                  buf += SourceNamedTree(tree, source)
+                traverseChildren(tree)
+              case tree: untpd.Inlined =>
+                traverse(tree.call)
+              case _ =>
+                traverseChildren(tree)
+            }
           }
-        }
-      }.traverse(topTree)
+        }.traverse(topTree)
+      case _: SourceImportTree => ()
     }
 
     buf.toList
@@ -304,9 +306,9 @@ object Interactive {
    *  the tree closest enclosing `pos` and ends with an element of `trees`.
    */
   def pathTo(trees: List[SourceTree], pos: SourcePosition)(implicit ctx: Context): List[Tree] = {
-    val allTrees = trees.flatMap(t => t.tree :: t.topLevelImports)
-    allTrees.find(t => t.pos.contains(pos.pos)) match {
-      case Some(tree) => pathTo(tree, pos.pos)
+    trees.find(t => t.pos.contains(pos)) match {
+      case Some(tree) =>
+        pathTo(tree.tree, pos.pos)
       case None => Nil
     }
   }

compiler/src/dotty/tools/dotc/interactive/InteractiveDriver.scala

@@ -67,14 +67,14 @@ class InteractiveDriver(val settings: List[String]) extends Driver {
     (fromSource ++ fromClassPath).distinct
   }
 
-  private def tree(className: TypeName, id: String)(implicit ctx: Context): Option[SourceTree] = {
+  private def tree(className: TypeName, id: String)(implicit ctx: Context): List[SourceTree] = {
     val clsd = ctx.base.staticRef(className)
     clsd match {
       case clsd: ClassDenotation =>
         clsd.ensureCompleted()
         SourceTree.fromSymbol(clsd.symbol.asClass, id)
       case _ =>
-        None
+        Nil
     }
   }
 
@@ -169,18 +169,19 @@ class InteractiveDriver(val settings: List[String]) extends Driver {
     names.toList
   }
 
-  private def topLevelClassTrees(topTree: Tree, source: SourceFile): List[SourceTree] = {
+  private def topLevelTrees(topTree: Tree, source: SourceFile): List[SourceTree] = {
     val trees = new mutable.ListBuffer[SourceTree]
 
-    def addTrees(topLevelImports: List[Import], tree: Tree): Unit = tree match {
+    def addTrees(tree: Tree): Unit = tree match {
       case PackageDef(_, stats) =>
-        val imports = stats.collect { case imp: Import => imp }
-        stats.foreach(addTrees(imports, _))
+        stats.foreach(addTrees)
+      case imp: Import =>
+        trees += SourceImportTree(imp, source)
       case tree: TypeDef =>
-        trees += SourceTree(topLevelImports, tree, source)
+        trees += SourceNamedTree(tree, source)
       case _ =>
     }
-    addTrees(Nil, topTree)
+    addTrees(topTree)
 
     trees.toList
   }
@@ -244,7 +245,7 @@ class InteractiveDriver(val settings: List[String]) extends Driver {
       val unit = ctx.run.units.head
       val t = unit.tpdTree
       cleanup(t)
-      myOpenedTrees(uri) = topLevelClassTrees(t, source)
+      myOpenedTrees(uri) = topLevelTrees(t, source)
       myCompilationUnits(uri) = unit
 
       reporter.removeBufferedMessages

compiler/src/dotty/tools/dotc/interactive/SourceTree.scala

@@ -9,10 +9,24 @@ import core._, core.Decorators.{sourcePos => _, _}
 import Contexts._, NameOps._, Symbols._
 import util._, util.Positions._
 
-/** A typechecked named `tree` coming from `source` */
-case class SourceTree(topLevelImports: List[tpd.Import], tree: tpd.NameTree, source: SourceFile) {
+/** A `tree` coming from `source` */
+sealed trait SourceTree {
+
+  /** The underlying tree. */
+  def tree: tpd.Tree
+
+  /** The source from which `tree` comes. */
+  def source: SourceFile
+
   /** The position of `tree` */
-  def pos(implicit ctx: Context): SourcePosition = source.atPos(tree.pos)
+  final def pos(implicit ctx: Context): SourcePosition = source.atPos(tree.pos)
+}
+
+/** An import coming from `source` */
+case class SourceImportTree(tree: tpd.Import, source: SourceFile) extends SourceTree
+
+/** A typechecked `tree` coming from `source` */
+case class SourceNamedTree(tree: tpd.NameTree, source: SourceFile) extends SourceTree {
 
   /** The position of the name in `tree` */
   def namePos(implicit ctx: Context): SourcePosition = {
@@ -39,22 +53,34 @@ case class SourceTree(topLevelImports: List[tpd.Import], tree: tpd.NameTree, sou
 }
 
 object SourceTree {
-  def fromSymbol(sym: ClassSymbol, id: String = "")(implicit ctx: Context): Option[SourceTree] = {
+  def fromSymbol(sym: ClassSymbol, id: String = "")(implicit ctx: Context): List[SourceTree] = {
     if (sym == defn.SourceFileAnnot || // FIXME: No SourceFile annotation on SourceFile itself
         sym.sourceFile == null) // FIXME: We cannot deal with external projects yet
-      None
+      Nil
     else {
       import ast.Trees._
-      def sourceTreeOfClass(tree: tpd.Tree): Option[SourceTree] = tree match {
+
+      def sourceTreeOfClass(tree: tpd.Tree): Option[SourceNamedTree] = tree match {
         case PackageDef(_, stats) =>
-          val imports = stats.collect { case imp: tpd.Import => imp }
-          stats.flatMap(sourceTreeOfClass(_).map(_.copy(topLevelImports = imports))).headOption
+          stats.flatMap(sourceTreeOfClass).headOption
         case tree: tpd.TypeDef if tree.symbol == sym =>
           val sourceFile = new SourceFile(sym.sourceFile, Codec.UTF8)
-          Some(SourceTree(Nil, tree, sourceFile))
-        case _ => None
+          Some(SourceNamedTree(tree, sourceFile))
+        case _ =>
+          None
+      }
+
+      def sourceImports(tree: tpd.Tree, sourceFile: SourceFile): List[SourceImportTree] = tree match {
+        case PackageDef(_, stats) => stats.flatMap(sourceImports(_, sourceFile))
+        case imp: tpd.Import => SourceImportTree(imp, sourceFile) :: Nil
+        case _ => Nil
+      }
+
+      val tree = sym.treeContaining(id)
+      sourceTreeOfClass(tree) match {
+        case Some(namedTree) => namedTree :: sourceImports(tree, namedTree.source)
+        case None => Nil
       }
-      sourceTreeOfClass(sym.treeContaining(id))
     }
   }
 }

compiler/src/dotty/tools/repl/ReplDriver.scala

@@ -9,7 +9,7 @@ import dotc.reporting.MessageRendering
 import dotc.reporting.diagnostic.MessageContainer
 import dotc.ast.untpd
 import dotc.ast.tpd
-import dotc.interactive.{ SourceTree, Interactive }
+import dotc.interactive.{ SourceNamedTree, Interactive }
 import dotc.core.Contexts.Context
 import dotc.{ CompilationUnit, Run }
 import dotc.core.Mode

language-server/src/dotty/tools/languageserver/DottyLanguageServer.scala

@@ -279,7 +279,7 @@ class DottyLanguageServer extends LanguageServer
               (driver.allTreesContaining(sym.name.sourceModuleName.toString),
                 Include.overriding | Include.overridden)
             else
-              (SourceTree.fromSymbol(sym.topLevelClass.asClass).toList,
+              (SourceTree.fromSymbol(sym.topLevelClass.asClass),
                 Include.overriding)
           Interactive.namedTrees(trees, include, sym)
       }