improvement: use pc for finding references of local symbols

scalameta · Dec 19, 2023 · a4d5d77 · a4d5d77
1 parent 67ec987
commit a4d5d77
Show file tree

Hide file tree

Showing 35 changed files with 1,357 additions and 895 deletions.
diff --git a/metals/src/main/scala/scala/meta/internal/metals/AdjustLspData.scala b/metals/src/main/scala/scala/meta/internal/metals/AdjustLspData.scala
@@ -50,6 +50,9 @@ trait AdjustLspData {
     diag
   }
 
+  def adjustLocation(location: Location): Location =
+    new Location(location.getUri(), adjustRange(location.getRange()))
+
   def adjustLocations(
       locations: java.util.List[Location]
   ): ju.List[Location]

diff --git a/metals/src/main/scala/scala/meta/internal/metals/Compilers.scala b/metals/src/main/scala/scala/meta/internal/metals/Compilers.scala
@@ -36,6 +36,7 @@ import scala.meta.pc.OffsetParams
 import scala.meta.pc.PresentationCompiler
 import scala.meta.pc.SymbolSearch
 import scala.meta.pc.SyntheticDecoration
+import scala.meta.pc.VirtualFileParams
 
 import ch.epfl.scala.bsp4j.BuildTargetIdentifier
 import ch.epfl.scala.bsp4j.CompileReport
@@ -46,6 +47,8 @@ import org.eclipse.lsp4j.CompletionParams
 import org.eclipse.lsp4j.Diagnostic
 import org.eclipse.lsp4j.DocumentHighlight
 import org.eclipse.lsp4j.InitializeParams
+import org.eclipse.lsp4j.Location
+import org.eclipse.lsp4j.ReferenceParams
 import org.eclipse.lsp4j.RenameParams
 import org.eclipse.lsp4j.SelectionRange
 import org.eclipse.lsp4j.SelectionRangeParams
@@ -683,6 +686,39 @@ class Compilers(
     }
   }.getOrElse(Future.successful(Nil.asJava))
 
+  def references(
+      params: ReferenceParams,
+      targetFiles: List[AbsolutePath],
+      token: CancelToken,
+  ): Future[List[Location]] = {
+    withPCAndAdjustLsp(params) { (pc, pos, adjust) =>
+      val targets = targetFiles.map { target =>
+        target.toURI.toString -> {
+          val (vFile, _, adjustLsp) =
+            sourceAdjustments(
+              target.toURI.toString(),
+              pc.scalaVersion(),
+            )
+          val params =
+            CompilerVirtualFileParams(target.toURI, vFile.text, token)
+          (params, adjustLsp)
+        }
+      }.toMap
+      val targetFilesParams: List[VirtualFileParams] =
+        targets.values.map(_._1).toList
+      pc.references(
+        CompilerOffsetParamsUtils.fromPos(pos, token),
+        targetFilesParams.asJava,
+        params.getContext().isIncludeDeclaration(),
+      ).asScala
+        .map(
+          _.asScala.toList.map(loc =>
+            targets(loc.getUri())._2.adjustLocation(loc)
+          )
+        )
+    }
+  }.getOrElse(Future.successful(Nil))
+
   def extractMethod(
       doc: TextDocumentIdentifier,
       range: LspRange,

diff --git a/metals/src/main/scala/scala/meta/internal/metals/MetalsLspService.scala b/metals/src/main/scala/scala/meta/internal/metals/MetalsLspService.scala
@@ -549,6 +549,59 @@ class MetalsLspService(
     implementationProvider,
   )
 
+  val worksheetProvider: WorksheetProvider = {
+    val worksheetPublisher =
+      if (clientConfig.isDecorationProvider)
+        new DecorationWorksheetPublisher(
+          clientConfig.isInlineDecorationProvider()
+        )
+      else
+        new WorkspaceEditWorksheetPublisher(buffers, trees)
+
+    register(
+      new WorksheetProvider(
+        folder,
+        buffers,
+        buildTargets,
+        languageClient,
+        () => userConfig,
+        statusBar,
+        diagnostics,
+        embedded,
+        worksheetPublisher,
+        compilations,
+        scalaVersionSelector,
+      )
+    )
+  }
+
+  private val symbolSearch: MetalsSymbolSearch = new MetalsSymbolSearch(
+    symbolDocs,
+    workspaceSymbols,
+    definitionProvider,
+  )
+
+  val compilers: Compilers = register(
+    new Compilers(
+      folder,
+      clientConfig,
+      () => userConfig,
+      buildTargets,
+      buffers,
+      symbolSearch,
+      embedded,
+      statusBar,
+      sh,
+      initializeParams,
+      () => excludedPackageHandler,
+      scalaVersionSelector,
+      trees,
+      mtagsResolver,
+      sourceMapper,
+      worksheetProvider,
+    )
+  )
+
   private val referencesProvider: ReferenceProvider = new ReferenceProvider(
     folder,
     semanticdbs,
@@ -557,6 +610,7 @@ class MetalsLspService(
     remote,
     trees,
     buildTargets,
+    compilers,
   )
 
   private val syntheticHoverProvider: SyntheticHoverProvider =
@@ -640,59 +694,6 @@ class MetalsLspService(
     clientConfig.icons,
   )
 
-  private val symbolSearch: MetalsSymbolSearch = new MetalsSymbolSearch(
-    symbolDocs,
-    workspaceSymbols,
-    definitionProvider,
-  )
-
-  val worksheetProvider: WorksheetProvider = {
-    val worksheetPublisher =
-      if (clientConfig.isDecorationProvider)
-        new DecorationWorksheetPublisher(
-          clientConfig.isInlineDecorationProvider()
-        )
-      else
-        new WorkspaceEditWorksheetPublisher(buffers, trees)
-
-    register(
-      new WorksheetProvider(
-        folder,
-        buffers,
-        buildTargets,
-        languageClient,
-        () => userConfig,
-        statusBar,
-        diagnostics,
-        embedded,
-        worksheetPublisher,
-        compilations,
-        scalaVersionSelector,
-      )
-    )
-  }
-
-  private val compilers: Compilers = register(
-    new Compilers(
-      folder,
-      clientConfig,
-      () => userConfig,
-      buildTargets,
-      buffers,
-      symbolSearch,
-      embedded,
-      statusBar,
-      sh,
-      initializeParams,
-      () => excludedPackageHandler,
-      scalaVersionSelector,
-      trees,
-      mtagsResolver,
-      sourceMapper,
-      worksheetProvider,
-    )
-  )
-
   private val renameProvider: RenameProvider = new RenameProvider(
     referencesProvider,
     implementationProvider,
@@ -1519,18 +1520,20 @@ class MetalsLspService(
   override def references(
       params: ReferenceParams
   ): CompletableFuture[util.List[Location]] =
-    CancelTokens { _ => referencesResult(params).flatMap(_.locations).asJava }
+    CancelTokens.future { _ =>
+      referencesResult(params).map(_.flatMap(_.locations).asJava)
+    }
 
   // Triggers a cascade compilation and tries to find new references to a given symbol.
   // It's not possible to stream reference results so if we find new symbols we notify the
   // user to run references again to see updated results.
   private def compileAndLookForNewReferences(
       params: ReferenceParams,
       result: List[ReferencesResult],
-  ): Unit = {
+  ): Future[Unit] = {
     val path = params.getTextDocument.getUri.toAbsolutePath
     val old = path.toInputFromBuffers(buffers)
-    compilations.cascadeCompileFiles(Seq(path)).foreach { _ =>
+    compilations.cascadeCompileFiles(Seq(path)).flatMap { _ =>
       val newBuffer = path.toInputFromBuffers(buffers)
       val newParams: Option[ReferenceParams] =
         if (newBuffer.text == old.text) Some(params)
@@ -1553,46 +1556,54 @@ class MetalsLspService(
             )
         }
       newParams match {
-        case None =>
+        case None => Future.unit
         case Some(p) =>
-          val newResult = referencesProvider.references(p)
-          val diff = newResult
-            .flatMap(_.locations)
-            .length - result.flatMap(_.locations).length
-          val diffSyms: Set[String] =
-            newResult.map(_.symbol).toSet -- result.map(_.symbol).toSet
-          if (diffSyms.nonEmpty && diff > 0) {
-            import scala.meta.internal.semanticdb.Scala._
-            val names =
-              diffSyms.map(sym => s"'${sym.desc.name.value}'").mkString(" and ")
-            val message =
-              s"Found new symbol references for $names, try running again."
-            scribe.info(message)
-            statusBar
-              .addMessage(clientConfig.icons.info + message)
+          for {
+            newResult <- referencesProvider.references(p)
+          } yield {
+            val diff = newResult
+              .flatMap(_.locations)
+              .length - result.flatMap(_.locations).length
+            val diffSyms: Set[String] =
+              newResult.map(_.symbol).toSet -- result.map(_.symbol).toSet
+            if (diffSyms.nonEmpty && diff > 0) {
+              import scala.meta.internal.semanticdb.Scala._
+              val names =
+                diffSyms
+                  .map(sym => s"'${sym.desc.name.value}'")
+                  .mkString(" and ")
+              val message =
+                s"Found new symbol references for $names, try running again."
+              scribe.info(message)
+              statusBar
+                .addMessage(clientConfig.icons.info + message)
+            }
           }
       }
     }
   }
 
-  def referencesResult(params: ReferenceParams): List[ReferencesResult] = {
+  def referencesResult(
+      params: ReferenceParams
+  ): Future[List[ReferencesResult]] = {
     val timer = new Timer(time)
-    val results: List[ReferencesResult] = referencesProvider.references(params)
-    if (clientConfig.initialConfig.statistics.isReferences) {
-      if (results.forall(_.symbol.isEmpty)) {
-        scribe.info(s"time: found 0 references in $timer")
-      } else {
-        scribe.info(
-          s"time: found ${results.flatMap(_.locations).length} references to symbol '${results
-              .map(_.symbol)
-              .mkString("and")}' in $timer"
-        )
+    referencesProvider.references(params).map { results =>
+      if (clientConfig.initialConfig.statistics.isReferences) {
+        if (results.forall(_.symbol.isEmpty)) {
+          scribe.info(s"time: found 0 references in $timer")
+        } else {
+          scribe.info(
+            s"time: found ${results.flatMap(_.locations).length} references to symbol '${results
+                .map(_.symbol)
+                .mkString("and")}' in $timer"
+          )
+        }
       }
+      if (results.nonEmpty) {
+        compileAndLookForNewReferences(params, results)
+      }
+      results
     }
-    if (results.nonEmpty) {
-      compileAndLookForNewReferences(params, results)
-    }
-    results
   }
 
   override def semanticTokensFull(
@@ -2535,21 +2546,22 @@ class MetalsLspService(
               positionParams.getPosition(),
               new ReferenceContext(false),
             )
-            val results = referencesResult(refParams)
-            if (results.flatMap(_.locations).isEmpty) {
-              // Fallback again to the original behavior that returns
-              // the definition location itself if no reference locations found,
-              // for avoiding the confusing messages like "No definition found ..."
-              definitionResult(positionParams, token)
-            } else {
-              Future.successful(
-                DefinitionResult(
-                  locations = results.flatMap(_.locations).asJava,
-                  symbol = results.head.symbol,
-                  definition = None,
-                  semanticdb = None,
+            referencesResult(refParams).flatMap { results =>
+              if (results.flatMap(_.locations).isEmpty) {
+                // Fallback again to the original behavior that returns
+                // the definition location itself if no reference locations found,
+                // for avoiding the confusing messages like "No definition found ..."
+                definitionResult(positionParams, token)
+              } else {
+                Future.successful(
+                  DefinitionResult(
+                    locations = results.flatMap(_.locations).asJava,
+                    symbol = results.head.symbol,
+                    definition = None,
+                    semanticdb = None,
+                  )
                 )
-              )
+              }
             }
           } else {
             definitionResult(positionParams, token)