IDE: Factor out logic to find references

Author: Duhemm

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

@@ -25,7 +25,9 @@ object Interactive {
     type Set = Int
     val overridden = 1 // include trees whose symbol is overridden by `sym`
     val overriding = 2 // include trees whose symbol overrides `sym` (but for performance only in same source file)
-    val references = 4 // include references and not just definitions
+    val references = 4 // include references
+    val definitions = 8 // include definitions
+    val linkedClass = 16 // include `symbol.linkedClass`
   }
 
   /** Does this tree define a symbol ? */
@@ -299,6 +301,35 @@ object Interactive {
     buf.toList
   }
 
+  /**
+   * Find trees that match `symbol` in `trees`.
+   *
+   * @param trees    The trees to inspect.
+   * @param includes Whether to include references, definitions, etc.
+   * @param symbol   The symbol for which we want to find references.
+   */
+  def findTreesMatching(trees: List[SourceTree],
+                        includes: Include.Set,
+                        symbol: Symbol)(implicit ctx: Context): List[SourceTree] = {
+    val linkedSym = symbol.linkedClass
+    val includeReferences  = (includes & Include.references) != 0
+    val includeDeclaration = (includes & Include.definitions) != 0
+    val includeLinkedClass = (includes & Include.linkedClass) != 0
+    val predicate: NameTree => Boolean = tree =>
+      (  tree.pos.isSourceDerived
+      && !tree.symbol.isConstructor
+      && (includeDeclaration || !Interactive.isDefinition(tree))
+      && (  Interactive.matchSymbol(tree, symbol, includes)
+         || (  includeDeclaration
+            && includeLinkedClass
+            && linkedSym.exists
+            && Interactive.matchSymbol(tree, linkedSym, includes)
+            )
+         )
+      )
+    namedTrees(trees, includeReferences, predicate)
+  }
+
   /** The reverse path to the node that closest encloses position `pos`,
    *  or `Nil` if no such path exists. If a non-empty path is returned it starts with
    *  the tree closest enclosing `pos` and ends with an element of `trees`.

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

@@ -263,7 +263,7 @@ class DottyLanguageServer extends LanguageServer
           case _ =>
             (Nil, Include.overriding)
         }
-      val defs = Interactive.namedTrees(trees, include, sym)
+      val defs = Interactive.findTreesMatching(trees, include | Include.definitions, sym)
       defs.map(d => location(d.namePos)).asJava
     }
   }
@@ -273,7 +273,6 @@ class DottyLanguageServer extends LanguageServer
     val driver = driverFor(uri)
     implicit val ctx = driver.currentCtx
 
-    val includeDeclaration = params.getContext.isIncludeDeclaration
     val pos = sourcePosition(driver, uri, params.getPosition)
     val sym = Interactive.enclosingSourceSymbol(driver.openedTrees(uri), pos)
 
@@ -283,11 +282,10 @@ class DottyLanguageServer extends LanguageServer
       // only need to look for trees in the target directory if the symbol is defined in the
       // current project
       val trees = driver.allTreesContaining(sym.name.sourceModuleName.toString)
-      val refs = Interactive.namedTrees(trees, includeReferences = true, tree =>
-        tree.pos.isSourceDerived
-          && (includeDeclaration || !Interactive.isDefinition(tree))
-          && !tree.symbol.isConstructor
-          && (Interactive.matchSymbol(tree, sym, Include.overriding)))
+      val includeDeclaration = params.getContext.isIncludeDeclaration
+      val includes =
+        Include.references | Include.overriding | (if (includeDeclaration) Include.definitions else 0)
+      val refs = Interactive.findTreesMatching(trees, includes, sym)
 
       refs.map(ref => location(ref.namePos)).asJava
     }
@@ -304,14 +302,10 @@ class DottyLanguageServer extends LanguageServer
     if (sym == NoSymbol) new WorkspaceEdit()
     else {
       val trees = driver.allTreesContaining(sym.name.sourceModuleName.toString)
-      val linkedSym = sym.linkedClass
       val newName = params.getNewName
-
-      val refs = Interactive.namedTrees(trees, includeReferences = true, tree =>
-        tree.pos.isSourceDerived
-        && !tree.symbol.isConstructor
-          && (Interactive.matchSymbol(tree, sym, Include.overriding)
-            || (linkedSym != NoSymbol && Interactive.matchSymbol(tree, linkedSym, Include.overriding))))
+      val includes =
+        Include.references | Include.definitions | Include.linkedClass | Include.overriding
+      val refs = Interactive.findTreesMatching(trees, includes, sym)
 
       val changes = refs.groupBy(ref => toUri(ref.source).toString).mapValues(_.map(ref => new TextEdit(range(ref.namePos), newName)).asJava)