improvement: look for definition in pc only for local symbols in the …

…current tree
scala · tgodzik · Feb 4, 2025 · Jan 31, 2025 · Jan 31, 2025 · Feb 3, 2025
commit 559edd49631715d3135000cfb183e292759ea9ae
diff --git a/presentation-compiler/src/main/dotty/tools/pc/PcDefinitionProvider.scala b/presentation-compiler/src/main/dotty/tools/pc/PcDefinitionProvider.scala
@@ -1,5 +1,6 @@
 package dotty.tools.pc
 
+import java.net.URI
 import java.nio.file.Paths
 import java.util.ArrayList
 
@@ -16,6 +17,7 @@ import dotty.tools.dotc.core.Contexts.Context
 import dotty.tools.d
8000
otc.core.Flags.{Exported, ModuleClass}
 import dotty.tools.dotc.core.Symbols.*
 import dotty.tools.dotc.interactive.Interactive
+import dotty.tools.dotc.interactive.Interactive.Include
 import dotty.tools.dotc.interactive.InteractiveDriver
 import dotty.tools.dotc.util.SourceFile
 import dotty.tools.dotc.util.SourcePosition
@@ -51,10 +53,10 @@ class PcDefinitionProvider(
     given ctx: Context = driver.localContext(params)
     val indexedContext = IndexedContext(ctx)
     val result =
-      if findTypeDef then findTypeDefinitions(path, pos, indexedContext)
-      else findDefinitions(path, pos, indexedContext)
+      if findTypeDef then findTypeDefinitions(path, pos, indexedContext, uri)
+      else findDefinitions(path, pos, indexedContext, uri)
 
-    if result.locations().nn.isEmpty() then fallbackToUntyped(pos)(using ctx)
+    if result.locations().nn.isEmpty() then fallbackToUntyped(pos, uri)(using ctx)
     else result
   end definitions
 
@@ -70,32 +72,35 @@ class PcDefinitionProvider(
    * @param pos cursor position
    * @return definition result
    */
-  private def fallbackToUntyped(pos: SourcePosition)(
+  private def fallbackToUntyped(pos: SourcePosition, uri: URI)(
     using ctx: Context
   ) =
     lazy val untpdPath = NavigateAST
       .untypedPath(pos.span)
       .collect { case t: untpd.Tree => t }
 
-    definitionsForSymbol(untpdPath.headOption.map(_.symbol).toList, pos)
+    definitionsForSymbol(untpdPath.headOption.map(_.symbol).toList, uri, pos)
   end fallbackToUntyped
 
   private def findDefinitions(
       path: List[Tree],
       pos: SourcePosition,
-      indexed: IndexedContext
+      indexed: IndexedContext,
+      uri: URI,
   ): DefinitionResult =
     import indexed.ctx
     definitionsForSymbol(
       MetalsInteractive.enclosingSymbols(path, pos, indexed),
+      uri,
       pos
     )
   end findDefinitions
 
   private def findTypeDefinitions(
       path: List[Tree],
       pos: SourcePosition,
-      indexed: IndexedContext
+      indexed: IndexedContext,
+      uri: URI,
   ): DefinitionResult =
     import indexed.ctx
     val enclosing = path.expandRangeToEnclosingApply(pos)
@@ -108,24 +113,25 @@ class PcDefinitionProvider(
       case Nil =>
         path.headOption match
           case Some(value: Literal) =>
-            definitionsForSymbol(List(value.typeOpt.widen.typeSymbol), pos)
+            definitionsForSymbol(List(value.typeOpt.widen.typeSymbol), uri, pos)
           case _ => DefinitionResultImpl.empty
       case _ =>
-        definitionsForSymbol(typeSymbols, pos)
+        definitionsForSymbol(typeSymbols, uri, pos)
 
   end findTypeDefinitions
 
   private def definitionsForSymbol(
       symbols: List[Symbol],
+      uri: URI,
       pos: SourcePosition
   )(using ctx: Context): DefinitionResult =
     symbols match
       case symbols @ (sym :: other) =>
         val isLocal = sym.source == pos.source
         if isLocal then
+          val include = Include.definitions | Include.local
           val (exportedDefs, otherDefs) =
-            Interactive.findDefinitions(List(sym), driver, false, false)
-              .filter(_.source == sym.source)
+            Interactive.findTreesMatching(driver.openedTrees(uri), include, sym)
               .partition(_.tree.symbol.is(Exported))
 
           otherDefs.headOption.orElse(exportedDefs.headOption)  match

diff --git a/presentation-compiler/test/dotty/tools/pc/tests/CompilerCachingSuite.scala b/presentation-compiler/test/dotty/tools/pc/tests/CompilerCachingSuite.scala
@@ -6,8 +6,10 @@ import dotty.tools.pc.ScalaPresentationCompiler
 import org.junit.{Before, Test}
 
 import scala.language.unsafeNulls
-import scala.meta.internal.metals.EmptyCancelToken
 import scala.meta.internal.metals.CompilerOffsetParams
+import scala.meta.internal.metals.EmptyCancelToken
+import scala.meta.internal.metals.EmptyReportContext
+import scala.meta.internal.metals.PcQueryContext
 import scala.meta.pc.OffsetParams
 import scala.concurrent.Future
 import scala.concurrent.Await
@@ -26,20 +28,22 @@ class CompilerCachingSuite extends BasePCSuite:
   private def checkCompilationCount(expected: Int): Unit =
     presentationCompiler match
       case pc: ScalaPresentationCompiler =>
-        val compilations = pc.compilerAccess.withNonInterruptableCompiler(None)(-1, EmptyCancelToken) { driver =>
+        val compilations = pc.compilerAccess.withNonInterruptableCompiler(-1, EmptyCancelToken) { driver =>
           driver.compiler().currentCtx.runId
-        }.get(timeout.length, timeout.unit)
+        }(emptyQueryContext).get(timeout.length, timeout.unit)
         assertEquals(expected, compilations, s"Expected $expected compilations but got $compilations")
       case _ => throw IllegalStateException("Presentation compiler should always be of type of ScalaPresentationCompiler")
 
  private def getContext(): Context =
     presentationCompiler match
       case pc: ScalaPresentationCompiler =>
-        pc.compilerAccess.withNonInterruptableCompiler(None)(null, EmptyCancelToken) { driver =>
+        pc.compilerAccess.withNonInterruptableCompiler(null, EmptyCancelToken) { driver =>
           driver.compiler().currentCtx
-        }.get(timeout.length, timeout.unit)
+        }(emptyQueryContext).get(timeout.length, timeout.unit)
       case _ => throw IllegalStateException("Presentation compiler should always be of type of ScalaPresentationCompiler")
 
+  private def emptyQueryContext = PcQueryContext(None, () => "")(using EmptyReportContext)
+
   @Before
   def beforeEach: Unit =
     presentationCompiler.restart()

diff --git a/presentation-compiler/test/dotty/tools/pc/utils/TestInlayHints.scala b/presentation-compiler/test/dotty/tools/pc/utils/TestInlayHints.scala
@@ -3,8 +3,10 @@ package dotty.tools.pc.utils
 import scala.collection.mutable.ListBuffer
 
 import scala.meta.internal.jdk.CollectionConverters._
+import scala.meta.internal.pc.InlayHints
 import dotty.tools.pc.utils.InteractiveEnrichments.*
 
+import com.google.gson.JsonElement
 import org.eclipse.lsp4j.InlayHint
 import org.eclipse.lsp4j.TextEdit
 import org.eclipse.{lsp4j => l}
@@ -31,7 +33,7 @@ object TestInlayHints {
       case Right(labelParts) => labelParts.asScala.map(_.getValue()).toList
     }
     val data =
-      inlayHint.getData().asInstanceOf[Array[Any]]
+      InlayHints.fromData(inlayHint.getData().asInstanceOf[JsonElement])._2
     buffer += "/*"
     labels.zip(data).foreach { case (label, data) =>
       buffer += label.nn
@@ -41,15 +43,13 @@ object TestInlayHints {
     buffer.toList.mkString
   }
 
-  private def readData(data: Any): List[String] = {
-    data match {
-      case data: String if data.isEmpty => Nil
-      case data: String => List("<<", data, ">>")
-      case data: l.Position =>
+  private def readData(data: Either[String, l.Position]): List[String] =
+    data match
+      case Left("") => Nil
+      case Left(data) => List("<<", data, ">>")
+      case Right(data) =>
         val str = s"(${data.getLine()}:${data.getCharacter()})"
         List("<<", str, ">>")
-    }
-  }
 
   def applyInlayHints(text: String, inlayHints: List[InlayHint]): String = {
     val textEdits = inlayHints.map { hint =>